fix(csharp): prevent Thrift buffer corruption on cancellation

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()

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

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);
+    }
+}

csharp/test/Unit/DatabricksOperationStatusPollerTests.cs

@@ -22,14 +22,18 @@
 */
 
 using System;
+using System.Diagnostics;
+using System.Net.Http;
 using System.Threading;
 using System.Threading.Tasks;
 using AdbcDrivers.Databricks.Reader;
 using AdbcDrivers.Databricks.Telemetry;
 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<IResponse>();
             _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<TGetOperationStatusReq>(), It.IsAny<CancellationToken>()), 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<TCLIService.IAsync>(MockBehavior.Strict);
+            server.Setup(c => c.GetOperationStatus(It.IsAny<TGetOperationStatusReq>(), It.IsAny<CancellationToken>()))
+                .ReturnsAsync(new TGetOperationStatusResp
+                {
+                    Status = new TStatus(TStatusCode.SUCCESS_STATUS),
+                    OperationState = TOperationState.FINISHED_STATE
+                });
+            server.Setup(c => c.CloseOperation(It.IsAny<TCloseOperationReq>(), It.IsAny<CancellationToken>()))
+                .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<TGetOperationStatusReq>(), It.IsAny<CancellationToken>()), Times.Never);
+            server.Verify(c => c.CloseOperation(It.IsAny<TCloseOperationReq>(), It.IsAny<CancellationToken>()), Times.Once);
+        }
+
+        [Fact]
+        public async Task RequestTimeoutDuringSerializationPreservesCloseOperation()
+        {
+            var server = new Mock<TCLIService.IAsync>(MockBehavior.Strict);
+            server.Setup(c => c.CloseOperation(It.IsAny<TCloseOperationReq>(), It.IsAny<CancellationToken>()))
+                .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<CancellationToken>(TaskCreationOptions.RunContinuationsAsynchronously);
+            _mockClient.Setup(c => c.GetOperationStatus(It.IsAny<TGetOperationStatusReq>(), It.IsAny<CancellationToken>()))
+                .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<bool>(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<TCloseOperationReq>(), It.IsAny<CancellationToken>()), Times.Once);
+        }
+
+        [Theory]
+        [InlineData(true)]
+        [InlineData(false)]
+        public async Task HttpCancellationOnlyReportsErrorForRequestTimeout(bool stopPoller)
+        {
+            var server = new Mock<TCLIService.IAsync>(MockBehavior.Strict);
+            server.Setup(c => c.GetOperationStatus(It.IsAny<TGetOperationStatusReq>(), It.IsAny<CancellationToken>()))
+                .ReturnsAsync(new TGetOperationStatusResp { Status = new TStatus(TStatusCode.SUCCESS_STATUS) });
+            server.Setup(c => c.CloseOperation(It.IsAny<TCloseOperationReq>(), It.IsAny<CancellationToken>()))
+                .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<IActivityTracer>();
+            tracer.Setup(t => t.Trace).Returns(trace);
+            var completed = new TaskCompletionSource<Activity>(TaskCreationOptions.RunContinuationsAsynchronously);
+            using var listener = new ActivityListener
+            {
+                ShouldListenTo = source => ReferenceEquals(source, trace.ActivitySource),
+                Sample = (ref ActivityCreationOptions<ActivityContext> _) => 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<TCloseOperationReq>(), It.IsAny<CancellationToken>()), Times.Once);
+        }
+
+        private static Task WaitForSignal(Task task)
+        {
+            return ThriftCancellationTests.WaitForSignal(task);
+        }
     }
 }