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/DatabricksStatement.cs

@@ -105,9 +105,40 @@ internal class DatabricksStatement : SparkStatement, IHiveServer2Statement
 
         public override long BatchSize { get; protected set; } = DatabricksBatchSizeDefault;
 
+        // Statement-lifetime cancellation for the CloudFetch pipeline, linked to the connection's
+        // shutdown token. This gives the full connection ⊃ statement ⊃ cloudfetch cancel cascade:
+        // connection dispose cancels every statement's downloads (via the link), and Cancel()/Dispose()
+        // on a single statement stops just its downloads. It is distinct from the base
+        // HiveServer2Statement._executeTokenSource, which is disposed when ExecuteQuery() returns and so
+        // cannot cover the later CloudFetch result-fetch phase.
+        // Not readonly: refreshed at the start of each execution by RefreshCloudFetchStatementCts()
+        // so a Cancel() on a prior execution doesn't poison the next CloudFetch read (see that method).
+        // Guarded by _cloudFetchStatementCtsLock: Cancel() is explicitly supported from another thread
+        // and can race the field swap performed by a concurrent re-execute (RefreshCloudFetchStatementCts).
+        private CancellationTokenSource _cloudFetchStatementCts;
+
+        // Serializes the field swap in RefreshCloudFetchStatementCts() against the reads +
+        // Cancel()/Dispose() of _cloudFetchStatementCts. Without it, a cross-thread Cancel() issued
+        // around a re-execute boundary can act on the source the swap is replacing (cancelling the
+        // wrong pipeline) or on the source the swap just disposed (silently swallowed) — either of
+        // which drops the cancel this PR exists to deliver.
+        private readonly object _cloudFetchStatementCtsLock = new object();
+
+        /// <summary>
+        /// Token cancelled when this statement is cancelled or disposed — and, via linkage to the
+        /// connection's shutdown token, when the connection is disposed. The CloudFetch download
+        /// manager links this into its pipeline source so any of those tears down in-flight downloads.
+        /// </summary>
+        internal CancellationToken CloudFetchStatementToken
+        {
+            get { lock (_cloudFetchStatementCtsLock) { return _cloudFetchStatementCts.Token; } }
+        }
+
         public DatabricksStatement(DatabricksConnection connection)
             : base(connection)
         {
+            _cloudFetchStatementCts = CancellationTokenSource.CreateLinkedTokenSource(connection.CloudFetchShutdownToken);
+
             // set the catalog name for legacy compatibility
             // TODO: use catalog and schema fields in hiveserver2 connection instead of DefaultNamespace so we don't need to cast
             var defaultNamespace = ((DatabricksConnection)Connection).DefaultNamespace;
@@ -137,6 +168,44 @@ public DatabricksStatement(DatabricksConnection connection)
             }
         }
 
+        /// <summary>
+        /// Recreates the statement-lifetime CloudFetch cancellation source (re-linked to the
+        /// connection's shutdown token) at the start of each execution. <see cref="AdbcStatement"/>
+        /// is reusable (settable <see cref="SqlQuery"/> + repeated Execute), and <see cref="Cancel"/>/
+        /// <see cref="Dispose(bool)"/> cancel this source permanently; without a refresh a
+        /// cancel-then-reexecute would start the next CloudFetch read with an already-cancelled
+        /// token. This mirrors how the base <c>HiveServer2Statement._executeTokenSource</c> is
+        /// refreshed per-execute so a statement stays reusable after cancel.
+        /// </summary>
+        internal void RefreshCloudFetchStatementCts()
+        {
+            // Swap the field under the lock so a concurrent cross-thread Cancel()/Dispose() either
+            // acts on the old source before the swap or on the new one after it — never on a torn read.
+            CancellationTokenSource previous;
+            lock (_cloudFetchStatementCtsLock)
+            {
+                previous = _cloudFetchStatementCts;
+                _cloudFetchStatementCts = CancellationTokenSource.CreateLinkedTokenSource(
+                    ((DatabricksConnection)Connection).CloudFetchShutdownToken);
+            }
+            // Release the prior source's registration on the connection shutdown token. Under normal
+            // AdbcStatement usage the previous result set is fully consumed/disposed before the next
+            // Execute, so no pipeline is still linked to the old token here. We dispose (rather than
+            // leak) it so a reused statement (repeated Execute) doesn't accumulate one linked-CTS
+            // registration on the connection shutdown token per execution for the connection's lifetime.
+            //
+            // Constraint: the CloudFetch pipeline's cancellation source is CreateLinkedTokenSource of
+            // this statement token (CloudFetchDownloadManager), whose only link to the connection
+            // shutdown token runs THROUGH this source. If a caller re-executes while still holding an
+            // OPEN (undisposed) reader from a prior execution, disposing the old source detaches that
+            // reader's still-running pipeline from the connection-dispose cascade. That reader's own
+            // disposal (StopAsync/Dispose) still tears its pipeline down; the connection-shutdown safety
+            // net only lapses for the narrow case of a reader that is never disposed AND whose
+            // connection is then disposed. Disposing a source does not cancel its already-created
+            // linked children.
+            previous?.Dispose();
+        }
+
         private StatementTelemetryContext? CreateTelemetryContext(Telemetry.Proto.Statement.Types.Type statementType)
         {
             var session = ((DatabricksConnection)Connection).TelemetrySession;
@@ -270,6 +339,9 @@ private async Task EnsureCatalogScopedAsync()
 
         public override QueryResult ExecuteQuery()
         {
+            // Refresh the CloudFetch cancellation source before base.ExecuteQuery() captures it into
+            // the reader, so a prior Cancel() doesn't leave the next read starting cancelled.
+            RefreshCloudFetchStatementCts();
             EnsureCatalogScopedAsync().GetAwaiter().GetResult();
             var ctx = IsMetadataCommand
                 ? CreateMetadataTelemetryContext()
@@ -301,6 +373,9 @@ public override QueryResult ExecuteQuery()
 
         public override async ValueTask<QueryResult> ExecuteQueryAsync()
         {
+            // Refresh the CloudFetch cancellation source before base.ExecuteQueryAsync() captures it
+            // into the reader, so a prior Cancel() doesn't leave the next read starting cancelled.
+            RefreshCloudFetchStatementCts();
             await EnsureCatalogScopedAsync().ConfigureAwait(false);
             var ctx = IsMetadataCommand
                 ? CreateMetadataTelemetryContext()
@@ -1428,6 +1503,23 @@ protected override void Dispose(bool disposing)
         {
             if (disposing)
             {
+                // Cancel this statement's CloudFetch pipeline before anything else so in-flight
+                // downloads stop promptly if the caller disposed the statement mid-stream.
+                //
+                // Cancel + dispose under the lock so this can't interleave with a concurrent
+                // re-execute's field swap (RefreshCloudFetchStatementCts) and act on a stale source.
+                lock (_cloudFetchStatementCtsLock)
+                {
+                    try { _cloudFetchStatementCts.Cancel(); } catch (ObjectDisposedException) { }
+
+                    // Dispose the CloudFetch statement CTS before the telemetry emission below:
+                    // it was already Cancel()ed just above and nothing in the telemetry blocks
+                    // depends on it, so releasing it here guarantees the linked-token registration
+                    // it holds on the connection's _cloudFetchShutdownCts is freed even if the
+                    // telemetry calls below throw. Mirrors the ordering in DatabricksConnection.Dispose.
+                    _cloudFetchStatementCts.Dispose();
+                }
+
                 if (PendingTelemetryContext != null)
                 {
                     // Emit telemetry now that results have been consumed
@@ -1489,6 +1581,20 @@ public override void Cancel()
             long startMs = _statementLifetimeStopwatch.ElapsedMilliseconds;
             try
             {
+                // Cancel the CloudFetch pipeline for this statement first. base.Cancel() only signals
+                // the per-execute token, which is already disposed once results are streaming, so
+                // without this a Cancel() during CloudFetch would leave the downloads running. Do it
+                // before base.Cancel() because base.Cancel() issues the remote CancelOperation RPC,
+                // which can throw on a network/transport failure and would otherwise skip this cleanup.
+                //
+                // Under the lock so a re-execute's field swap (RefreshCloudFetchStatementCts) can't
+                // race this read: we cancel whichever source is current, never a torn/half-swapped one.
+                // The lock scopes only the field access — base.Cancel()'s remote RPC runs outside it.
+                lock (_cloudFetchStatementCtsLock)
+                {
+                    try { _cloudFetchStatementCts.Cancel(); } catch (ObjectDisposedException) { }
+                }
+
                 base.Cancel();
             }
             catch (Exception ex)

csharp/src/Reader/CloudFetch/CloudFetchDownloadManager.cs

@@ -77,6 +77,27 @@ public CloudFetchDownloadManager(
         /// <inheritdoc />
         public bool HasMoreResults => !_downloader.IsCompleted || !_resultQueue.IsCompleted;
 
+        /// <inheritdoc />
+        public CancellationToken PipelineToken
+        {
+            get
+            {
+                var cts = _cancellationTokenSource;
+                if (cts == null)
+                {
+                    return CancellationToken.None;
+                }
+                try
+                {
+                    return cts.Token;
+                }
+                catch (ObjectDisposedException)
+                {
+                    return CancellationToken.None;
+                }
+            }
+        }
+
         /// <inheritdoc />
         public async Task<IDownloadResult?> GetNextDownloadedFileAsync(CancellationToken cancellationToken)
         {
@@ -120,7 +141,7 @@ public CloudFetchDownloadManager(
         }
 
         /// <inheritdoc />
-        public async Task StartAsync()
+        public async Task StartAsync(CancellationToken cancellationToken = default)
         {
             ThrowIfDisposed();
 
@@ -129,8 +150,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/CloudFetchReader.cs

@@ -99,8 +99,17 @@ public CloudFetchReader(
             {
                 ThrowIfDisposed();
 
+                // Observe the pipeline's cancellation (statement cancel / connection dispose) in
+                // addition to the caller's token, so the read stops promptly even while draining
+                // already-buffered chunks — not only when it next blocks for a download.
+                CancellationToken pipelineToken = this.downloadManager?.PipelineToken ?? CancellationToken.None;
+                using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, pipelineToken);
+                CancellationToken token = linkedCts.Token;
+
                 while (true)
                 {
+                    token.ThrowIfCancellationRequested();
+
                     // Check global row limit first (used by SEA with manifest.TotalRowCount)
                     if (_totalExpectedRows > 0 && _rowsRead >= _totalExpectedRows)
                     {
@@ -126,7 +135,7 @@ public CloudFetchReader(
                     // If we have a current reader, try to read the next batch
                     if (this.currentReader != null)
                     {
-                        RecordBatch? next = await this.currentReader.ReadNextRecordBatchAsync(cancellationToken);
+                        RecordBatch? next = await this.currentReader.ReadNextRecordBatchAsync(token);
                         if (next != null)
                         {
                             // Apply row count limiting: trim the batch if it would exceed expected rows
@@ -151,7 +160,7 @@ public CloudFetchReader(
                         try
                         {
                             // Get the next downloaded file
-                            this.currentDownloadResult = await this.downloadManager.GetNextDownloadedFileAsync(cancellationToken);
+                            this.currentDownloadResult = await this.downloadManager.GetNextDownloadedFileAsync(token);
                             if (this.currentDownloadResult == null)
                             {
                                 Activity.Current?.AddEvent("cloudfetch.reader_no_more_files", [

csharp/src/Reader/CloudFetch/CloudFetchReaderFactory.cs

@@ -108,8 +108,12 @@ public static CloudFetchReader CreateThriftReader(
                 resultQueue,
                 config);
 
-            // Start the download manager
-            downloadManager.StartAsync().Wait();
+            // Start the download manager, linked to the statement's CloudFetch token (which is itself
+            // linked to the connection's shutdown token). This gives the full cancel cascade: closing
+            // the connection, or cancelling/disposing this statement, tears down the pipeline and
+            // unblocks the reader. This is the Thrift path, so the statement is always a
+            // DatabricksStatement — same assumption as the DatabricksConnection cast above.
+            downloadManager.StartAsync(((DatabricksStatement)statement).CloudFetchStatementToken).Wait();
 
             // Add telemetry tag for compression
             Activity.Current?.SetTag(StatementExecutionEvent.ResultCompressionEnabled, isLz4Compressed);
@@ -207,7 +211,16 @@ public static CloudFetchReader CreateStatementExecutionReader(
                 resultQueue,
                 config);
 
-            // Start the download manager
+            // Start the download manager with no external cancellation token.
+            //
+            // Unlike the Thrift path above, the SEA/StatementExecution reader is not linked into a
+            // connection-shutdown cascade: StatementExecutionConnection/StatementExecutionStatement
+            // have no CloudFetchShutdownToken equivalent (this factory receives only an
+            // ITracingStatement, not a DatabricksConnection), so there is no token to thread through.
+            // The connection-dispose ⊃ CloudFetch teardown fix in this PR therefore covers the Thrift
+            // path only, which is where the reported flake occurs. Extending the same cascade to SEA
+            // means introducing a shutdown-token cascade on the StatementExecution types and plumbing
+            // it here — a larger, protocol-specific change that is intentionally left out of scope.
             downloadManager.StartAsync().Wait();
 
             // Add telemetry tag for compression