From 31bbbcaf122d671a6d87ae72b0adb28b132d00fb Mon Sep 17 00:00:00 2001 From: Alex Soffronow-Pagonidis Date: Wed, 26 Aug 2026 10:35:21 +0200 Subject: [PATCH] Bound reads with an idle deadline and tell the server when a result is abandoned ReadTimeout was parsed, stored and read by nothing. It now bounds every read of an operation, armed immediately before each read from the transport and disarmed as soon as that read returns, so it measures silence rather than duration: a result that streams for an hour never trips it, and neither does a consumer that holds a block longer than the deadline. TimeSpan.Zero disables it, as it does for the pool's limits. Giving up on a result now sends the Cancel packet before closing the connection, so the server stops rather than finishing a query nobody reads. That covers cancellation, a read that gave up, and a consumer that breaks out of the enumeration. The insert row phase is excluded: a block is part-written there, so an appended Cancel would be read as more block bytes. Co-Authored-By: Claude --- .../Client/ClickHouseTcpClientOptionsTests.cs | 13 +- ...ickHouseTcpCancellationIntegrationTests.cs | 119 ++++++++++ .../ClickHouseTcpConnectionInsertTests.cs | 91 ++++++++ .../ClickHouseTcpConnectionQueryTests.cs | 203 ++++++++++++++++- .../Protocol/ClickHouseTcpConnectionTests.cs | 18 ++ .../Protocol/ReadBufferDeadlineTests.cs | 134 +++++++++++ .../Utilities/QueryLog.cs | 68 ++++++ .../Utilities/ScriptedDuplexStream.cs | 11 +- .../Client/ClickHouseTcpClientOptions.cs | 22 +- .../ClickHouseTcpConnectionStringBuilder.cs | 2 +- .../Client/IConnectionFactory.cs | 2 +- .../Protocol/ClickHouseTcpConnection.cs | 215 +++++++++++++----- .../Protocol/IdleReadDeadline.cs | 76 +++++++ ClickHouse.Driver.Tcp/Protocol/ReadBuffer.cs | 55 +++-- 14 files changed, 947 insertions(+), 82 deletions(-) create mode 100644 ClickHouse.Driver.Tcp.Tests/Integration/ClickHouseTcpCancellationIntegrationTests.cs create mode 100644 ClickHouse.Driver.Tcp.Tests/Protocol/ReadBufferDeadlineTests.cs create mode 100644 ClickHouse.Driver.Tcp.Tests/Utilities/QueryLog.cs create mode 100644 ClickHouse.Driver.Tcp/Protocol/IdleReadDeadline.cs diff --git a/ClickHouse.Driver.Tcp.Tests/Client/ClickHouseTcpClientOptionsTests.cs b/ClickHouse.Driver.Tcp.Tests/Client/ClickHouseTcpClientOptionsTests.cs index f7ab9c837..f6a64faaa 100644 --- a/ClickHouse.Driver.Tcp.Tests/Client/ClickHouseTcpClientOptionsTests.cs +++ b/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(() => 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)] diff --git a/ClickHouse.Driver.Tcp.Tests/Integration/ClickHouseTcpCancellationIntegrationTests.cs b/ClickHouse.Driver.Tcp.Tests/Integration/ClickHouseTcpCancellationIntegrationTests.cs new file mode 100644 index 000000000..2add38632 --- /dev/null +++ b/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(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(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)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 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 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; + } +} diff --git a/ClickHouse.Driver.Tcp.Tests/Protocol/ClickHouseTcpConnectionInsertTests.cs b/ClickHouse.Driver.Tcp.Tests/Protocol/ClickHouseTcpConnectionInsertTests.cs index 6a4863a23..0824ad6ea 100644 --- a/ClickHouse.Driver.Tcp.Tests/Protocol/ClickHouseTcpConnectionInsertTests.cs +++ b/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(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( + 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( + 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(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() { diff --git a/ClickHouse.Driver.Tcp.Tests/Protocol/ClickHouseTcpConnectionQueryTests.cs b/ClickHouse.Driver.Tcp.Tests/Protocol/ClickHouseTcpConnectionQueryTests.cs index 29498bf11..fdbc94eed 100644 --- a/ClickHouse.Driver.Tcp.Tests/Protocol/ClickHouseTcpConnectionQueryTests.cs +++ b/ClickHouse.Driver.Tcp.Tests/Protocol/ClickHouseTcpConnectionQueryTests.cs @@ -3,6 +3,7 @@ using System.IO; using System.Threading; using System.Threading.Tasks; +using ClickHouse.Driver.Compression; using ClickHouse.Driver.Tcp.Format; using ClickHouse.Driver.Tcp.Protocol; using ClickHouse.Driver.Tcp.Tests.Utilities; @@ -155,6 +156,200 @@ await DataPacketAsync(new ulong[] { 2 }), Assert.That(connection.State, Is.EqualTo(TcpConnectionState.Terminated)); } + [Test] + public async Task QueryAsync_EnumerationAbandonedBeforeEndOfStream_SendsCancelBeforeTerminating() + { + byte[] script = Concat( + await ServerHelloBytesAsync(54476), + await DataPacketAsync(new ulong[] { 1 }), + await DataPacketAsync(new ulong[] { 2 }), + EndOfStreamPacket()); + var transport = new ScriptedDuplexStream(script); + using var connection = new ClickHouseTcpConnection(transport, socket: null); + await connection.HandshakeAsync(Handshake, None); + + var writtenBeforeAbandoning = 0; + await foreach (Block block in connection.QueryAsync("SELECT 1", cancellationToken: None)) + { + _ = block; + writtenBeforeAbandoning = transport.Written.Length; + break; + } + + // An incomplete response triggers Cancel before the connection closes. + // Cancel has no body, so it adds one byte after the request. + Assert.Multiple(() => + { + Assert.That(transport.Written, Has.Length.EqualTo(writtenBeforeAbandoning + 1)); + AssertCancelSent(transport); + }); + } + + [Test] + public async Task QueryAsync_CancelledWhileReadingResponse_SendsCancelBeforeTerminating() + { + // The request is flushed, then the first response read blocks. + 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 drain = DrainAsync(connection, cts.Token); + await cts.CancelAsync(); + + Assert.CatchAsync(async () => await drain); + Assert.Multiple(() => + { + AssertCancelSent(transport); + Assert.That(connection.State, Is.EqualTo(TcpConnectionState.Terminated)); + }); + } + + [Test] + public async Task QueryAsync_ServerGoesSilentPastReadTimeout_ThrowsTimeoutAndSendsCancel() + { + 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(async () => await DrainAsync(connection)); + + Assert.Multiple(() => + { + Assert.That(thrown.Message, Does.Contain("ReadTimeout")); + AssertCancelSent(transport); + Assert.That(connection.State, Is.EqualTo(TcpConnectionState.Terminated)); + }); + } + + [Test] + public async Task QueryAsync_CompressedAndTheServerStopsInsideABlock_ThrowsTimeoutNamingReadTimeout() + { + // A stalled compressed frame must time out through the underlying transport buffer. + // The decoder must propagate TimeoutException without reporting a malformed frame. + byte[] script = Concat( + await ServerHelloBytesAsync(54476), + await BytesAsync(w => + { + w.WriteVarUInt((ulong)ServerPacketType.Data); + w.WriteString(string.Empty); // The envelope is never framed; the frames begin after the table name. + })); + var transport = new ScriptedDuplexStream(script, blockWhenExhausted: true); + using var connection = new ClickHouseTcpConnection(transport, socket: null, Lz4Compressor.Default, readTimeout: TimeSpan.FromMilliseconds(200)); + await connection.HandshakeAsync(Handshake, None); + + var thrown = Assert.CatchAsync(async () => await DrainAsync(connection)); + + Assert.Multiple(() => + { + Assert.That(thrown.Message, Does.Contain("ReadTimeout")); + Assert.That(connection.State, Is.EqualTo(TcpConnectionState.Terminated)); + }); + } + + [Test] + public async Task QueryAsync_CallerCancelsWhileTheDeadlineIsArmed_ReportsCancellationNotTimeout() + { + // Use a long read timeout so the caller token triggers cancellation first. + var transport = new ScriptedDuplexStream(await ServerHelloBytesAsync(54476), blockWhenExhausted: true); + using var connection = new ClickHouseTcpConnection(transport, socket: null, readTimeout: TimeSpan.FromSeconds(30)); + await connection.HandshakeAsync(Handshake, None); + + using var cts = new CancellationTokenSource(); + Task drain = DrainAsync(connection, cts.Token); + await cts.CancelAsync(); + + Assert.CatchAsync(async () => await drain); + } + + [Test] + public async Task HandshakeAsync_SlowerThanReadTimeout_CompletesBecauseTheDeadlineCoversResponsesOnly() + { + // The handshake uses DialTimeout; ReadTimeout must remain inactive during connection establishment. + var transport = new ScriptedDuplexStream(await ServerHelloBytesAsync(54476), maxChunk: 2, readDelay: TimeSpan.FromMilliseconds(20)); + using var connection = new ClickHouseTcpConnection(transport, socket: null, readTimeout: TimeSpan.FromMilliseconds(50)); + + await connection.HandshakeAsync(Handshake, None); + + Assert.That(connection.State, Is.EqualTo(TcpConnectionState.Ready)); + } + + [Test] + public async Task QueryAsync_SecondQueryOnTheSameConnection_RearmsTheDeadlineForItsOwnReads() + { + // Verify that successive queries create separate deadline token sources on the same connection. + byte[] script = Concat( + await ServerHelloBytesAsync(54476), + await DataPacketAsync(new ulong[] { 1 }), + EndOfStreamPacket(), + await DataPacketAsync(new ulong[] { 2 }), + EndOfStreamPacket()); + using var connection = new ClickHouseTcpConnection(new ScriptedDuplexStream(script, maxChunk: 1), socket: null, readTimeout: TimeSpan.FromMilliseconds(200)); + await connection.HandshakeAsync(Handshake, None); + + List first = await MaterializeAsync(connection); + await Task.Delay(TimeSpan.FromMilliseconds(300)); + List second = await MaterializeAsync(connection); + + Assert.Multiple(() => + { + CollectionAssert.AreEqual(new ulong[] { 1 }, first[0]); + CollectionAssert.AreEqual(new ulong[] { 2 }, second[0]); + Assert.That(connection.State, Is.EqualTo(TcpConnectionState.Ready)); + }); + } + + [Test] + public async Task QueryAsync_ResponseSlowerOverallThanReadTimeout_CompletesBecauseTheDeadlineMeasuresSilence() + { + // Two-byte reads delayed by 20 ms make the query exceed 250 ms in total while each read stays below + // the timeout. Only reads after the handshake use ReadTimeout. + byte[] script = Concat( + await ServerHelloBytesAsync(54476), + await DataPacketAsync(new ulong[] { 1, 2, 3 }), + EndOfStreamPacket()); + var transport = new ScriptedDuplexStream(script, maxChunk: 2, readDelay: TimeSpan.FromMilliseconds(20)); + using var connection = new ClickHouseTcpConnection(transport, socket: null, readTimeout: TimeSpan.FromMilliseconds(250)); + await connection.HandshakeAsync(Handshake, None); + + var rows = await MaterializeAsync(connection); + + Assert.Multiple(() => + { + Assert.That(rows, Has.Count.EqualTo(1)); + CollectionAssert.AreEqual(new ulong[] { 1, 2, 3 }, rows[0]); + Assert.That(connection.State, Is.EqualTo(TcpConnectionState.Ready)); + }); + } + + [Test] + public async Task QueryAsync_ConsumerHoldsABlockPastReadTimeout_IsNotTreatedAsASilentServer() + { + // The read timer must be disarmed while the consumer holds a yielded block. + // Limit reads to one byte so the handshake cannot prefetch the entire response and bypass timed reads. + byte[] script = Concat( + await ServerHelloBytesAsync(54476), + await DataPacketAsync(new ulong[] { 1 }), + await DataPacketAsync(new ulong[] { 2 }), + EndOfStreamPacket()); + using var connection = new ClickHouseTcpConnection(new ScriptedDuplexStream(script, maxChunk: 1), socket: null, readTimeout: TimeSpan.FromMilliseconds(150)); + await connection.HandshakeAsync(Handshake, None); + + var blocks = 0; + await foreach (Block block in connection.QueryAsync("SELECT 1", cancellationToken: None)) + { + _ = block; + blocks++; + await Task.Delay(TimeSpan.FromMilliseconds(300)); + } + + Assert.Multiple(() => + { + Assert.That(blocks, Is.EqualTo(2)); + Assert.That(connection.State, Is.EqualTo(TcpConnectionState.Ready)); + }); + } + [Test] public async Task QueryAsync_TokenAlreadyCancelled_ThrowsWithoutClaimingConnection() { @@ -422,6 +617,10 @@ public async Task QueryAsync_AfterTerminate_ThrowsObjectDisposed() Assert.ThrowsAsync(async () => await DrainAsync(connection)); } + // Cancel is a one-byte packet type with no body; it must be the last byte written. + private static void AssertCancelSent(ScriptedDuplexStream transport) + => Assert.That(transport.Written[^1], Is.EqualTo((byte)ClientPacketType.Cancel), "the Cancel packet should be the last thing written"); + private static async Task ConnectedAsync(byte[] script) { var connection = new ClickHouseTcpConnection(new ScriptedDuplexStream(script), socket: null); @@ -443,9 +642,9 @@ private static async Task> MaterializeAsync(ClickHouseTcpConnectio } // Enumerates the response without reading block contents (for tests that assert an exception or state). - private static async Task DrainAsync(ClickHouseTcpConnection connection) + private static async Task DrainAsync(ClickHouseTcpConnection connection, CancellationToken cancellationToken = default) { - await foreach (Block block in connection.QueryAsync("SELECT 1", cancellationToken: None)) + await foreach (Block block in connection.QueryAsync("SELECT 1", cancellationToken: cancellationToken)) { _ = block; } diff --git a/ClickHouse.Driver.Tcp.Tests/Protocol/ClickHouseTcpConnectionTests.cs b/ClickHouse.Driver.Tcp.Tests/Protocol/ClickHouseTcpConnectionTests.cs index 544889ce4..3226e43e5 100644 --- a/ClickHouse.Driver.Tcp.Tests/Protocol/ClickHouseTcpConnectionTests.cs +++ b/ClickHouse.Driver.Tcp.Tests/Protocol/ClickHouseTcpConnectionTests.cs @@ -173,6 +173,24 @@ public async Task PingAsync_CancelledWhileAwaitingPong_TerminatesConnection() Assert.That(connection.State, Is.EqualTo(TcpConnectionState.Terminated)); } + [Test] + public async Task PingAsync_ServerNeverAnswersWithinReadTimeout_ThrowsTimeoutAndTerminates() + { + // Verify that ReadTimeout applies while waiting for Pong. + byte[] script = await ServerHelloBytesAsync(54476); + using var connection = new ClickHouseTcpConnection( + new ScriptedDuplexStream(script, blockWhenExhausted: true), socket: null, readTimeout: TimeSpan.FromMilliseconds(200)); + await connection.HandshakeAsync(Handshake, None); + + var thrown = Assert.CatchAsync(async () => await connection.PingAsync(None)); + + Assert.Multiple(() => + { + Assert.That(thrown.Message, Does.Contain("ReadTimeout")); + Assert.That(connection.State, Is.EqualTo(TcpConnectionState.Terminated)); + }); + } + [Test] public async Task PingAsync_AfterTerminate_ThrowsObjectDisposed() { diff --git a/ClickHouse.Driver.Tcp.Tests/Protocol/ReadBufferDeadlineTests.cs b/ClickHouse.Driver.Tcp.Tests/Protocol/ReadBufferDeadlineTests.cs new file mode 100644 index 000000000..09907a227 --- /dev/null +++ b/ClickHouse.Driver.Tcp.Tests/Protocol/ReadBufferDeadlineTests.cs @@ -0,0 +1,134 @@ +using System; +using System.IO; +using System.Threading; +using System.Threading.Tasks; +using System.Threading.Tasks.Sources; +using ClickHouse.Driver.Tcp.Protocol; + +namespace ClickHouse.Driver.Tcp.Tests.Protocol; + +[TestFixture] +public class ReadBufferDeadlineTests +{ + [TestCase(false)] + [TestCase(true)] + public async Task ReadIntoAsync_TimeoutAfterSuccessfulRead_NextReadUsesFreshDeadline(bool cancelCaller) + { + using var caller = new CancellationTokenSource(); + var deadline = new IdleReadDeadline(TimeSpan.FromMilliseconds(200)); + deadline.Begin(caller.Token); + using var stream = new DelayedContinuationStream(cancelCaller ? caller.Cancel : null); + using var buffer = new ReadBuffer(stream, deadline: deadline); + try + { + var first = new byte[1]; + Task read = buffer.ReadIntoAsync(first, caller.Token).AsTask(); + + // The first read succeeds, but its continuation is held until the actual timer cancels its token. + // This forces the completion/timeout race without relying on thread-pool scheduling. + await stream.TimeoutObserved.WaitAsync(TimeSpan.FromSeconds(10)); + stream.ResumeContinuation(); + await read; + Assert.That(first[0], Is.EqualTo(42)); + + if (cancelCaller) + { + Task nextRead = buffer.ReadIntoAsync(new byte[1], caller.Token).AsTask(); + Assert.CatchAsync(async () => + await nextRead.WaitAsync(TimeSpan.FromSeconds(10))); + } + else + { + var second = new byte[1]; + await buffer.ReadIntoAsync(second, caller.Token); + Assert.That(second[0], Is.EqualTo(43)); + + // The replacement must still enforce ReadTimeout when the next transport read stalls. + var timeout = Assert.ThrowsAsync(async () => + await buffer.ReadIntoAsync(new byte[1], caller.Token).AsTask().WaitAsync(TimeSpan.FromSeconds(10))); + Assert.That(timeout.Message, Does.Contain("ReadTimeout")); + } + } + finally + { + deadline.End(); + } + } + + // Completes the first read successfully and holds its await continuation until the test releases it. + private sealed class DelayedContinuationStream : MemoryStream, IValueTaskSource + { + private readonly Action cancelSecondRead; + private readonly TaskCompletionSource timeoutObserved = new(TaskCreationOptions.RunContinuationsAsynchronously); + private CancellationTokenRegistration registration; + private Action continuation; + private object continuationState; + private bool completed; + private int reads; + private int result; + + internal DelayedContinuationStream(Action cancelSecondRead) + : base(new byte[] { 42, 43 }) + { + this.cancelSecondRead = cancelSecondRead; + } + + internal Task TimeoutObserved => timeoutObserved.Task; + + internal void ResumeContinuation() => continuation(continuationState); + + public override ValueTask ReadAsync(Memory buffer, CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + if (++reads == 1) + { + result = Read(buffer.Span); + registration = cancellationToken.Register(() => timeoutObserved.TrySetResult()); + return new ValueTask(this, 0); + } + + if (reads == 2 && cancelSecondRead is null) + { + return base.ReadAsync(buffer, cancellationToken); + } + + Task pending = WaitForCancellationAsync(cancellationToken); + if (reads == 2) + { + // Cancel only after the pending read has registered on its deadline token. + cancelSecondRead(); + Assert.That(cancellationToken.IsCancellationRequested, Is.True, "Caller cancellation must reach the replacement read token synchronously."); + } + + return new ValueTask(pending); + } + + public ValueTaskSourceStatus GetStatus(short token) + => completed ? ValueTaskSourceStatus.Succeeded : ValueTaskSourceStatus.Pending; + + public int GetResult(short token) => result; + + public void OnCompleted(Action callback, object state, short token, ValueTaskSourceOnCompletedFlags flags) + { + continuation = callback; + continuationState = state; + completed = true; + } + + protected override void Dispose(bool disposing) + { + if (disposing) + { + registration.Dispose(); + } + + base.Dispose(disposing); + } + + private static async Task WaitForCancellationAsync(CancellationToken cancellationToken) + { + await Task.Delay(Timeout.InfiniteTimeSpan, cancellationToken); + return 0; + } + } +} diff --git a/ClickHouse.Driver.Tcp.Tests/Utilities/QueryLog.cs b/ClickHouse.Driver.Tcp.Tests/Utilities/QueryLog.cs new file mode 100644 index 000000000..bf31e1b1a --- /dev/null +++ b/ClickHouse.Driver.Tcp.Tests/Utilities/QueryLog.cs @@ -0,0 +1,68 @@ +using System; +using System.Threading; +using System.Threading.Tasks; +using ClickHouse.Driver.Tcp.Client; + +namespace ClickHouse.Driver.Tcp.Tests.Utilities; + +/// Reads system.query_log through the TCP client, retrying until the record is available. +/// +/// Query log records may be queued after the response reaches the client, so a single flush can miss them. +/// Retries allow cancelled queries time to stop and produce a log record. Only query_log is flushed +/// because the framework suites share one server. +/// +internal static class QueryLog +{ + /// Number of flush-and-read attempts before giving up. + internal const int MaxAttempts = 40; + + private static readonly TimeSpan RetryDelay = TimeSpan.FromMilliseconds(100); + + /// + /// Returns the first column of the first matching row, retrying until a non-null value is available. + /// Fails the test when the retry limit is reached. + /// + /// + /// A missing row and a NULL value are indistinguishable here, so select an expression that is never NULL for + /// a row that exists. + /// + /// Client to run the flush and the lookup on. + /// Lookup returning one row with the value under test in its first column. + /// The value read once the row became visible. + internal static async Task ScalarAsync(ClickHouseTcpClient client, string sql) + { + for (var attempt = 1; attempt <= MaxAttempts; attempt++) + { + await client.ExecuteAsync("SYSTEM FLUSH LOGS query_log", cancellationToken: CancellationToken.None); + + object value = await ReadFirstAsync(client, sql); + if (value is not null) + { + return value; + } + + if (attempt < MaxAttempts) + { + await Task.Delay(RetryDelay); + } + } + + string message = $"No system.query_log row appeared after {MaxAttempts} flush attempts, so the value under test could not be determined. Query: {sql}"; + Assert.Fail(message); + + // Assert.Fail can return inside Assert.Multiple; throw to prevent returning a missing value. + throw new InvalidOperationException(message); + } + + // Consume the full response to keep the connection reusable between lookups. + private static async Task ReadFirstAsync(ClickHouseTcpClient client, string sql) + { + object first = null; + await foreach (object[] row in client.QueryAsync(sql, cancellationToken: CancellationToken.None)) + { + first ??= row[0]; + } + + return first; + } +} diff --git a/ClickHouse.Driver.Tcp.Tests/Utilities/ScriptedDuplexStream.cs b/ClickHouse.Driver.Tcp.Tests/Utilities/ScriptedDuplexStream.cs index 2210830ee..b05521215 100644 --- a/ClickHouse.Driver.Tcp.Tests/Utilities/ScriptedDuplexStream.cs +++ b/ClickHouse.Driver.Tcp.Tests/Utilities/ScriptedDuplexStream.cs @@ -16,14 +16,16 @@ internal sealed class ScriptedDuplexStream : Stream private readonly byte[] script; private readonly int maxChunk; private readonly bool blockWhenExhausted; + private readonly TimeSpan readDelay; private readonly MemoryStream sink = new(); private int position; - public ScriptedDuplexStream(byte[] script, int maxChunk = int.MaxValue, bool blockWhenExhausted = false) + public ScriptedDuplexStream(byte[] script, int maxChunk = int.MaxValue, bool blockWhenExhausted = false, TimeSpan readDelay = default) { this.script = script; this.maxChunk = maxChunk < 1 ? 1 : maxChunk; this.blockWhenExhausted = blockWhenExhausted; + this.readDelay = readDelay; } /// The bytes the connection has written (the client → server side of the exchange). @@ -70,6 +72,13 @@ public override async ValueTask ReadAsync(Memory buffer, Cancellation await Task.Delay(Timeout.Infinite, cancellationToken).ConfigureAwait(false); } + // Combine delayed reads with maxChunk to test responses whose total duration exceeds ReadTimeout + // while each read completes within it. + if (readDelay > TimeSpan.Zero) + { + await Task.Delay(readDelay, cancellationToken).ConfigureAwait(false); + } + return Read(buffer.Span); } diff --git a/ClickHouse.Driver.Tcp/Client/ClickHouseTcpClientOptions.cs b/ClickHouse.Driver.Tcp/Client/ClickHouseTcpClientOptions.cs index 31d34ab16..3eaaa1d02 100644 --- a/ClickHouse.Driver.Tcp/Client/ClickHouseTcpClientOptions.cs +++ b/ClickHouse.Driver.Tcp/Client/ClickHouseTcpClientOptions.cs @@ -175,10 +175,15 @@ public sealed record ClickHouseTcpClientOptions public TimeSpan DialTimeout { get; init; } = DefaultDialTimeout; /// - /// The idle deadline for reading a response — reset each time a packet arrives — so a long streaming query - /// is not killed for taking a long time overall. Defaults to 300s. Stored but not yet enforced; the - /// idle-deadline read loop lands in a later change. + /// Maximum time to wait for a transport read during an operation. On expiry, the operation throws + /// and discards the connection. Defaults to 300 seconds; + /// disables this timeout. /// + /// + /// The timer starts before each transport read and stops when that read completes. Total query duration + /// and time spent processing a returned block are unrestricted by this timeout. + /// Connection establishment, including the handshake, uses . + /// public TimeSpan ReadTimeout { get; init; } = DefaultReadTimeout; /// @@ -447,7 +452,16 @@ internal void Validate() RequireUsableTimeout(DialTimeout, nameof(DialTimeout)); - RequireUsableTimeout(ReadTimeout, nameof(ReadTimeout)); + // Zero disables the read timeout. + if (ReadTimeout < TimeSpan.Zero) + { + throw new ArgumentOutOfRangeException(nameof(ReadTimeout), ReadTimeout, "ReadTimeout must not be negative; use TimeSpan.Zero to disable the deadline."); + } + + if (ReadTimeout.TotalMilliseconds > int.MaxValue) + { + throw new ArgumentOutOfRangeException(nameof(ReadTimeout), ReadTimeout, $"ReadTimeout must not exceed {TimeSpan.FromMilliseconds(int.MaxValue)} (about 24.8 days)."); + } if (MaxSendBufferBytes <= 0) { diff --git a/ClickHouse.Driver.Tcp/Client/ClickHouseTcpConnectionStringBuilder.cs b/ClickHouse.Driver.Tcp/Client/ClickHouseTcpConnectionStringBuilder.cs index 577229963..c30ae981f 100644 --- a/ClickHouse.Driver.Tcp/Client/ClickHouseTcpConnectionStringBuilder.cs +++ b/ClickHouse.Driver.Tcp/Client/ClickHouseTcpConnectionStringBuilder.cs @@ -171,7 +171,7 @@ public TimeSpan DialTimeout set => this["DialTimeout"] = value.TotalSeconds; } - /// The idle read deadline, in seconds. Defaults to 300. + /// Maximum time per transport read during an operation, in seconds. Defaults to 300; 0 disables it. public TimeSpan ReadTimeout { get => GetTimeSpanSecondsOrDefault("ReadTimeout", ClickHouseTcpClientOptions.DefaultReadTimeout); diff --git a/ClickHouse.Driver.Tcp/Client/IConnectionFactory.cs b/ClickHouse.Driver.Tcp/Client/IConnectionFactory.cs index 835bcb6e3..c10605f0e 100644 --- a/ClickHouse.Driver.Tcp/Client/IConnectionFactory.cs +++ b/ClickHouse.Driver.Tcp/Client/IConnectionFactory.cs @@ -68,7 +68,7 @@ public async ValueTask CreateAsync(CancellationToken ca try { ClickHouseTcpConnection connection = await ClickHouseTcpConnection.ConnectAsync( - options.Host, options.ResolvedPort, options.ToHandshakeParameters(), tls, linked.Token, options.Compressor).ConfigureAwait(false); + options.Host, options.ResolvedPort, options.ToHandshakeParameters(), tls, linked.Token, options.Compressor, options.ReadTimeout).ConfigureAwait(false); activity?.SetSuccess(); if (logger is not null) diff --git a/ClickHouse.Driver.Tcp/Protocol/ClickHouseTcpConnection.cs b/ClickHouse.Driver.Tcp/Protocol/ClickHouseTcpConnection.cs index 5ff3ac3d1..abc3f8c47 100644 --- a/ClickHouse.Driver.Tcp/Protocol/ClickHouseTcpConnection.cs +++ b/ClickHouse.Driver.Tcp/Protocol/ClickHouseTcpConnection.cs @@ -58,11 +58,17 @@ internal sealed class ClickHouseTcpConnection : IDisposable, IAsyncDisposable // for timezone-less DateTime/DateTime64 result columns. private const string SessionTimezoneSetting = "session_timezone"; + // Limits the Cancel flush before closing the connection and releasing its pool lease. + private static readonly TimeSpan CancelSendTimeout = TimeSpan.FromSeconds(2); + private readonly Socket socket; private readonly Stream stream; private readonly ClickHouseBinaryReader reader; private readonly ClickHouseBinaryWriter writer; + // Timeout per transport read, or null when disabled. + private readonly IdleReadDeadline readDeadline; + // Null means every query on this connection is uncompressed. Compression is per-query on the wire, but the // codec is a client-level option today, so it is fixed for a connection's life; a per-query override would // move this to the operation entry points. @@ -85,12 +91,20 @@ internal sealed class ClickHouseTcpConnection : IDisposable, IAsyncDisposable /// The duplex transport stream (a network stream in production). /// The underlying socket, closed on termination; null when the stream owns teardown. /// Frame codec for this connection's queries, or null to run them uncompressed. - internal ClickHouseTcpConnection(Stream stream, Socket socket, IClickHouseCompressor compressor = null) + /// + /// Timeout per transport read during an operation. disables it; + /// scripted streams use this default. + /// + internal ClickHouseTcpConnection(Stream stream, Socket socket, IClickHouseCompressor compressor = null, TimeSpan readTimeout = default) { this.stream = stream; this.socket = socket; this.compressor = compressor; - reader = new ClickHouseBinaryReader(stream); + readDeadline = readTimeout == TimeSpan.Zero ? null : new IdleReadDeadline(readTimeout); + + // Attach the deadline to the transport buffer. The frame decoder reads through this buffer, + // so compressed reads use the same timeout. + reader = new ClickHouseBinaryReader(new ReadBuffer(stream, deadline: readDeadline), ownsBuffer: true); writer = new ClickHouseBinaryWriter(stream); state = TcpConnectionState.Handshaking; } @@ -122,9 +136,7 @@ internal ClickHouseTcpConnection(Stream stream, Socket socket, IClickHouseCompre /// up, which on Linux takes about fifteen minutes. That is inherent to a client-side check, so the pool does not /// rely on this alone: it also refuses a connection that has sat idle past IdleTimeout, which covers the /// common case of an intermediary dropping a connection nobody was using. Neither catches a drop that strikes a - /// connection in active use. The answer to that is an idle read deadline rather than a stricter probe, and - /// that deadline does not exist yet: ReadTimeout is parsed and stored but nothing enforces it, so a - /// caller's own is currently the only bound on such a stall. + /// connection in active use; ReadTimeout bounds each transport read during an operation. /// /// internal bool IsReusable @@ -217,7 +229,8 @@ public static async ValueTask ConnectAsync( ClientHandshakeParameters handshake, TlsParameters tls, CancellationToken cancellationToken, - IClickHouseCompressor compressor = null) + IClickHouseCompressor compressor = null, + TimeSpan readTimeout = default) { ArgumentNullException.ThrowIfNull(host); ArgumentNullException.ThrowIfNull(handshake); @@ -261,8 +274,9 @@ public static async ValueTask ConnectAsync( } // HandshakeAsync terminates the connection (closing this socket) on any failure, so a throw here needs - // no extra cleanup. - var connection = new ClickHouseTcpConnection(transport, socket, compressor); + // no extra cleanup. The handshake itself runs under the caller's connect deadline rather than + // readTimeout, so the two never stack on the one exchange. + var connection = new ClickHouseTcpConnection(transport, socket, compressor, readTimeout); await connection.HandshakeAsync(handshake, cancellationToken).ConfigureAwait(false); return connection; } @@ -278,54 +292,63 @@ public static async ValueTask ConnectAsync( /// The server replied with an Exception. /// The server replied with something other than Pong or Exception. /// The connection failed while the ping was in flight. + /// A transport read exceeded the connection's ReadTimeout. public async ValueTask PingAsync(CancellationToken cancellationToken) { cancellationToken.ThrowIfCancellationRequested(); BeginOperation(); - ServerPacketType reply; + BeginRead(cancellationToken); try { - writer.WriteClientPacketType(ClientPacketType.Ping); - await writer.FlushAsync(cancellationToken).ConfigureAwait(false); + ServerPacketType reply; + try + { + writer.WriteClientPacketType(ClientPacketType.Ping); + await writer.FlushAsync(cancellationToken).ConfigureAwait(false); - // A Ping is only ever sent on an idle connection, never mid-query, so no Progress or other - // interleaved packet can precede the reply — unlike a query response, which the read loop drains. - // A single read therefore suffices; anything but Pong or a (complete) Exception is a violation. - reply = await reader.ReadServerPacketTypeAsync(cancellationToken).ConfigureAwait(false); - } - catch - { - // The failed/cancelled I/O has unwound, but the stream position is unknown; discard the connection. - Terminate(); - throw; - } + // A Ping is only ever sent on an idle connection, never mid-query, so no Progress or other + // interleaved packet can precede the reply — unlike a query response, which the read loop drains. + // A single read therefore suffices; anything but Pong or a (complete) Exception is a violation. + reply = await reader.ReadServerPacketTypeAsync(cancellationToken).ConfigureAwait(false); + } + catch + { + // The failed/cancelled I/O has unwound, but the stream position is unknown; discard the connection. + Terminate(); + throw; + } - switch (reply) - { - case ServerPacketType.Pong: - state = TcpConnectionState.Ready; - return; + switch (reply) + { + case ServerPacketType.Pong: + state = TcpConnectionState.Ready; + return; - case ServerPacketType.Exception: - ClickHouseTcpServerException exception; - try - { - exception = await ClickHouseTcpServerException.ReadAsync(reader, cancellationToken).ConfigureAwait(false); - } - catch - { - Terminate(); - throw; - } + case ServerPacketType.Exception: + ClickHouseTcpServerException exception; + try + { + exception = await ClickHouseTcpServerException.ReadAsync(reader, cancellationToken).ConfigureAwait(false); + } + catch + { + Terminate(); + throw; + } - Terminate(); - throw exception; + Terminate(); + throw exception; - default: - Terminate(); - throw new ClickHouseTcpProtocolException( - $"Unexpected packet type {reply} ({(ulong)reply}) in response to Ping; expected Pong or Exception."); + default: + Terminate(); + throw new ClickHouseTcpProtocolException( + $"Unexpected packet type {reply} ({(ulong)reply}) in response to Ping; expected Pong or Exception."); + } + } + finally + { + EndRead(); } } @@ -357,6 +380,10 @@ public async ValueTask PingAsync(CancellationToken cancellationToken) /// } /// /// + /// + /// Dispose the enumerator to release the connection, the current block's buffers, and the read deadline's + /// cancellation registration. Stopping enumeration without disposal skips this cleanup. + /// /// /// The SQL text. /// Per-query settings as textual values, or null for none. @@ -371,6 +398,7 @@ public async ValueTask PingAsync(CancellationToken cancellationToken) /// The server reported an error while executing the query. /// The server sent an unexpected packet. /// The connection failed while the response was being read. + /// A transport read exceeded the connection's ReadTimeout. /// was cancelled. internal async IAsyncEnumerable QueryAsync( string sql, @@ -391,7 +419,9 @@ internal async IAsyncEnumerable QueryAsync( NegotiatedProtocol negotiated = server.Negotiated; ClickHouseTcpServerException pending = null; Block current = null; - bool completed = false; + bool responseCompleted = false; + bool reusable = false; + bool flushedWholePackets = false; // Encode the Query packet into the write buffer before any of it reaches the socket. A failure here is a // client-side error (e.g. parameters on a protocol revision that predates them): nothing has been sent, @@ -407,6 +437,7 @@ internal async IAsyncEnumerable QueryAsync( throw; } + BeginRead(cancellationToken); try { // The end-of-input marker is written here rather than above, because framing it is not buffer-only @@ -416,6 +447,9 @@ internal async IAsyncEnumerable QueryAsync( await WriteEndOfInputBlockAsync(cancellationToken).ConfigureAwait(false); await writer.FlushAsync(cancellationToken).ConfigureAwait(false); + // The request is fully flushed. Cancel can now be sent at a packet boundary if the response is incomplete. + flushedWholePackets = true; + while (true) { // Resuming here means the consumer has advanced past the previously yielded block, so its @@ -430,13 +464,15 @@ internal async IAsyncEnumerable QueryAsync( if (packet == ServerPacketType.EndOfStream) { - completed = true; + responseCompleted = true; + reusable = true; break; } if (packet == ServerPacketType.Exception) { pending = await ClickHouseTcpServerException.ReadAsync(reader, cancellationToken).ConfigureAwait(false); + responseCompleted = true; break; } @@ -464,15 +500,24 @@ internal async IAsyncEnumerable QueryAsync( } finally { + // Release the deadline registration before any subsequent cleanup can throw. + EndRead(); + // Release the last yielded block (still current) on end-of-stream, early disposal, or error. current?.Dispose(); - if (completed) + if (reusable) { state = TcpConnectionState.Ready; } else { + // A response that has not reached a terminal packet may still be running on the server. + if (!responseCompleted) + { + await TrySendCancelAsync(flushedWholePackets).ConfigureAwait(false); + } + Terminate(); } } @@ -536,6 +581,7 @@ internal async IAsyncEnumerable QueryAsync( /// The server reported an error while executing the insert. /// The connection failed while the blocks were being sent or the response read. /// The server sent an unexpected packet, or no schema block. + /// A transport read exceeded the connection's ReadTimeout. /// was cancelled. internal ValueTask InsertAsync( string sql, @@ -614,8 +660,11 @@ private async ValueTask InsertCoreAsync( Exception buildFailure = null; IReadOnlyList values = null; IInsertColumnSource source = null; - bool completed = false; + bool responseCompleted = false; + bool reusable = false; + bool flushedWholePackets = false; string mismatchError = null; + BeginRead(cancellationToken); try { // The empty end-of-input block must follow the Query: the server waits for it before sending the @@ -623,11 +672,13 @@ private async ValueTask InsertCoreAsync( Query.Write(writer, negotiated, clientMetadata, queryId, sql, settings, parameters, compressor is not null); await WriteEndOfInputBlockAsync(cancellationToken).ConfigureAwait(false); await writer.FlushAsync(cancellationToken).ConfigureAwait(false); + flushedWholePackets = true; // Drain metadata until the schema block (the first Data packet) or a terminal packet. (Block schema, ClickHouseTcpServerException error) = await ReadToNextDataBlockAsync(negotiated, readContext, telemetry, callbacks, cancellationToken).ConfigureAwait(false); if (schema is null) { + responseCompleted = true; if (error is null) { // Clean end-of-stream with no schema: the server never opened the row-stream phase (e.g. @@ -637,7 +688,7 @@ private async ValueTask InsertCoreAsync( } // The Exception packet does not say whether the server accepted the query and returned to its - // request loop, so leave completed false and retire the connection in the finally below. + // request loop, so leave reusable false and retire the connection in the finally below. pending = error; } else @@ -675,29 +726,41 @@ private async ValueTask InsertCoreAsync( } } - // Always run, even after a factory failure: the row stream has to be closed for the server to - // finish the insert. A gather failure is deferred the same way, so the stream still closes cleanly. + // Disable Cancel while streaming rows: a failed write may leave a partial Data packet. + // Always close the row stream, including after factory or gather failures. A successful return + // restores the packet boundary and allows cancellation while reading the acknowledgement. + flushedWholePackets = false; Exception gatherFailure = await StreamInsertRowsAsync(plan, source, rowCount, maxRowsPerBlock, maxSendBufferBytes, negotiated, cancellationToken).ConfigureAwait(false); + flushedWholePackets = true; buildFailure ??= gatherFailure; // A clean acknowledgement leaves the connection reusable. A server Exception is parked for the // caller but retires the connection, because its packet does not prove the server will accept // another request. pending = await DrainToEndOfStreamAsync(negotiated, readContext, telemetry, callbacks, cancellationToken).ConfigureAwait(false); - completed = pending is null; + responseCompleted = true; + reusable = pending is null; } } finally { + // Release the deadline registration before any subsequent cleanup can throw. + EndRead(); + // Only the factory's source is ours to release; a caller's own columns outlive the insert. source?.Dispose(); - if (completed) + if (reusable) { state = TcpConnectionState.Ready; } else { + if (!responseCompleted) + { + await TrySendCancelAsync(flushedWholePackets).ConfigureAwait(false); + } + Terminate(); } } @@ -1394,4 +1457,50 @@ private void BeginOperation() $"The connection is busy ({state}); a single connection carries one in-flight operation at a time."); } } + + /// Initializes the operation's read deadline. Pair with in a finally block. + /// + /// Reads and writes receive the caller's token. The transport buffer applies the deadline token only + /// to the individual stream read, so a completed read's timeout cannot cancel later reads or writes. + /// + /// The caller's token for this operation. + private void BeginRead(CancellationToken cancellationToken) + => readDeadline?.Begin(cancellationToken); + + /// Closes the idle read deadline opened by . + private void EndRead() => readDeadline?.End(); + + /// + /// Attempts to send Cancel before closing the connection. Delivery failures are suppressed to preserve + /// the operation's original exception. + /// + /// + /// Whether the last successful flush ended at a packet boundary. Must be false after a partial or failed + /// flush: Cancel would be interpreted as packet data, and resetting the writer cannot retract sent bytes. + /// + /// A task that completes after the send attempt succeeds, fails, or times out. + private async ValueTask TrySendCancelAsync(bool flushedWholePackets) + { + // AbortTransport may have already closed the socket; skip the write in that case. + if (!flushedWholePackets || state == TcpConnectionState.Terminated) + { + return; + } + + try + { + // Discard anything the interrupted operation left buffered, so Cancel is the whole of what goes out. + writer.Reset(); + writer.WriteClientPacketType(ClientPacketType.Cancel); + + // Use an independent timeout because the operation token may already be cancelled. + // The pool lease remains held until this flush completes or times out. + using var deadline = new CancellationTokenSource(CancelSendTimeout); + await writer.FlushAsync(deadline.Token).ConfigureAwait(false); + } + catch (Exception e) when (e is not (OutOfMemoryException or StackOverflowException)) + { + // Ignore Cancel delivery failures; the caller closes the connection next. + } + } } diff --git a/ClickHouse.Driver.Tcp/Protocol/IdleReadDeadline.cs b/ClickHouse.Driver.Tcp/Protocol/IdleReadDeadline.cs new file mode 100644 index 000000000..6e322c5ea --- /dev/null +++ b/ClickHouse.Driver.Tcp/Protocol/IdleReadDeadline.cs @@ -0,0 +1,76 @@ +using System; +using System.Threading; + +namespace ClickHouse.Driver.Tcp.Protocol; + +/// Applies a timeout to each transport read within an operation. +/// +/// The timer is armed before each read and disarmed when it completes. It remains disarmed while the consumer +/// processes a returned block. Each operation must pair with . +/// This class is not thread-safe. +/// +internal sealed class IdleReadDeadline +{ + private readonly TimeSpan timeout; + private CancellationTokenSource source; + private CancellationToken callerToken; + + /// Initializes a deadline of . + /// + /// Positive read timeout, validated by against the timer's + /// supported range. Connections with a zero timeout do not create a deadline. + /// + internal IdleReadDeadline(TimeSpan timeout) => this.timeout = timeout; + + /// Whether the deadline token is cancelled without caller cancellation. + internal bool Elapsed => source is { IsCancellationRequested: true } && !callerToken.IsCancellationRequested; + + /// Initializes cancellation for an operation. Pair with in a finally block. + /// The caller's token for this operation. + internal void Begin(CancellationToken operationToken) + { + callerToken = operationToken; + source = CancellationTokenSource.CreateLinkedTokenSource(operationToken); + } + + /// Disposes the operation token source. Safe to call without a matching . + internal void End() + { + source?.Dispose(); + source = null; + callerToken = default; + } + + /// Arms the timer and returns the token for this transport read. + /// The caller's token, used while no operation is active (during the handshake). + /// The read's deadline token, or the caller's token when no operation is active. + internal CancellationToken Arm(CancellationToken cancellationToken) + { + if (source is null) + { + return cancellationToken; + } + + source.CancelAfter(timeout); + return source.Token; + } + + /// Disarms the timer and prepares the token source for the next read. + internal void Disarm() + { + if (source is null || source.TryReset()) + { + return; + } + + // TryReset fails if cancellation was requested or a timer callback was queued. Replace the source + // so that callback can only cancel the completed read's token. + source.Dispose(); + source = CancellationTokenSource.CreateLinkedTokenSource(callerToken); + } + + /// Creates the exception for a transport read timeout. + /// A timeout naming the option that set the deadline. + internal TimeoutException ToException() + => new($"The server sent nothing for {timeout.TotalSeconds:0.###}s while a response was being read (ReadTimeout)."); +} diff --git a/ClickHouse.Driver.Tcp/Protocol/ReadBuffer.cs b/ClickHouse.Driver.Tcp/Protocol/ReadBuffer.cs index b02d81631..be95780f0 100644 --- a/ClickHouse.Driver.Tcp/Protocol/ReadBuffer.cs +++ b/ClickHouse.Driver.Tcp/Protocol/ReadBuffer.cs @@ -27,6 +27,7 @@ internal sealed class ReadBuffer : IDisposable private readonly Stream stream; private readonly bool readsFromTransport; + private readonly IdleReadDeadline deadline; private byte[] buffer; private int capacity; private int head; // index of the first valid byte @@ -42,9 +43,13 @@ internal sealed class ReadBuffer : IDisposable /// Whether is the connection itself, so that a failed read is the transport failing. /// False for an adapter stream, whose own layer decides what its failures mean. /// + /// + /// Timeout for each transport read, or null to use only the supplied cancellation token. + /// Adapter buffers omit this deadline; their underlying transport buffer enforces it. + /// /// is null. /// is below . - public ReadBuffer(Stream stream, int capacity = 16384, bool readsFromTransport = true) + public ReadBuffer(Stream stream, int capacity = 16384, bool readsFromTransport = true, IdleReadDeadline deadline = null) { this.stream = stream ?? throw new ArgumentNullException(nameof(stream)); if (capacity < MaxContiguous) @@ -53,6 +58,7 @@ public ReadBuffer(Stream stream, int capacity = 16384, bool readsFromTransport = } this.readsFromTransport = readsFromTransport; + this.deadline = deadline; buffer = ArrayPool.Shared.Rent(capacity); this.capacity = buffer.Length; // Rent may return a larger array; use all of it. } @@ -67,6 +73,7 @@ public ReadBuffer(Stream stream, int capacity = 16384, bool readsFromTransport = /// The number of contiguous bytes that must be available; must not exceed . /// A token to observe for cancellation. /// exceeds the buffer capacity. + /// A transport read exceeded the configured timeout. /// The stream ended before enough bytes arrived, or the read failed. public async ValueTask EnsureAsync(int needed, CancellationToken cancellationToken) { @@ -137,6 +144,7 @@ public ReadOnlySpan ReadSpan(int count) /// /// The region to fill completely with consumed bytes. /// A token to observe for cancellation. + /// A transport read exceeded the configured timeout. /// The stream ended before the destination was filled, or the read failed. public async ValueTask ReadIntoAsync(Memory destination, CancellationToken cancellationToken) { @@ -150,21 +158,7 @@ public async ValueTask ReadIntoAsync(Memory destination, CancellationToken while (!destination.IsEmpty) { - int read; - try - { - read = await stream.ReadAsync(destination, cancellationToken).ConfigureAwait(false); - } - catch (Exception e) when (readsFromTransport && TransportFailure.IsTransportFailure(e)) - { - throw TransportFailure.Read(e); - } - - if (read == 0) - { - throw TransportFailure.EndOfStream(); - } - + int read = await ReadFromStreamAsync(destination, cancellationToken).ConfigureAwait(false); destination = destination.Slice(read); } } @@ -204,21 +198,46 @@ private void Advance(int n) private async ValueTask FillOnceAsync(CancellationToken cancellationToken) { int writeStart = head + buffered; + int read = await ReadFromStreamAsync(buffer.AsMemory(writeStart, capacity - writeStart), cancellationToken).ConfigureAwait(false); + buffered += read; + } + + /// + /// Performs one stream read with the configured timeout and translates timeout, EOF, and transport failures. + /// + /// Where to put the bytes; filled in part or in whole. + /// A token to observe for cancellation. + /// The number of bytes read, always at least one. + /// A transport read exceeded the configured timeout. + /// The stream ended, or the read failed. + private async ValueTask ReadFromStreamAsync(Memory destination, CancellationToken cancellationToken) + { int read; + CancellationToken readToken = deadline?.Arm(cancellationToken) ?? cancellationToken; try { - read = await stream.ReadAsync(buffer.AsMemory(writeStart, capacity - writeStart), cancellationToken).ConfigureAwait(false); + read = await stream.ReadAsync(destination, readToken).ConfigureAwait(false); + } + catch (Exception e) when (deadline is { Elapsed: true } && (e is OperationCanceledException || TransportFailure.IsTransportFailure(e))) + { + // Translate timeout cancellation to TimeoutException, including transport exceptions raised during + // cancellation (for example, by SslStream). Caller cancellation does not count as a timeout. + throw deadline.ToException(); } catch (Exception e) when (readsFromTransport && TransportFailure.IsTransportFailure(e)) { throw TransportFailure.Read(e); } + finally + { + deadline?.Disarm(); + } if (read == 0) { throw TransportFailure.EndOfStream(); } - buffered += read; + return read; } }