From 4d2937522da904febfe8ca9be77a6c1d7c62f0ae Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Tue, 13 Oct 2026 03:46:20 +0000 Subject: [PATCH 1/2] fix(csharp): drain active status polls before disposal Signed-off-by: Vu Anh Phung --- .../Reader/DatabricksOperationStatusPoller.cs | 4 +- .../DatabricksOperationStatusPollerTests.cs | 150 ++++++++++++++++++ 2 files changed, 153 insertions(+), 1 deletion(-) diff --git a/csharp/src/Reader/DatabricksOperationStatusPoller.cs b/csharp/src/Reader/DatabricksOperationStatusPoller.cs index 2b6170ce4..a56b8ffb3 100644 --- a/csharp/src/Reader/DatabricksOperationStatusPoller.cs +++ b/csharp/src/Reader/DatabricksOperationStatusPoller.cs @@ -135,7 +135,9 @@ async Task PollOnceAsync(CancellationToken ct, Activity? activity) TOperationHandle? operationHandle = _response.OperationHandle; if (operationHandle == null) return; - using CancellationTokenSource timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(ct); + ct.ThrowIfCancellationRequested(); + // Stopping must not interrupt serialization into Thrift's shared transport buffer. + using CancellationTokenSource timeoutCts = new CancellationTokenSource(); timeoutCts.CancelAfter(TimeSpan.FromSeconds(_requestTimeoutSeconds)); TGetOperationStatusReq request = new TGetOperationStatusReq(operationHandle); diff --git a/csharp/test/Unit/DatabricksOperationStatusPollerTests.cs b/csharp/test/Unit/DatabricksOperationStatusPollerTests.cs index 67551105a..459d6a9b7 100644 --- a/csharp/test/Unit/DatabricksOperationStatusPollerTests.cs +++ b/csharp/test/Unit/DatabricksOperationStatusPollerTests.cs @@ -22,6 +22,9 @@ */ using System; +using System.IO; +using System.Net; +using System.Net.Http; using System.Threading; using System.Threading.Tasks; using AdbcDrivers.Databricks.Reader; @@ -30,6 +33,10 @@ using AdbcDrivers.HiveServer2.Hive2; using Apache.Hive.Service.Rpc.Thrift; using Moq; +using Thrift; +using Thrift.Protocol; +using Thrift.Transport; +using Thrift.Transport.Client; using Xunit; using Xunit.Abstractions; @@ -309,5 +316,148 @@ public async Task UsesCustomRequestTimeout() Assert.True(pollCount > 0, "Should have polled at least once"); _mockClient.Verify(c => c.GetOperationStatus(It.IsAny(), It.IsAny()), Times.AtLeastOnce); } + + [Theory] + [InlineData(1, false)] + [InlineData(4, false)] + [InlineData(10, false)] + [InlineData(20, false)] + [InlineData(1, true)] + [InlineData(4, true)] + [InlineData(10, true)] + [InlineData(20, true)] + public async Task StoppingDuringSerializationPreservesCloseOperation(int pauseAfterWrite, bool cancelExternally) + { + _operationHandle.OperationId.Secret = new byte[16]; + _operationHandle.HasResultSet = true; + var server = new Mock(MockBehavior.Strict); + server.Setup(c => c.GetOperationStatus(It.IsAny(), It.IsAny())) + .ReturnsAsync(new TGetOperationStatusResp + { + Status = new TStatus(TStatusCode.SUCCESS_STATUS), + OperationState = TOperationState.FINISHED_STATE + }); + server.Setup(c => c.CloseOperation(It.IsAny(), It.IsAny())) + .ReturnsAsync(new TCloseOperationResp { Status = new TStatus(TStatusCode.SUCCESS_STATUS) }); + + using var httpClient = new HttpClient(new ThriftHandler(server.Object)) + { + BaseAddress = new Uri("http://localhost/") + }; + using var transport = new PausingHttpTransport(httpClient, pauseAfterWrite); + using var client = new ThreadSafeClient(new TCLIService.Client(new TBinaryProtocol(transport))); + _mockStatement.Setup(s => s.Client).Returns(client); + using var cancellation = new CancellationTokenSource(); + using var poller = new DatabricksOperationStatusPoller(_mockStatement.Object, _mockResponse.Object); + + poller.Start(cancellation.Token); + try + { + await WaitForSignal(transport.Paused.Task); + if (cancelExternally) + { + cancellation.Cancel(); + } + else + { + poller.Stop(); + } + } + finally + { + transport.Resume.TrySetResult(true); + poller.Dispose(); + } + + await client.CloseOperation(new TCloseOperationReq(_operationHandle)); + + server.Verify(c => c.GetOperationStatus(It.IsAny(), It.IsAny()), Times.Once); + server.Verify(c => c.CloseOperation(It.IsAny(), It.IsAny()), Times.Once); + } + + [Fact] + public async Task StoppedPollStillHonorsRequestTimeout() + { + var started = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var timedOut = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + _mockClient.Setup(c => c.GetOperationStatus(It.IsAny(), It.IsAny())) + .Returns(async (TGetOperationStatusReq request, CancellationToken token) => + { + started.TrySetResult(token); + try + { + await Task.Delay(Timeout.Infinite, token); + } + catch (OperationCanceledException) + { + timedOut.TrySetResult(true); + throw; + } + return new TGetOperationStatusResp(); + }); + using var poller = new DatabricksOperationStatusPoller( + _mockStatement.Object, _mockResponse.Object, requestTimeoutSeconds: 3); + + poller.Start(); + await WaitForSignal(started.Task); + CancellationToken requestToken = await started.Task; + poller.Stop(); + Assert.False(requestToken.IsCancellationRequested); + await WaitForSignal(timedOut.Task); + poller.Dispose(); + + _mockClient.Verify(c => c.GetOperationStatus(It.IsAny(), It.IsAny()), Times.Once); + } + + private static async Task WaitForSignal(Task task) + { + Assert.Same(task, await Task.WhenAny(task, Task.Delay(TimeSpan.FromSeconds(10)))); + await task; + } + + private sealed class PausingHttpTransport : THttpTransport + { + private readonly int _pauseAfterWrite; + private int _writes; + + public TaskCompletionSource Paused { get; } = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + public TaskCompletionSource Resume { get; } = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + public PausingHttpTransport(HttpClient client, int pauseAfterWrite) + : base(client, new TConfiguration()) + { + _pauseAfterWrite = pauseAfterWrite; + } + + public override async Task WriteAsync(byte[] buffer, int offset, int length, CancellationToken cancellationToken) + { + await base.WriteAsync(buffer, offset, length, cancellationToken); + if (++_writes == _pauseAfterWrite) + { + Paused.TrySetResult(true); + await Resume.Task; + } + } + } + + private sealed class ThriftHandler : HttpMessageHandler + { + private readonly TCLIService.AsyncProcessor _processor; + + public ThriftHandler(TCLIService.IAsync server) + { + _processor = new TCLIService.AsyncProcessor(server); + } + + protected override async Task SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) + { + using var input = new MemoryStream(await request.Content!.ReadAsByteArrayAsync()); + using var output = new MemoryStream(); + using var transport = new TStreamTransport(input, output, new TConfiguration()); + Assert.True(await _processor.ProcessAsync(new TBinaryProtocol(transport), new TBinaryProtocol(transport), cancellationToken)); + Assert.Equal(input.Length, input.Position); + return new HttpResponseMessage(HttpStatusCode.OK) { Content = new ByteArrayContent(output.ToArray()) }; + } + } } } From b6d58d103ef85a8c95b34bae118dfde838a4d21b Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Tue, 13 Oct 2026 04:46:44 +0000 Subject: [PATCH 2/2] fix(csharp): protect Thrift serialization from cancellation Signed-off-by: Vu Anh Phung --- csharp/src/DatabricksConnection.cs | 2 +- .../Reader/DatabricksOperationStatusPoller.cs | 7 +- csharp/src/UninterruptibleWriteProtocol.cs | 77 ++++++ .../DatabricksOperationStatusPollerTests.cs | 158 ++++++------ csharp/test/Unit/ThriftCancellationTests.cs | 226 ++++++++++++++++++ 5 files changed, 394 insertions(+), 76 deletions(-) create mode 100644 csharp/src/UninterruptibleWriteProtocol.cs create mode 100644 csharp/test/Unit/ThriftCancellationTests.cs diff --git a/csharp/src/DatabricksConnection.cs b/csharp/src/DatabricksConnection.cs index 5bc615768..ff1cee8f0 100644 --- a/csharp/src/DatabricksConnection.cs +++ b/csharp/src/DatabricksConnection.cs @@ -201,7 +201,7 @@ private void LogConnectionProperties(Activity? activity) protected override TCLIService.IAsync CreateTCLIServiceClient(TProtocol protocol) { - return new ThreadSafeClient(new TCLIService.Client(protocol)); + return new ThreadSafeClient(new TCLIService.Client(protocol, new UninterruptibleWriteProtocol(protocol))); } private void ValidateProperties() diff --git a/csharp/src/Reader/DatabricksOperationStatusPoller.cs b/csharp/src/Reader/DatabricksOperationStatusPoller.cs index a56b8ffb3..29baac323 100644 --- a/csharp/src/Reader/DatabricksOperationStatusPoller.cs +++ b/csharp/src/Reader/DatabricksOperationStatusPoller.cs @@ -135,9 +135,7 @@ async Task PollOnceAsync(CancellationToken ct, Activity? activity) TOperationHandle? operationHandle = _response.OperationHandle; if (operationHandle == null) return; - ct.ThrowIfCancellationRequested(); - // Stopping must not interrupt serialization into Thrift's shared transport buffer. - using CancellationTokenSource timeoutCts = new CancellationTokenSource(); + using CancellationTokenSource timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(ct); timeoutCts.CancelAfter(TimeSpan.FromSeconds(_requestTimeoutSeconds)); TGetOperationStatusReq request = new TGetOperationStatusReq(operationHandle); @@ -188,6 +186,9 @@ async Task PollOnceAsync(CancellationToken ct, Activity? activity) } catch (Exception ex) { + // Thrift wraps canceled HTTP requests in TTransportException. + if (ct.IsCancellationRequested) return; + consecutiveFailures++; // Log the error but continue polling. Transient errors (e.g. ObjectDisposedException diff --git a/csharp/src/UninterruptibleWriteProtocol.cs b/csharp/src/UninterruptibleWriteProtocol.cs new file mode 100644 index 000000000..a3b1f8141 --- /dev/null +++ b/csharp/src/UninterruptibleWriteProtocol.cs @@ -0,0 +1,77 @@ +/* +* Copyright (c) 2025 ADBC Drivers Contributors +* +* Licensed under the Apache License, Version 2.0 (the "License"); +* you may not use this file except in compliance with the License. +* You may obtain a copy of the License at +* +* http://www.apache.org/licenses/LICENSE-2.0 +* +* Unless required by applicable law or agreed to in writing, software +* distributed under the License is distributed on an "AS IS" BASIS, +* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +* See the License for the specific language governing permissions and +* limitations under the License. +*/ + +using System; +using System.Threading; +using System.Threading.Tasks; +using Thrift.Protocol; +using Thrift.Protocol.Entities; + +namespace AdbcDrivers.Databricks +{ + // Complete in-memory serialization so cancellable FlushAsync always resets the HTTP request buffer. + internal sealed class UninterruptibleWriteProtocol : TProtocolDecorator + { + public UninterruptibleWriteProtocol(TProtocol protocol) : base(protocol) + { + } + + public override Task WriteMessageBeginAsync(TMessage message, CancellationToken cancellationToken) + => base.WriteMessageBeginAsync(message, CancellationToken.None); + public override Task WriteMessageEndAsync(CancellationToken cancellationToken) + => base.WriteMessageEndAsync(CancellationToken.None); + public override Task WriteStructBeginAsync(TStruct @struct, CancellationToken cancellationToken) + => base.WriteStructBeginAsync(@struct, CancellationToken.None); + public override Task WriteStructEndAsync(CancellationToken cancellationToken) + => base.WriteStructEndAsync(CancellationToken.None); + public override Task WriteFieldBeginAsync(TField field, CancellationToken cancellationToken) + => base.WriteFieldBeginAsync(field, CancellationToken.None); + public override Task WriteFieldEndAsync(CancellationToken cancellationToken) + => base.WriteFieldEndAsync(CancellationToken.None); + public override Task WriteFieldStopAsync(CancellationToken cancellationToken) + => base.WriteFieldStopAsync(CancellationToken.None); + public override Task WriteMapBeginAsync(TMap map, CancellationToken cancellationToken) + => base.WriteMapBeginAsync(map, CancellationToken.None); + public override Task WriteMapEndAsync(CancellationToken cancellationToken) + => base.WriteMapEndAsync(CancellationToken.None); + public override Task WriteListBeginAsync(TList list, CancellationToken cancellationToken) + => base.WriteListBeginAsync(list, CancellationToken.None); + public override Task WriteListEndAsync(CancellationToken cancellationToken) + => base.WriteListEndAsync(CancellationToken.None); + public override Task WriteSetBeginAsync(TSet set, CancellationToken cancellationToken) + => base.WriteSetBeginAsync(set, CancellationToken.None); + public override Task WriteSetEndAsync(CancellationToken cancellationToken) + => base.WriteSetEndAsync(CancellationToken.None); + public override Task WriteBoolAsync(bool value, CancellationToken cancellationToken) + => base.WriteBoolAsync(value, CancellationToken.None); + public override Task WriteByteAsync(sbyte value, CancellationToken cancellationToken) + => base.WriteByteAsync(value, CancellationToken.None); + public override Task WriteI16Async(short value, CancellationToken cancellationToken) + => base.WriteI16Async(value, CancellationToken.None); + public override Task WriteI32Async(int value, CancellationToken cancellationToken) + => base.WriteI32Async(value, CancellationToken.None); + public override Task WriteI64Async(long value, CancellationToken cancellationToken) + => base.WriteI64Async(value, CancellationToken.None); + public override Task WriteDoubleAsync(double value, CancellationToken cancellationToken) + => base.WriteDoubleAsync(value, CancellationToken.None); + public override Task WriteStringAsync(string value, CancellationToken cancellationToken) + => base.WriteStringAsync(value, CancellationToken.None); + public override Task WriteBinaryAsync(byte[] value, CancellationToken cancellationToken) + => base.WriteBinaryAsync(value, CancellationToken.None); + public override Task WriteUuidAsync(Guid value, CancellationToken cancellationToken) + => base.WriteUuidAsync(value, CancellationToken.None); + } +} diff --git a/csharp/test/Unit/DatabricksOperationStatusPollerTests.cs b/csharp/test/Unit/DatabricksOperationStatusPollerTests.cs index 459d6a9b7..3380ea867 100644 --- a/csharp/test/Unit/DatabricksOperationStatusPollerTests.cs +++ b/csharp/test/Unit/DatabricksOperationStatusPollerTests.cs @@ -22,8 +22,7 @@ */ using System; -using System.IO; -using System.Net; +using System.Diagnostics; using System.Net.Http; using System.Threading; using System.Threading.Tasks; @@ -32,11 +31,9 @@ using AdbcDrivers.Databricks.Telemetry.Proto; using AdbcDrivers.HiveServer2.Hive2; using Apache.Hive.Service.Rpc.Thrift; +using Apache.Arrow.Adbc.Tracing; using Moq; -using Thrift; using Thrift.Protocol; -using Thrift.Transport; -using Thrift.Transport.Client; using Xunit; using Xunit.Abstractions; @@ -60,7 +57,7 @@ public DatabricksOperationStatusPollerTests(ITestOutputHelper outputHelper) _mockResponse = new Mock(); _operationHandle = new TOperationHandle { - OperationId = new THandleIdentifier { Guid = new byte[] { 1, 2, 3, 4 } }, + OperationId = new THandleIdentifier { Guid = new byte[] { 1, 2, 3, 4 }, Secret = new byte[16] }, OperationType = TOperationType.EXECUTE_STATEMENT }; @@ -322,10 +319,12 @@ public async Task UsesCustomRequestTimeout() [InlineData(4, false)] [InlineData(10, false)] [InlineData(20, false)] + [InlineData(28, false)] [InlineData(1, true)] [InlineData(4, true)] [InlineData(10, true)] [InlineData(20, true)] + [InlineData(28, true)] public async Task StoppingDuringSerializationPreservesCloseOperation(int pauseAfterWrite, bool cancelExternally) { _operationHandle.OperationId.Secret = new byte[16]; @@ -345,10 +344,11 @@ public async Task StoppingDuringSerializationPreservesCloseOperation(int pauseAf BaseAddress = new Uri("http://localhost/") }; using var transport = new PausingHttpTransport(httpClient, pauseAfterWrite); - using var client = new ThreadSafeClient(new TCLIService.Client(new TBinaryProtocol(transport))); + using var connection = new ThriftTestConnection(); + using var client = connection.CreateClient(new TBinaryProtocol(transport)); _mockStatement.Setup(s => s.Client).Returns(client); using var cancellation = new CancellationTokenSource(); - using var poller = new DatabricksOperationStatusPoller(_mockStatement.Object, _mockResponse.Object); + using var poller = new DatabricksOperationStatusPoller(_mockStatement.Object, _mockResponse.Object, _heartbeatIntervalSeconds); poller.Start(cancellation.Token); try @@ -371,93 +371,107 @@ public async Task StoppingDuringSerializationPreservesCloseOperation(int pauseAf await client.CloseOperation(new TCloseOperationReq(_operationHandle)); - server.Verify(c => c.GetOperationStatus(It.IsAny(), It.IsAny()), Times.Once); + server.Verify(c => c.GetOperationStatus(It.IsAny(), It.IsAny()), Times.Never); server.Verify(c => c.CloseOperation(It.IsAny(), It.IsAny()), Times.Once); } [Fact] - public async Task StoppedPollStillHonorsRequestTimeout() + public async Task RequestTimeoutDuringSerializationPreservesCloseOperation() { + var server = new Mock(MockBehavior.Strict); + server.Setup(c => c.CloseOperation(It.IsAny(), It.IsAny())) + .ReturnsAsync(new TCloseOperationResp { Status = new TStatus(TStatusCode.SUCCESS_STATUS) }); + using var httpClient = new HttpClient(new ThriftHandler(server.Object)) { BaseAddress = new Uri("http://localhost/") }; + using var transport = new PausingHttpTransport(httpClient, pauseAfterWrite: 20); + using var connection = new ThriftTestConnection(); + using var client = connection.CreateClient(new TBinaryProtocol(transport)); var started = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); - var timedOut = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); _mockClient.Setup(c => c.GetOperationStatus(It.IsAny(), It.IsAny())) - .Returns(async (TGetOperationStatusReq request, CancellationToken token) => + .Returns((TGetOperationStatusReq request, CancellationToken token) => { started.TrySetResult(token); - try - { - await Task.Delay(Timeout.Infinite, token); - } - catch (OperationCanceledException) - { - timedOut.TrySetResult(true); - throw; - } - return new TGetOperationStatusResp(); + return client.GetOperationStatus(request, token); }); using var poller = new DatabricksOperationStatusPoller( - _mockStatement.Object, _mockResponse.Object, requestTimeoutSeconds: 3); - + _mockStatement.Object, _mockResponse.Object, _heartbeatIntervalSeconds, requestTimeoutSeconds: 1); poller.Start(); - await WaitForSignal(started.Task); - CancellationToken requestToken = await started.Task; - poller.Stop(); - Assert.False(requestToken.IsCancellationRequested); - await WaitForSignal(timedOut.Task); - poller.Dispose(); - - _mockClient.Verify(c => c.GetOperationStatus(It.IsAny(), It.IsAny()), Times.Once); - } - - private static async Task WaitForSignal(Task task) - { - Assert.Same(task, await Task.WhenAny(task, Task.Delay(TimeSpan.FromSeconds(10)))); - await task; - } - - private sealed class PausingHttpTransport : THttpTransport - { - private readonly int _pauseAfterWrite; - private int _writes; - - public TaskCompletionSource Paused { get; } = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); - public TaskCompletionSource Resume { get; } = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); - - public PausingHttpTransport(HttpClient client, int pauseAfterWrite) - : base(client, new TConfiguration()) + try { - _pauseAfterWrite = pauseAfterWrite; + await WaitForSignal(transport.Paused.Task); + await WaitForSignal(started.Task); + var cancelled = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + using var registration = (await started.Task).Register(() => cancelled.TrySetResult(true)); + await WaitForSignal(cancelled.Task); } - - public override async Task WriteAsync(byte[] buffer, int offset, int length, CancellationToken cancellationToken) + finally { - await base.WriteAsync(buffer, offset, length, cancellationToken); - if (++_writes == _pauseAfterWrite) - { - Paused.TrySetResult(true); - await Resume.Task; - } + poller.Stop(); + transport.Resume.TrySetResult(true); + poller.Dispose(); } + await client.CloseOperation(new TCloseOperationReq(_operationHandle)); + server.Verify(c => c.CloseOperation(It.IsAny(), It.IsAny()), Times.Once); } - private sealed class ThriftHandler : HttpMessageHandler + [Theory] + [InlineData(true)] + [InlineData(false)] + public async Task HttpCancellationOnlyReportsErrorForRequestTimeout(bool stopPoller) { - private readonly TCLIService.AsyncProcessor _processor; - - public ThriftHandler(TCLIService.IAsync server) + var server = new Mock(MockBehavior.Strict); + server.Setup(c => c.GetOperationStatus(It.IsAny(), It.IsAny())) + .ReturnsAsync(new TGetOperationStatusResp { Status = new TStatus(TStatusCode.SUCCESS_STATUS) }); + server.Setup(c => c.CloseOperation(It.IsAny(), It.IsAny())) + .ReturnsAsync(new TCloseOperationResp { Status = new TStatus(TStatusCode.SUCCESS_STATUS) }); + using var handler = new ThriftHandler(server.Object, blockFirstRequest: true); + using var httpClient = new HttpClient(handler) { BaseAddress = new Uri("http://localhost/") }; + using var transport = new PausingHttpTransport(httpClient, pauseAfterWrite: int.MaxValue); + using var connection = new ThriftTestConnection(); + using var client = connection.CreateClient(new TBinaryProtocol(transport)); + _mockStatement.Setup(s => s.Client).Returns(client); + using var trace = new ActivityTrace(nameof(HttpCancellationOnlyReportsErrorForRequestTimeout)); + var tracer = new Mock(); + tracer.Setup(t => t.Trace).Returns(trace); + var completed = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + using var listener = new ActivityListener + { + ShouldListenTo = source => ReferenceEquals(source, trace.ActivitySource), + Sample = (ref ActivityCreationOptions _) => ActivitySamplingResult.AllData, + ActivityStopped = activity => completed.TrySetResult(activity) + }; + ActivitySource.AddActivityListener(listener); + using var poller = new DatabricksOperationStatusPoller( + _mockStatement.Object, _mockResponse.Object, _heartbeatIntervalSeconds, + requestTimeoutSeconds: stopPoller ? 30 : 1, activityTracer: tracer.Object); + poller.Start(); + try { - _processor = new TCLIService.AsyncProcessor(server); + await WaitForSignal(handler.Started.Task); + if (stopPoller) poller.Stop(); + await WaitForSignal(completed.Task); } - - protected override async Task SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) + finally + { + handler.Resume.TrySetResult(true); + poller.Dispose(); + } + Activity activity = await completed.Task; + if (stopPoller) + { + Assert.DoesNotContain(activity.Events, e => e.Name == "operation_status_poller.poll_error"); + } + else { - using var input = new MemoryStream(await request.Content!.ReadAsByteArrayAsync()); - using var output = new MemoryStream(); - using var transport = new TStreamTransport(input, output, new TConfiguration()); - Assert.True(await _processor.ProcessAsync(new TBinaryProtocol(transport), new TBinaryProtocol(transport), cancellationToken)); - Assert.Equal(input.Length, input.Position); - return new HttpResponseMessage(HttpStatusCode.OK) { Content = new ByteArrayContent(output.ToArray()) }; + Assert.Contains(activity.Events, e => e.Name == "operation_status_poller.poll_error"); } + Assert.NotEqual(ActivityStatusCode.Error, activity.Status); + await client.CloseOperation(new TCloseOperationReq(_operationHandle)); + server.Verify(c => c.CloseOperation(It.IsAny(), It.IsAny()), Times.Once); + } + + private static Task WaitForSignal(Task task) + { + return ThriftCancellationTests.WaitForSignal(task); } } } diff --git a/csharp/test/Unit/ThriftCancellationTests.cs b/csharp/test/Unit/ThriftCancellationTests.cs new file mode 100644 index 000000000..55e0a9f2c --- /dev/null +++ b/csharp/test/Unit/ThriftCancellationTests.cs @@ -0,0 +1,226 @@ +/* +* Copyright (c) 2025 ADBC Drivers Contributors +* +* Licensed under the Apache License, Version 2.0 (the "License"); +* you may not use this file except in compliance with the License. +* You may obtain a copy of the License at +* +* http://www.apache.org/licenses/LICENSE-2.0 +* +* Unless required by applicable law or agreed to in writing, software +* distributed under the License is distributed on an "AS IS" BASIS, +* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +* See the License for the specific language governing permissions and +* limitations under the License. +*/ + +using System; +using System.Collections.Generic; +using System.IO; +using System.Net; +using System.Net.Http; +using System.Threading; +using System.Threading.Tasks; +using AdbcDrivers.HiveServer2.Hive2; +using AdbcDrivers.HiveServer2.Spark; +using Apache.Hive.Service.Rpc.Thrift; +using Moq; +using Thrift; +using Thrift.Protocol; +using Thrift.Transport; +using Thrift.Transport.Client; +using Xunit; + +namespace AdbcDrivers.Databricks.Tests.Unit +{ + public class ThriftCancellationTests + { + [Theory] + [InlineData("GetOperationStatus", 1)] + [InlineData("GetOperationStatus", 4)] + [InlineData("GetOperationStatus", 10)] + [InlineData("GetOperationStatus", 20)] + [InlineData("FetchResults", 1)] + [InlineData("FetchResults", 4)] + [InlineData("FetchResults", 10)] + [InlineData("FetchResults", 20)] + [InlineData("GetResultSetMetadata", 1)] + [InlineData("GetResultSetMetadata", 4)] + [InlineData("GetResultSetMetadata", 10)] + [InlineData("GetResultSetMetadata", 20)] + [InlineData("ExecuteStatement", 1)] + [InlineData("ExecuteStatement", 4)] + [InlineData("ExecuteStatement", 10)] + [InlineData("ExecuteStatement", 20)] + public async Task CancellationDuringSerializationPreservesConnection(string operation, int pauseAfterWrite) + { + var handle = new TOperationHandle + { + OperationId = new THandleIdentifier { Guid = new byte[16], Secret = new byte[16] }, + OperationType = TOperationType.EXECUTE_STATEMENT, + HasResultSet = true + }; + var server = new Mock(MockBehavior.Strict); + server.Setup(c => c.CancelOperation(It.IsAny(), It.IsAny())) + .ReturnsAsync(new TCancelOperationResp { Status = new TStatus(TStatusCode.SUCCESS_STATUS) }); + server.Setup(c => c.CloseOperation(It.IsAny(), It.IsAny())) + .ReturnsAsync(new TCloseOperationResp { Status = new TStatus(TStatusCode.SUCCESS_STATUS) }); + using var handler = new ThriftHandler(server.Object); + using var httpClient = new HttpClient(handler) { BaseAddress = new Uri("http://localhost/") }; + using var transport = new PausingHttpTransport(httpClient, pauseAfterWrite); + using var connection = new ThriftTestConnection(); + using var client = connection.CreateClient(new TBinaryProtocol(transport)); + using var cancellation = new CancellationTokenSource(); + + Task request = operation switch + { + "GetOperationStatus" => client.GetOperationStatus(new TGetOperationStatusReq(handle), cancellation.Token), + "FetchResults" => client.FetchResults(new TFetchResultsReq(handle, TFetchOrientation.FETCH_NEXT, 1000), cancellation.Token), + "GetResultSetMetadata" => client.GetResultSetMetadata(new TGetResultSetMetadataReq(handle), cancellation.Token), + "ExecuteStatement" => client.ExecuteStatement(new TExecuteStatementReq(new TSessionHandle(handle.OperationId), "SELECT 1") + { + Parameters = new List(), + RunAsync = true + }, cancellation.Token), + _ => throw new ArgumentException(operation) + }; + + try + { + await WaitForSignal(transport.Paused.Task); + cancellation.Cancel(); + } + finally + { + transport.Resume.TrySetResult(true); + } + var exception = await Record.ExceptionAsync(() => request); + if (exception is TTransportException transportException) + { + Assert.Equal(TTransportException.ExceptionType.Interrupted, transportException.Type); + exception = transportException.InnerException; + } + Assert.IsAssignableFrom(exception); + Assert.Equal(0, handler.RequestCount); + + await client.CancelOperation(new TCancelOperationReq(handle)); + await client.CloseOperation(new TCloseOperationReq(handle)); + + server.Verify(c => c.CancelOperation(It.IsAny(), It.IsAny()), Times.Once); + server.Verify(c => c.CloseOperation(It.IsAny(), It.IsAny()), Times.Once); + } + + [Fact] + public async Task CancellationWhileWaitingForClientDoesNotWriteRequest() + { + var handle = new TOperationHandle { OperationId = new THandleIdentifier { Guid = new byte[16], Secret = new byte[16] } }; + var server = new Mock(MockBehavior.Strict); + server.Setup(c => c.GetOperationStatus(It.IsAny(), It.IsAny())) + .ReturnsAsync(new TGetOperationStatusResp { Status = new TStatus(TStatusCode.SUCCESS_STATUS) }); + server.Setup(c => c.CloseOperation(It.IsAny(), It.IsAny())) + .ReturnsAsync(new TCloseOperationResp { Status = new TStatus(TStatusCode.SUCCESS_STATUS) }); + using var handler = new ThriftHandler(server.Object, blockFirstRequest: true); + using var httpClient = new HttpClient(handler) { BaseAddress = new Uri("http://localhost/") }; + using var transport = new PausingHttpTransport(httpClient, pauseAfterWrite: int.MaxValue); + using var connection = new ThriftTestConnection(); + using var client = connection.CreateClient(new TBinaryProtocol(transport)); + using var cancellation = new CancellationTokenSource(); + + Task activeRequest = client.GetOperationStatus(new TGetOperationStatusReq(handle)); + try + { + await WaitForSignal(handler.Started.Task); + int writes = transport.WriteCount; + Task queuedRequest = client.GetResultSetMetadata(new TGetResultSetMetadataReq(handle), cancellation.Token); + cancellation.Cancel(); + await Assert.ThrowsAnyAsync(() => queuedRequest); + Assert.Equal(writes, transport.WriteCount); + Assert.Equal(1, handler.RequestCount); + } + finally + { + handler.Resume.TrySetResult(true); + await activeRequest; + } + await client.CloseOperation(new TCloseOperationReq(handle)); + server.Verify(c => c.CloseOperation(It.IsAny(), It.IsAny()), Times.Once); + } + + internal static async Task WaitForSignal(Task task) + { + Assert.Same(task, await Task.WhenAny(task, Task.Delay(TimeSpan.FromSeconds(10)))); + await task; + } + } + + internal sealed class ThriftTestConnection : DatabricksConnection + { + public ThriftTestConnection() : base(new Dictionary + { + [SparkParameters.HostName] = "localhost", + [SparkParameters.Token] = "test-token" + }) + { + } + + public ThreadSafeClient CreateClient(TProtocol protocol) => (ThreadSafeClient)CreateTCLIServiceClient(protocol); + } + + internal sealed class PausingHttpTransport : THttpTransport + { + private readonly int _pauseAfterWrite; + public int WriteCount { get; private set; } + public TaskCompletionSource Paused { get; } = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + public TaskCompletionSource Resume { get; } = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + public PausingHttpTransport(HttpClient client, int pauseAfterWrite) : base(client, new TConfiguration()) + { + _pauseAfterWrite = pauseAfterWrite; + } + + public override async Task WriteAsync(byte[] buffer, int offset, int length, CancellationToken cancellationToken) + { + await base.WriteAsync(buffer, offset, length, cancellationToken); + if (++WriteCount == _pauseAfterWrite) + { + Paused.TrySetResult(true); + await Resume.Task; + } + } + } + + internal sealed class ThriftHandler : HttpMessageHandler + { + private readonly TCLIService.AsyncProcessor _processor; + private readonly bool _blockFirstRequest; + private int _requestCount; + public int RequestCount => Volatile.Read(ref _requestCount); + public TaskCompletionSource Started { get; } = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + public TaskCompletionSource Resume { get; } = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + public ThriftHandler(TCLIService.IAsync server, bool blockFirstRequest = false) + { + _processor = new TCLIService.AsyncProcessor(server); + _blockFirstRequest = blockFirstRequest; + } + + protected override async Task SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + if (Interlocked.Increment(ref _requestCount) == 1 && _blockFirstRequest) + { + Started.TrySetResult(true); + var cancelled = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + using var registration = cancellationToken.Register(() => cancelled.TrySetResult(true)); + await Task.WhenAny(Resume.Task, cancelled.Task); + cancellationToken.ThrowIfCancellationRequested(); + } + using var input = new MemoryStream(await request.Content!.ReadAsByteArrayAsync()); + using var output = new MemoryStream(); + using var transport = new TStreamTransport(input, output, new TConfiguration()); + Assert.True(await _processor.ProcessAsync(new TBinaryProtocol(transport), new TBinaryProtocol(transport), cancellationToken)); + Assert.Equal(input.Length, input.Position); + return new HttpResponseMessage(HttpStatusCode.OK) { Content = new ByteArrayContent(output.ToArray()) }; + } + } +}