TCP Q3/Q4: an idle read deadline, and telling the server when a result is abandoned

ClickHouse.Driver.Tcp.Tests/Client/ClickHouseTcpClientOptionsTests.cs

@@ -73,13 +73,22 @@ public void Validate_EmptyDatabase_ThrowsArgumentException()
     }
 
     [Test]
-    public void Validate_NonPositiveReadTimeout_ThrowsArgumentOutOfRangeException()
+    public void Validate_NegativeReadTimeout_ThrowsArgumentOutOfRangeException()
     {
-        var options = new ClickHouseTcpClientOptions { ReadTimeout = TimeSpan.Zero };
+        var options = new ClickHouseTcpClientOptions { ReadTimeout = TimeSpan.FromSeconds(-1) };
 
         Assert.Throws<ArgumentOutOfRangeException>(() => options.Validate());
     }
 
+    [Test]
+    public void Validate_ZeroReadTimeout_IsAccepted()
+    {
+        // Zero disables the read timeout.
+        var options = new ClickHouseTcpClientOptions { ReadTimeout = TimeSpan.Zero };
+
+        Assert.DoesNotThrow(() => options.Validate());
+    }
+
     [TestCase(0)]
     [TestCase(-1)]
     [TestCase(65536)]

ClickHouse.Driver.Tcp.Tests/Integration/ClickHouseTcpCancellationIntegrationTests.cs

@@ -0,0 +1,119 @@
+using System;
+using System.Collections.Generic;
+using System.Globalization;
+using System.Threading;
+using System.Threading.Tasks;
+using ClickHouse.Driver.Tcp.Client;
+using ClickHouse.Driver.Tcp.Format;
+using ClickHouse.Driver.Tcp.Tests.Utilities;
+using ClickHouse.Driver.Tcp.Types;
+
+namespace ClickHouse.Driver.Tcp.Tests.Integration;
+
+// Verify cancellation through query_log error code 735 (QUERY_WAS_CANCELLED_BY_CLIENT),
+// which confirms that the server processed the Cancel packet.
+[TestFixture]
+[Category("Integration")]
+public class ClickHouseTcpCancellationIntegrationTests
+{
+    private const int QueryWasCancelledByClient = 735;
+
+    private static readonly CancellationToken None = CancellationToken.None;
+
+    [Test]
+    public async Task StreamAsync_CancelledMidResult_StopsTheQueryOnTheServer()
+    {
+        await using ClickHouseTcpClient client = TcpServerFixture.CreateClient();
+        string queryId = Guid.NewGuid().ToString();
+
+        using var cts = new CancellationTokenSource();
+        Assert.CatchAsync<OperationCanceledException>(async () =>
+        {
+            await foreach (Block block in Unbounded(client, queryId, cts.Token))
+            {
+                _ = block;
+                await cts.CancelAsync();
+            }
+        });
+
+        Assert.That(await CancelledByClientAsync(client, queryId), Is.True);
+    }
+
+    [Test]
+    public async Task StreamAsync_EnumerationAbandonedEarly_StopsTheQueryOnTheServer()
+    {
+        await using ClickHouseTcpClient client = TcpServerFixture.CreateClient();
+        string queryId = Guid.NewGuid().ToString();
+
+        await foreach (Block block in Unbounded(client, queryId, None))
+        {
+            _ = block;
+            break;
+        }
+
+        Assert.That(await CancelledByClientAsync(client, queryId), Is.True);
+    }
+
+    [Test]
+    public async Task StreamAsync_CancelledMidResult_ReturnsThePoolSlotForTheNextOperation()
+    {
+        // A single pool slot verifies that cancellation leaves capacity for the next query.
+        ClickHouseTcpClientOptions options = TcpServerFixture.Options() with { MaxPoolSize = 1 };
+        await using var client = new ClickHouseTcpClient(options);
+
+        using var cts = new CancellationTokenSource();
+        Assert.CatchAsync<OperationCanceledException>(async () =>
+        {
+            await foreach (Block block in Unbounded(client, Guid.NewGuid().ToString(), cts.Token))
+            {
+                _ = block;
+                await cts.CancelAsync();
+            }
+        });
+
+        var answer = 0;
+        await foreach (Block block in client.StreamAsync("SELECT 42", cancellationToken: None))
+        {
+            answer = ((IColumn<byte>)block[0]).Values[0];
+        }
+
+        Assert.That(answer, Is.EqualTo(42));
+    }
+
+    [Test]
+    public async Task StreamAsync_QueryLongerThanReadTimeout_SurvivesBecauseTheDeadlineMeasuresSilence()
+    {
+        // The query takes roughly two seconds, producing ten-row blocks every 200 ms with a one-second read timeout.
+        // Select the sleepEachRow result so the planner must evaluate it.
+        ClickHouseTcpClientOptions options = TcpServerFixture.Options() with { ReadTimeout = TimeSpan.FromSeconds(1) };
+        await using var client = new ClickHouseTcpClient(options);
+
+        var rows = 0;
+        await foreach (Block block in client.StreamAsync(
+            "SELECT number, sleepEachRow(0.02) FROM system.numbers LIMIT 100 SETTINGS max_block_size = 10",
+            cancellationToken: None))
+        {
+            rows += block.RowCount;
+        }
+
+        Assert.That(rows, Is.EqualTo(100));
+    }
+
+    // An unbounded query remains active when the client stops reading. Small, delayed blocks
+    // allow the server to read Cancel between writes.
+    private static IAsyncEnumerable<Block> Unbounded(ClickHouseTcpClient client, string queryId, CancellationToken cancellationToken)
+        => client.StreamAsync(
+            "SELECT number, sleepEachRow(0.02) FROM system.numbers SETTINGS max_block_size = 10",
+            new ClickHouseTcpQueryOptions { QueryId = queryId },
+            cancellationToken);
+
+    // Wait for the query log record and check for the client-cancellation error code.
+    private static async Task<bool> CancelledByClientAsync(ClickHouseTcpClient client, string queryId)
+    {
+        object code = await QueryLog.ScalarAsync(
+            client,
+            $"SELECT exception_code FROM system.query_log WHERE query_id = '{queryId}' AND type != 'QueryStart' ORDER BY event_time_microseconds DESC LIMIT 1");
+
+        return Convert.ToInt32(code, CultureInfo.InvariantCulture) == QueryWasCancelledByClient;
+    }
+}

ClickHouse.Driver.Tcp.Tests/Protocol/ClickHouseTcpConnectionInsertTests.cs

@@ -192,6 +192,97 @@ public async Task InsertAsync_TokenAlreadyCancelled_ThrowsWithoutClaimingConnect
         });
     }
 
+    [Test]
+    public async Task InsertAsync_CancelledWhileAwaitingTheSchemaBlock_SendsCancelBeforeTerminating()
+    {
+        // The request is fully flushed before the schema read blocks, so Cancel can be sent at a packet boundary.
+        var transport = new ScriptedDuplexStream(await ServerHelloBytesAsync(54476), blockWhenExhausted: true);
+        using var connection = new ClickHouseTcpConnection(transport, socket: null);
+        await connection.HandshakeAsync(Handshake, None);
+
+        using var cts = new CancellationTokenSource();
+        Task insert = connection.InsertAsync("INSERT INTO t VALUES", Columns(UInt64Column(1)), cancellationToken: cts.Token).AsTask();
+        await cts.CancelAsync();
+
+        Assert.CatchAsync<OperationCanceledException>(async () => await insert);
+        Assert.Multiple(() =>
+        {
+            Assert.That(transport.Written[^1], Is.EqualTo((byte)ClientPacketType.Cancel));
+            Assert.That(connection.State, Is.EqualTo(TcpConnectionState.Terminated));
+        });
+    }
+
+    [Test]
+    public async Task InsertAsync_ServerGoesSilentAwaitingTheSchemaBlock_ThrowsTimeoutAndSendsCancel()
+    {
+        // Verify that ReadTimeout applies while waiting for the insert schema.
+        var transport = new ScriptedDuplexStream(await ServerHelloBytesAsync(54476), blockWhenExhausted: true);
+        using var connection = new ClickHouseTcpConnection(transport, socket: null, readTimeout: TimeSpan.FromMilliseconds(200));
+        await connection.HandshakeAsync(Handshake, None);
+
+        var thrown = Assert.CatchAsync<TimeoutException>(
+            async () => await connection.InsertAsync("INSERT INTO t VALUES", Columns(UInt64Column(1)), cancellationToken: None));
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(thrown.Message, Does.Contain("ReadTimeout"));
+            Assert.That(transport.Written[^1], Is.EqualTo((byte)ClientPacketType.Cancel));
+            Assert.That(connection.State, Is.EqualTo(TcpConnectionState.Terminated));
+        });
+    }
+
+    [Test]
+    public async Task InsertAsync_ServerGoesSilentDrainingTheAcknowledgement_ThrowsTimeout()
+    {
+        // Verify that ReadTimeout also applies while waiting for the insert acknowledgement.
+        byte[] script = Concat(
+            await ServerHelloBytesAsync(54476),
+            await SchemaBlockAsync(("x", "UInt64")));
+        var transport = new ScriptedDuplexStream(script, blockWhenExhausted: true);
+        using var connection = new ClickHouseTcpConnection(transport, socket: null, readTimeout: TimeSpan.FromMilliseconds(200));
+        await connection.HandshakeAsync(Handshake, None);
+
+        var thrown = Assert.CatchAsync<TimeoutException>(
+            async () => await connection.InsertAsync("INSERT INTO t VALUES", Columns(UInt64Column(1)), cancellationToken: None));
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(thrown.Message, Does.Contain("ReadTimeout"));
+            Assert.That(connection.State, Is.EqualTo(TcpConnectionState.Terminated));
+        });
+    }
+
+    [Test]
+    public async Task InsertAsync_CancelledWhileStreamingRows_LeavesTheTruncatedBlockWithoutAppendingCancel()
+    {
+        // Cancellation during the row phase suppresses Cancel because a write may have left a partial Data packet.
+        // The connection must close without appending a packet.
+        byte[] script = Concat(
+            await ServerHelloBytesAsync(54476),
+            await SchemaBlockAsync(("x", "UInt64")),
+            EndOfStreamPacket());
+        var transport = new ScriptedDuplexStream(script);
+        using var connection = new ClickHouseTcpConnection(transport, socket: null);
+        await connection.HandshakeAsync(Handshake, None);
+
+        using var cts = new CancellationTokenSource();
+        Assert.CatchAsync<OperationCanceledException>(async () => await connection.InsertAsync(
+            "INSERT INTO t VALUES",
+            rowCount: 1,
+            buildColumns: _ =>
+            {
+                cts.Cancel();
+                return new StubInsertColumnSource(UInt64Column(1));
+            },
+            cancellationToken: cts.Token));
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(transport.Written[^1], Is.Not.EqualTo((byte)ClientPacketType.Cancel));
+            Assert.That(connection.State, Is.EqualTo(TcpConnectionState.Terminated));
+        });
+    }
+
     [Test]
     public async Task InsertAsync_ColumnCountDisagreesWithSchema_ThrowsArgumentButStaysReady()
     {