fix(csharp): tear down in-flight CloudFetch on connection dispose and statement cancel/dispose

csharp/src/DatabricksConnection.cs

@@ -122,6 +122,20 @@ internal class DatabricksConnection : SparkHttpConnection
         // Shared HttpClient for CloudFetch downloads (created once, reused across queries)
         private HttpClient? _cloudFetchHttpClient;
 
+        // Connection-scoped shutdown signal for in-flight CloudFetch pipelines. Cancelled at the
+        // start of Dispose so any active download manager (which links this token into its own
+        // cancellation source) tears down promptly, unblocking a reader parked on a download.
+        // Without this, closing a connection mid-CloudFetch cannot cancel the pipeline (the reader
+        // that owns it is blocked), so the reader spins on download retries until the retry budget
+        // expires (minutes) — see CloseConnection_DuringCloudFetch_ShouldNotHang.
+        private readonly CancellationTokenSource _cloudFetchShutdownCts = new CancellationTokenSource();
+
+        /// <summary>
+        /// Token cancelled when this connection is disposed. CloudFetch download managers link this
+        /// into their pipeline cancellation source so connection close tears down in-flight downloads.
+        /// </summary>
+        internal CancellationToken CloudFetchShutdownToken => _cloudFetchShutdownCts.Token;
+
         // Telemetry
         private IConnectionTelemetry _telemetry = NoOpConnectionTelemetry.Instance;
         // Stopwatch covering the connection lifetime; started at construction and used to
@@ -1183,6 +1197,13 @@ protected override void Dispose(bool disposing)
         {
             if (disposing)
             {
+                // Signal in-flight CloudFetch pipelines to stop before tearing down the transport.
+                // The download manager links this token into its own cancellation source, so this
+                // cancels the download loop and faults any in-flight download — unblocking a reader
+                // parked on GetNextDownloadedFileAsync / DownloadCompletedTask instead of leaving it
+                // to spin on retries against the about-to-be-disposed HttpClient.
+                try { _cloudFetchShutdownCts.Cancel(); } catch (ObjectDisposedException) { }
+
                 // Dispose the shared CloudFetch HttpClient before closing the session so any
                 // in-flight CloudFetch HTTP work is torn down first (PR #385: concurrent Dispose deadlock fix).
                 _cloudFetchHttpClient?.Dispose();
@@ -1210,6 +1231,13 @@ protected override void Dispose(bool disposing)
                 {
                     closeStopwatch.Stop();
                     closeSessionElapsedMs = closeStopwatch.ElapsedMilliseconds;
+
+                    // Dispose the CloudFetch shutdown CTS first so cleanup is unconditional:
+                    // it was already Cancel()ed at the top of Dispose and nothing below depends
+                    // on it, so releasing it here guarantees it happens even if the telemetry
+                    // calls below throw.
+                    _cloudFetchShutdownCts.Dispose();
+
                     EmitDeleteSessionTelemetry(closeSessionElapsedMs, closeSessionError);
 
                     // Clean up telemetry client.

csharp/src/Reader/CloudFetch/CloudFetchDownloadManager.cs

@@ -120,7 +120,7 @@ public CloudFetchDownloadManager(
         }
 
         /// <inheritdoc />
-        public async Task StartAsync()
+        public async Task StartAsync(CancellationToken cancellationToken = default)
         {
             ThrowIfDisposed();
 
@@ -129,8 +129,9 @@ public async Task StartAsync()
                 throw new InvalidOperationException("Download manager is already started.");
             }
 
-            // Create a new cancellation token source
-            _cancellationTokenSource = new CancellationTokenSource();
+            // Link the caller's token (e.g. the connection's shutdown token) so cancelling it tears
+            // down the fetcher + downloader and unblocks any reader waiting on the result queue.
+            _cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
 
             // Start the result fetcher
             await _resultFetcher.StartAsync(_cancellationTokenSource.Token).ConfigureAwait(false);

csharp/src/Reader/CloudFetch/CloudFetchReaderFactory.cs

@@ -108,8 +108,9 @@ public static CloudFetchReader CreateThriftReader(
                 resultQueue,
                 config);
 
-            // Start the download manager
-            downloadManager.StartAsync().Wait();
+            // Start the download manager, linked to the connection's shutdown token so closing the
+            // connection mid-stream tears down the pipeline and unblocks the reader.
+            downloadManager.StartAsync(connection.CloudFetchShutdownToken).Wait();
 
             // Add telemetry tag for compression
             Activity.Current?.SetTag(StatementExecutionEvent.ResultCompressionEnabled, isLz4Compressed);

csharp/src/Reader/CloudFetch/ICloudFetchInterfaces.cs

@@ -274,8 +274,12 @@ internal interface ICloudFetchDownloadManager : IDisposable
         /// <summary>
         /// Starts the download manager.
         /// </summary>
+        /// <param name="cancellationToken">
+        /// Token linked into the pipeline's cancellation source; cancelling it (e.g. on connection
+        /// dispose) tears down the fetcher and downloader and unblocks any waiting reader.
+        /// </param>
         /// <returns>A task representing the asynchronous operation.</returns>
-        Task StartAsync();
+        Task StartAsync(CancellationToken cancellationToken = default);
 
         /// <summary>
         /// Stops the download manager.

csharp/test/Unit/Reader/CloudFetch/CloudFetchDownloadManagerTests.cs

@@ -16,11 +16,14 @@
 
 using System;
 using System.Collections.Concurrent;
+using System.Net.Http;
 using System.Threading;
 using System.Threading.Tasks;
 using AdbcDrivers.Databricks;
 using AdbcDrivers.Databricks.Reader.CloudFetch;
+using AdbcDrivers.HiveServer2.Hive2;
 using Apache.Arrow.Adbc;
+using Apache.Arrow.Adbc.Tracing;
 using Moq;
 using Xunit;
 
@@ -184,5 +187,80 @@ public async Task GetNextDownloadedFileAsync_DownloaderReturnsResult_FetcherHasE
             downloadQueue.Dispose();
             resultQueue.Dispose();
         }
+
+        /// <summary>
+        /// Regression for the CloudFetch-dispose hang (CloseConnection_DuringCloudFetch_ShouldNotHang):
+        /// cancelling the token passed to StartAsync (the connection's shutdown token) must tear down
+        /// the pipeline and unblock a reader parked in GetNextDownloadedFileAsync.
+        ///
+        /// Uses a REAL CloudFetchDownloader whose fetcher never enqueues anything, so a call to
+        /// GetNextDownloadedFileAsync blocks on the (empty, not-yet-completed) result queue — exactly
+        /// the state the reader is in when a connection is disposed mid-stream. Before the fix, the
+        /// download manager's internal CTS was linked to nothing, so cancelling an external token had
+        /// no effect and the read blocked indefinitely; with the fix the linked token cancels the
+        /// download loop, which completes the result queue and unblocks the read.
+        /// </summary>
+        [Fact]
+        public async Task StartAsync_TokenCancelled_UnblocksReaderWaitingForNextFile()
+        {
+            // Arrange — a real downloader backed by a tracer-providing statement mock.
+            var mockStatement = new Mock<IHiveServer2Statement>();
+            mockStatement.Setup(s => s.Trace).Returns(new ActivityTrace("TestActivitySource"));
+            mockStatement.Setup(s => s.TraceParent).Returns((string?)null);
+            mockStatement.Setup(s => s.AssemblyVersion).Returns("1.0.0");
+            mockStatement.Setup(s => s.AssemblyName).Returns("TestAssembly");
+
+            // Fetcher never enqueues any download, so the downloader's loop parks on the empty
+            // download queue and the reader parks on the empty result queue.
+            var mockFetcher = new Mock<ICloudFetchResultFetcher>();
+            mockFetcher.Setup(f => f.HasError).Returns(false);
+            mockFetcher.Setup(f => f.Error).Returns((Exception?)null);
+            mockFetcher.Setup(f => f.StartAsync(It.IsAny<CancellationToken>())).Returns(Task.CompletedTask);
+            mockFetcher.Setup(f => f.StopAsync()).Returns(Task.CompletedTask);
+
+            var mockMemoryManager = new Mock<ICloudFetchMemoryBufferManager>();
+            var downloadQueue = new BlockingCollection<IDownloadResult>(new ConcurrentQueue<IDownloadResult>(), 10);
+            var resultQueue = new BlockingCollection<IDownloadResult>(new ConcurrentQueue<IDownloadResult>(), 10);
+            using var httpClient = new HttpClient();
+
+            var downloader = new CloudFetchDownloader(
+                mockStatement.Object,
+                downloadQueue,
+                resultQueue,
+                mockMemoryManager.Object,
+                httpClient,
+                mockFetcher.Object,
+                maxParallelDownloads: 3,
+                isLz4Compressed: false);
+
+            var config = new CloudFetchConfiguration();
+            var manager = new CloudFetchDownloadManager(
+                mockFetcher.Object,
+                downloader,
+                mockMemoryManager.Object,
+                downloadQueue,
+                resultQueue,
+                config);
+
+            using var shutdownCts = new CancellationTokenSource();
+            await manager.StartAsync(shutdownCts.Token);
+
+            // Reader parks waiting for the next downloaded file (token None, like the real reader).
+            var readTask = Task.Run(() => manager.GetNextDownloadedFileAsync(CancellationToken.None));
+
+            // Give the read a moment to reach the blocking Take, then simulate connection dispose.
+            await Task.Delay(200);
+            Assert.False(readTask.IsCompleted, "read should still be blocked before the token is cancelled");
+
+            shutdownCts.Cancel();
+
+            // Act & Assert — the read must unblock promptly (returns null: clean end of results).
+            var completed = await Task.WhenAny(readTask, Task.Delay(5000));
+            Assert.Same(readTask, completed);
+            Assert.Null(await readTask);
+
+            // Cleanup
+            manager.Dispose();
+        }
     }
 }