diff --git a/csharp/src/DatabricksConnection.cs b/csharp/src/DatabricksConnection.cs index 5bc61576..ff1cee8f 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 2b6170ce..29baac32 100644 --- a/csharp/src/Reader/DatabricksOperationStatusPoller.cs +++ b/csharp/src/Reader/DatabricksOperationStatusPoller.cs @@ -186,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 00000000..a3b1f814 --- /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 67551105..3380ea86 100644 --- a/csharp/test/Unit/DatabricksOperationStatusPollerTests.cs +++ b/csharp/test/Unit/DatabricksOperationStatusPollerTests.cs @@ -22,6 +22,8 @@ */ using System; +using System.Diagnostics; +using System.Net.Http; using System.Threading; using System.Threading.Tasks; using AdbcDrivers.Databricks.Reader; @@ -29,7 +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.Protocol; using Xunit; using Xunit.Abstractions; @@ -53,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 }; @@ -309,5 +313,165 @@ 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(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]; + _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 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, _heartbeatIntervalSeconds); + + 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.Never); + server.Verify(c => c.CloseOperation(It.IsAny(), It.IsAny()), Times.Once); + } + + [Fact] + 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); + _mockClient.Setup(c => c.GetOperationStatus(It.IsAny(), It.IsAny())) + .Returns((TGetOperationStatusReq request, CancellationToken token) => + { + started.TrySetResult(token); + return client.GetOperationStatus(request, token); + }); + using var poller = new DatabricksOperationStatusPoller( + _mockStatement.Object, _mockResponse.Object, _heartbeatIntervalSeconds, requestTimeoutSeconds: 1); + poller.Start(); + try + { + 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); + } + finally + { + 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); + } + + [Theory] + [InlineData(true)] + [InlineData(false)] + public async Task HttpCancellationOnlyReportsErrorForRequestTimeout(bool stopPoller) + { + 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 + { + await WaitForSignal(handler.Started.Task); + if (stopPoller) poller.Stop(); + await WaitForSignal(completed.Task); + } + 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 + { + 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 00000000..55e0a9f2 --- /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()) }; + } + } +}