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

csharp/src/DatabricksConnection.cs

@@ -122,6 +122,37 @@ 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
+        {
+            get
+            {
+                // Defensive against a read after Dispose(bool) has disposed the source: return
+                // CancellationToken.None rather than throwing, matching the sibling
+                // DatabricksStatement.CloudFetchStatementToken and CloudFetchDownloadManager.PipelineToken
+                // so the three tokens behave symmetrically.
+                try
+                {
+                    return _cloudFetchShutdownCts.Token;
+                }
+                catch (ObjectDisposedException)
+                {
+                    return CancellationToken.None;
+                }
+            }
+        }
+
         // Telemetry
         private IConnectionTelemetry _telemetry = NoOpConnectionTelemetry.Instance;
         // Stopwatch covering the connection lifetime; started at construction and used to
@@ -1183,6 +1214,26 @@ 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.
+                // Best-effort: teardown must continue even if Cancel() throws. Cancel() invokes
+                // registered cancellation callbacks synchronously and rethrows a faulting one wrapped
+                // in AggregateException (not ObjectDisposedException); letting that escape here would
+                // skip the HttpClient/session teardown below and leak them.
+                try { _cloudFetchShutdownCts.Cancel(); }
+                catch (Exception ex)
+                {
+                    Activity.Current?.AddEvent(new ActivityEvent("cloudfetch.shutdown.cancel.error",
+                        tags: new ActivityTagsCollection
+                        {
+                            { "error.type", ex.GetType().Name },
+                            { "error.message", ex.Message }
+                        }));
+                }
+
                 // 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 +1261,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,56 @@ 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)
+                {
+                    // Defensive against a read after Dispose(bool) has disposed the source: return
+                    // CancellationToken.None rather than throwing, matching the sibling
+                    // CloudFetchDownloadManager.PipelineToken so the two tokens behave symmetrically.
+                    try
+                    {
+                        return _cloudFetchStatementCts.Token;
+                    }
+                    catch (ObjectDisposedException)
+                    {
+                        return CancellationToken.None;
+                    }
+                }
+            }
+        }
+
         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 +184,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 +355,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 +389,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 +1519,42 @@ 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)
+                {
+                    // Best-effort: Cancel() runs cancellation callbacks synchronously and rethrows a
+                    // faulting one wrapped in AggregateException (not ObjectDisposedException); letting
+                    // that escape would skip the CTS Dispose and telemetry emission below.
+                    try { _cloudFetchStatementCts.Cancel(); }
+                    catch (ObjectDisposedException)
+                    {
+                        // Expected on a repeated Dispose(): the source was already disposed below on
+                        // the first pass. Dispose(bool) has no idempotency guard, so this is a normal
+                        // double-dispose, not an error — swallow silently (no error event), matching
+                        // the CloudFetchStatementToken getter's ObjectDisposedException handling.
+                    }
+                    catch (Exception ex)
+                    {
+                        Activity.Current?.AddEvent(new ActivityEvent("cloudfetch.statement.cancel.error",
+                            tags: new ActivityTagsCollection
+                            {
+                                { "error.type", ex.GetType().Name },
+                                { "error.message", ex.Message }
+                            }));
+                    }
+
+                    // 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 +1616,32 @@ 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)
+                {
+                    // Best-effort: don't let a faulting cancellation callback (surfaced as
+                    // AggregateException, not ObjectDisposedException) skip base.Cancel() and the
+                    // telemetry emission in the finally below.
+                    try { _cloudFetchStatementCts.Cancel(); }
+                    catch (Exception ex)
+                    {
+                        Activity.Current?.AddEvent(new ActivityEvent("cloudfetch.statement.cancel.error",
+                            tags: new ActivityTagsCollection
+                            {
+                                { "error.type", ex.GetType().Name },
+                                { "error.message", ex.Message }
+                            }));
+                    }
+                }
+
                 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);