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,23 @@ 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()
+    {
+        // The opt-out, as it is for the pool's limits: a caller reading a stream that is legitimately silent for
+        // arbitrarily long has to be able to say so.
+        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,125 @@
+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;
+
+// Giving up on a result has to reach the server, not just the client: without the Cancel packet the server keeps
+// running the query and writing into a socket nobody reads. Error code 735, QUERY_WAS_CANCELLED_BY_CLIENT, is
+// raised only where the server reads that packet, so a query logged with it ended because the client asked rather
+// than because the connection went away.
+[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()
+    {
+        // MaxPoolSize 1, so a lost permit or a connection put back broken leaves nothing for the next query to
+        // run on: the assertion below could not pass by chance on a fresh connection.
+        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()
+    {
+        // Roughly two seconds of work delivered in 200ms blocks, under a one-second deadline. The deadline bounds
+        // the gap between packets, not the response, so only the blocks have to fit inside it. The rows are
+        // selected rather than counted: an aggregate the planner can answer without them prunes the sleep away
+        // and the query returns instantly, proving nothing.
+        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));
+    }
+
+    // A result with no end, so the server is certainly still producing it when the client stops reading. Slow and
+    // small on purpose: a result that saturates the socket blocks the server in a write, where it reads nothing.
+    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);
+
+    // Waits for the query's own log record, which the server writes once the query has actually stopped, and
+    // reports whether it ended because the client cancelled it.
+    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,101 @@ public async Task InsertAsync_TokenAlreadyCancelled_ThrowsWithoutClaimingConnect
         });
     }
 
+    [Test]
+    public async Task InsertAsync_CancelledWhileAwaitingTheSchemaBlock_SendsCancelBeforeTerminating()
+    {
+        // The Query and the end-of-input block are written, then the read for the schema blocks. The server has
+        // consumed everything the client sent, so a Cancel lands on 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()
+    {
+        // The schema read has to carry the deadline's token. With the caller's it is unbounded, and an insert
+        // against a server that stopped answering waits for TCP rather than for ReadTimeout.
+        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()
+    {
+        // The other read the insert makes: the drain to end-of-stream, after every row has gone out. A regression
+        // that bounds only the schema read leaves this one waiting for TCP.
+        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()
+    {
+        // Cancelling from the column factory takes the insert down inside the row stream, where a block is
+        // part-written. A Cancel appended there is read as more block bytes rather than as a packet, so none is
+        // sent and the connection is simply closed.
+        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()
     {