diff --git a/csharp/src/DatabricksConnection.cs b/csharp/src/DatabricksConnection.cs index 5bc615768..1fd05bab3 100644 --- a/csharp/src/DatabricksConnection.cs +++ b/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(); + + /// + /// 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. + /// + 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,33 @@ 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 (ObjectDisposedException) + { + // Expected on a repeated Dispose(): the source was already disposed in the + // finally 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 DatabricksStatement.Dispose's ObjectDisposedException handling. + } + 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 +1268,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. diff --git a/csharp/src/DatabricksStatement.cs b/csharp/src/DatabricksStatement.cs index cd2427441..13d7e79f5 100644 --- a/csharp/src/DatabricksStatement.cs +++ b/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(); + + /// + /// 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. + /// + 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) } } + /// + /// Recreates the statement-lifetime CloudFetch cancellation source (re-linked to the + /// connection's shutdown token) at the start of each execution. + /// is reusable (settable + repeated Execute), and / + /// 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 HiveServer2Statement._executeTokenSource is + /// refreshed per-execute so a statement stays reusable after cancel. + /// + 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 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,39 @@ 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 (ObjectDisposedException) + { + // Cancel() is explicitly supported from another thread, so it can race (or + // follow) a Dispose() that already disposed this source. A disposed-source + // Cancel() throwing ObjectDisposedException is benign teardown, not an error + // — swallow silently (no error event), matching Dispose(bool)'s 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 } + })); + } + } + base.Cancel(); } catch (Exception ex) diff --git a/csharp/src/Reader/CloudFetch/CloudFetchDownloadManager.cs b/csharp/src/Reader/CloudFetch/CloudFetchDownloadManager.cs index 0df802b2c..be833fd4b 100644 --- a/csharp/src/Reader/CloudFetch/CloudFetchDownloadManager.cs +++ b/csharp/src/Reader/CloudFetch/CloudFetchDownloadManager.cs @@ -77,6 +77,27 @@ public CloudFetchDownloadManager( /// public bool HasMoreResults => !_downloader.IsCompleted || !_resultQueue.IsCompleted; + /// + public CancellationToken PipelineToken + { + get + { + var cts = _cancellationTokenSource; + if (cts == null) + { + return CancellationToken.None; + } + try + { + return cts.Token; + } + catch (ObjectDisposedException) + { + return CancellationToken.None; + } + } + } + /// public async Task GetNextDownloadedFileAsync(CancellationToken cancellationToken) { @@ -120,7 +141,7 @@ public CloudFetchDownloadManager( } /// - 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); diff --git a/csharp/src/Reader/CloudFetch/CloudFetchReader.cs b/csharp/src/Reader/CloudFetch/CloudFetchReader.cs index 2d5435aac..f0a8425d6 100644 --- a/csharp/src/Reader/CloudFetch/CloudFetchReader.cs +++ b/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,15 @@ public CloudFetchReader( try { // Get the next downloaded file - this.currentDownloadResult = await this.downloadManager.GetNextDownloadedFileAsync(cancellationToken); + this.currentDownloadResult = await this.downloadManager.GetNextDownloadedFileAsync(token); + + // Distinguish a cancelled null from a genuine end-of-results null. + // On cancellation (statement Cancel() / connection Dispose()) the + // download manager returns null without surfacing an error, so without + // this check a cancel between chunks would look like a clean EOF and + // silently present a truncated result set as a completed query. + token.ThrowIfCancellationRequested(); + if (this.currentDownloadResult == null) { Activity.Current?.AddEvent("cloudfetch.reader_no_more_files", [ @@ -176,7 +193,11 @@ public CloudFetchReader( new("chunk_row_count", this.currentDownloadResult.RowCount) ]); - await this.currentDownloadResult.DownloadCompletedTask; + // Wait for this chunk's download to complete, but observe the token so + // teardown is prompt even if a download hangs without honoring its own + // token. Task.WaitAsync(token) isn't available on netstandard2.0/net472, + // so race the download against an infinite delay tied to the token. + await AwaitWithCancellationAsync(this.currentDownloadResult.DownloadCompletedTask, token); // Track bytes downloaded for telemetry _totalBytesDownloaded += this.currentDownloadResult.Size; @@ -219,6 +240,44 @@ public CloudFetchReader( }); } + /// + /// Awaits while observing , throwing + /// if the token fires first. Provides the + /// equivalent of Task.WaitAsync(token), which is unavailable on netstandard2.0/net472, + /// so a hung download can't leave the reader parked after a statement cancel / dispose. + /// + internal static async Task AwaitWithCancellationAsync(Task task, CancellationToken token) + { + if (task.IsCompleted || !token.CanBeCanceled) + { + await task; + return; + } + + // Race the download against an infinite delay tied to the token. The linked CTS + // lets us cancel the delay once the download wins so we don't leak a pending timer. + using var delayCts = CancellationTokenSource.CreateLinkedTokenSource(token); + var delayTask = Task.Delay(Timeout.Infinite, delayCts.Token); + var completed = await Task.WhenAny(task, delayTask).ConfigureAwait(false); + if (completed != task) + { + // The token won the race, so we're about to abandon the download without + // awaiting it. Attach a continuation that observes its eventual fault (e.g. a + // download failing against a torn-down HttpClient during cancel/dispose) so the + // exception doesn't resurface later via TaskScheduler.UnobservedTaskException. + _ = task.ContinueWith( + t => { _ = t.Exception; }, + CancellationToken.None, + TaskContinuationOptions.OnlyOnFaulted | TaskContinuationOptions.ExecuteSynchronously, + TaskScheduler.Default); + token.ThrowIfCancellationRequested(); + } + + // Cancel the delay to release its timer, then observe the download's result/exception. + delayCts.Cancel(); + await task; + } + /// /// Cleans up the current reader and download result, resetting chunk-level tracking. /// diff --git a/csharp/src/Reader/CloudFetch/CloudFetchReaderFactory.cs b/csharp/src/Reader/CloudFetch/CloudFetchReaderFactory.cs index 949c49003..f59c7818e 100644 --- a/csharp/src/Reader/CloudFetch/CloudFetchReaderFactory.cs +++ b/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,8 +211,18 @@ public static CloudFetchReader CreateStatementExecutionReader( resultQueue, config); - // Start the download manager - downloadManager.StartAsync().Wait(); + // Start the download manager linked to the statement's CloudFetch token (itself linked to + // the SEA connection's shutdown token), mirroring the Thrift path: closing the connection, + // or cancelling/disposing this statement, tears down the pipeline and unblocks the reader. + // This is the SEA path, so the statement is always a StatementExecutionStatement; guard the + // cast so an unexpected future caller fails with a clear message rather than InvalidCastException. + if (statement is not StatementExecutionStatement seaStatement) + { + throw new InvalidOperationException( + $"{nameof(CreateStatementExecutionReader)} requires a {nameof(StatementExecutionStatement)}, but received {statement?.GetType().Name ?? "null"}."); + } + + downloadManager.StartAsync(seaStatement.CloudFetchStatementToken).Wait(); // Add telemetry tag for compression Activity.Current?.SetTag(StatementExecutionEvent.ResultCompressionEnabled, isLz4Compressed); diff --git a/csharp/src/Reader/CloudFetch/ICloudFetchInterfaces.cs b/csharp/src/Reader/CloudFetch/ICloudFetchInterfaces.cs index e20de220c..253f256c7 100644 --- a/csharp/src/Reader/CloudFetch/ICloudFetchInterfaces.cs +++ b/csharp/src/Reader/CloudFetch/ICloudFetchInterfaces.cs @@ -274,8 +274,12 @@ internal interface ICloudFetchDownloadManager : IDisposable /// /// Starts the download manager. /// + /// + /// 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. + /// /// A task representing the asynchronous operation. - Task StartAsync(); + Task StartAsync(CancellationToken cancellationToken = default); /// /// Stops the download manager. @@ -288,6 +292,14 @@ internal interface ICloudFetchDownloadManager : IDisposable /// bool HasMoreResults { get; } + /// + /// The pipeline's cancellation token (linked to the token passed to ). + /// The reader observes this so a cancel/dispose stops the read promptly even while draining + /// already-buffered chunks — not just when it next blocks for a download. Returns + /// before start or after stop/dispose. + /// + CancellationToken PipelineToken { get; } + /// /// Gets the aggregated chunk metrics from the downloader. /// Returns a snapshot of current metrics that can be safely passed to telemetry. diff --git a/csharp/src/StatementExecution/StatementExecutionConnection.cs b/csharp/src/StatementExecution/StatementExecutionConnection.cs index 88d0ea42c..c98f05017 100644 --- a/csharp/src/StatementExecution/StatementExecutionConnection.cs +++ b/csharp/src/StatementExecution/StatementExecutionConnection.cs @@ -16,6 +16,7 @@ using System; using System.Collections.Generic; +using System.Diagnostics; using System.Linq; using System.Net.Http; using System.Threading; @@ -28,6 +29,8 @@ using Apache.Arrow; using Apache.Arrow.Adbc; using Apache.Arrow.Adbc.Tracing; +using Apache.Arrow.Adbc.Telemetry.Traces.Listeners; +using Apache.Arrow.Adbc.Telemetry.Traces.Listeners.FileListener; using Apache.Arrow.Ipc; using Apache.Arrow.Types; using static Apache.Arrow.Adbc.AdbcConnection; @@ -48,6 +51,35 @@ internal class StatementExecutionConnection : TracingConnection, IGetObjectsData private string? _schema; private readonly HttpClient _httpClient; private readonly HttpClient _cloudFetchHttpClient; // Separate HttpClient without auth headers for CloudFetch downloads + + // Connection-scoped shutdown signal for in-flight CloudFetch pipelines (SEA). Mirrors the + // Thrift DatabricksConnection: cancelled at the start of Dispose so any active download + // manager (which links this token into its own cancellation source) tears down, and + // statements link their per-statement CloudFetch token to it for the full + // connection ⊃ statement ⊃ cloudfetch cancel cascade. + private readonly CancellationTokenSource _cloudFetchShutdownCts = new CancellationTokenSource(); + + /// + /// Token cancelled when this connection is disposed. SEA CloudFetch download managers link + /// this into their pipeline source so connection close tears down in-flight downloads. + /// + internal CancellationToken CloudFetchShutdownToken + { + get + { + // Defensive against a read after Dispose() has disposed the source: return + // CancellationToken.None rather than throwing, matching the Thrift DatabricksConnection. + try + { + return _cloudFetchShutdownCts.Token; + } + catch (ObjectDisposedException) + { + return CancellationToken.None; + } + } + } + private readonly IReadOnlyDictionary _properties; private readonly bool _ownsHttpClient; @@ -99,6 +131,13 @@ internal class StatementExecutionConnection : TracingConnection, IGetObjectsData /// Connection properties. /// Optional shared memory stream manager. /// Optional shared LZ4 buffer pool. + // Per-connection id tagged onto the shared ActivitySource so the file trace listener + // captures only this connection's activities. Statements reuse this same source + // (TracingStatement copies connection.Trace), so their spans match too. Mirrors the + // file-listener wiring HiveServer2Connection does for the Thrift path. + private readonly string _traceInstanceId = Guid.NewGuid().ToString("N"); + private readonly FileActivityListener? _fileActivityListener; + public StatementExecutionConnection( IReadOnlyDictionary properties, Microsoft.IO.RecyclableMemoryStreamManager? memoryStreamManager = null, @@ -135,6 +174,11 @@ private StatementExecutionConnection( _properties = properties ?? throw new ArgumentNullException(nameof(properties)); _ownsHttpClient = ownsHttpClient; + // Activate the adbcfile trace listener for this connection when the exporter is + // requested (adbc.traces.exporter / OTEL_TRACES_EXPORTER = adbcfile). Without this the + // SEA path produced no local trace file at all — only the Thrift path wired it up. + TryInitTracerProvider(out _fileActivityListener); + // Parse configuration - check for URI first (same as Thrift protocol) properties.TryGetValue(AdbcOptions.Uri, out var uri); properties.TryGetValue(SparkParameters.HostName, out var hostName); @@ -1176,6 +1220,29 @@ public override void Dispose() { this.TraceActivity(activity => { + // Signal in-flight CloudFetch pipelines to stop before tearing down the transport, + // so a reader parked on a download unblocks instead of failing on the disposed client. + // Best-effort: Cancel() runs cancellation callbacks synchronously and rethrows a + // faulting one wrapped in AggregateException (not ObjectDisposedException); letting + // that escape would skip the HttpClient/session teardown below and leak them. + try { _cloudFetchShutdownCts.Cancel(); } + catch (ObjectDisposedException) + { + // Expected on a repeated Dispose(): the source was already disposed below on + // the first pass. Dispose() has no idempotency guard, so this is a normal + // double-dispose, not an error — swallow silently (no error event), matching + // DatabricksStatement.Dispose's ObjectDisposedException handling. + } + catch (Exception ex) + { + activity?.AddEvent(new System.Diagnostics.ActivityEvent("cloudfetch.shutdown.cancel.error", + tags: new System.Diagnostics.ActivityTagsCollection + { + { "error.type", ex.GetType().Name }, + { "error.message", ex.Message } + })); + } + activity?.SetTag("session_id", _sessionId); activity?.SetTag("warehouse_id", _warehouseId); @@ -1209,8 +1276,13 @@ public override void Dispose() // Dispose the CloudFetch HTTP client (we always own it) _cloudFetchHttpClient.Dispose(); + _cloudFetchShutdownCts.Dispose(); + _sessionLock.Dispose(); }); + + // Deactivate the file trace listener last, so the dispose span above is still captured. + _fileActivityListener?.Dispose(); } // TracingConnection provides IActivityTracer implementation @@ -1218,5 +1290,26 @@ public override void Dispose() public override string AssemblyVersion => GetType().Assembly.GetName().Version?.ToString() ?? "1.0.0"; public override string AssemblyName => "AdbcDrivers.Databricks"; + + /// + /// Tags the shared ActivitySource with a per-connection id so the file trace listener + /// (see ) captures only this connection's spans. + /// Statements reuse this same source, so their spans carry the tag too. + /// + public override IEnumerable>? GetActivitySourceTags(IReadOnlyDictionary properties) + { + IEnumerable>? tags = base.GetActivitySourceTags(properties); + tags ??= []; + tags = tags.Concat([new(_traceInstanceId, null)]); + return tags; + } + + private bool TryInitTracerProvider(out FileActivityListener? fileActivityListener) + { + _properties.TryGetValue(ListenersOptions.Exporter, out string? exporterOption); + // This listener only listens for activity from this specific connection instance. + bool shouldListenTo(ActivitySource source) => source.Tags?.Any(t => ReferenceEquals(t.Key, _traceInstanceId)) == true; + return FileActivityListener.TryActivateFileListener(AssemblyName, exporterOption, out fileActivityListener, shouldListenTo: shouldListenTo); + } } } diff --git a/csharp/src/StatementExecution/StatementExecutionStatement.cs b/csharp/src/StatementExecution/StatementExecutionStatement.cs index 06b625f26..30e6bbe41 100644 --- a/csharp/src/StatementExecution/StatementExecutionStatement.cs +++ b/csharp/src/StatementExecution/StatementExecutionStatement.cs @@ -25,6 +25,7 @@ using AdbcDrivers.Databricks.Reader.CloudFetch; using AdbcDrivers.Databricks.StatementExecution.MetadataCommands; using AdbcDrivers.Databricks.Result; +using AdbcDrivers.Databricks.Telemetry.TagDefinitions; using AdbcDrivers.HiveServer2; using AdbcDrivers.HiveServer2.Hive2; using Apache.Arrow; @@ -89,6 +90,39 @@ internal class StatementExecutionStatement : TracingStatement private readonly object _cancelLock = new(); private CancellationTokenSource? _executeCts; + // Statement-lifetime CloudFetch cancellation, linked to the connection's shutdown token, so + // the connection ⊃ statement ⊃ cloudfetch cancel cascade holds on the SEA path too. Distinct + // from _executeCts, which is disposed when execution returns and so cannot cover the later + // CloudFetch result-fetch phase. Refreshed per-execute (Cancel()/Dispose() cancel it + // permanently, so a reused statement needs a fresh one) via RefreshCloudFetchStatementCts(). + private CancellationTokenSource _cloudFetchStatementCts; + + /// + /// 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. + /// + internal CancellationToken CloudFetchStatementToken + { + get + { + // Under _cancelLock so a concurrent re-execute's field swap (RefreshCloudFetchStatementCts) + // can't be read half-applied; guard ObjectDisposedException → None for a read after + // Dispose. Mirrors the Thrift DatabricksStatement. + lock (_cancelLock) + { + try + { + return _cloudFetchStatementCts.Token; + } + catch (ObjectDisposedException) + { + return CancellationToken.None; + } + } + } + } + // Metadata command support private bool _isMetadataCommand; private bool _escapePatternWildcards; @@ -178,6 +212,7 @@ public StatementExecutionStatement( : base(connection) { _connection = connection ?? throw new ArgumentNullException(nameof(connection)); + _cloudFetchStatementCts = CancellationTokenSource.CreateLinkedTokenSource(connection.CloudFetchShutdownToken); _client = client ?? throw new ArgumentNullException(nameof(client)); _sessionId = sessionId; _warehouseId = warehouseId ?? throw new ArgumentNullException(nameof(warehouseId)); @@ -301,6 +336,10 @@ public async Task ExecuteQueryAsync( CancellationToken cancellationToken = default, bool isMetadataExecution = false) { + // Refresh the CloudFetch cancellation source before the reader captures it, so a prior + // Cancel() doesn't leave the next read starting cancelled (the statement is reusable). + RefreshCloudFetchStatementCts(); + if (_isMetadataCommand) { return await ExecuteMetadataCommandAsync(cancellationToken).ConfigureAwait(false); @@ -378,7 +417,18 @@ private async Task EnsureCatalogScopedAsync(CancellationToken cancellationToken) _connection.UpdateCurrentCatalog(catalog); } - private async Task ExecuteQueryInternalAsync(CancellationToken cancellationToken, bool isMetadataExecution) + private Task ExecuteQueryInternalAsync(CancellationToken cancellationToken, bool isMetadataExecution) + { + // Wrap execute + poll + reader-creation in a named span so the SEA path emits the same + // kind of local trace as Thrift's HiveServer2Statement.ExecuteStatementAsync. Without + // this, the whole SEA query path is invisible in the adbcfile trace (only Dispose and + // the cloudfetch.* read events show up). + return this.TraceActivityAsync( + activity => ExecuteQueryInternalCoreAsync(activity, cancellationToken, isMetadataExecution), + activityName: nameof(StatementExecutionStatement) + "." + nameof(ExecuteQueryInternalAsync)); + } + + private async Task ExecuteQueryInternalCoreAsync(Activity? activity, CancellationToken cancellationToken, bool isMetadataExecution) { // If the caller explicitly scoped this statement to a catalog, set the session's // current catalog first via USE CATALOG so a 2-level `schema`.`table` name @@ -421,6 +471,11 @@ private async Task ExecuteQueryInternalAsync(CancellationToken canc // CANCELED: user canceled; can come from explicit cancel call, or timeout with on_wait_timeout=CANCEL // CLOSED: execution successful, and statement closed; result no longer available for fetch var state = response.Status?.State; + activity?.SetTag(StatementExecutionEvent.StatementId, response.StatementId); + activity?.SetTag(StatementExecutionEvent.ResultFormat, _resultFormat); + activity?.SetTag(StatementExecutionEvent.ResultCompressionEnabled, !string.IsNullOrEmpty(_resultCompression)); + activity?.AddEvent(new ActivityEvent("statement.execute.submitted", + tags: new ActivityTagsCollection { { "initial_state", state ?? "(null)" } })); if (state == "PENDING" || state == "RUNNING") { response = await PollWithTimeoutAsync(response.StatementId, cancellationToken).ConfigureAwait(false); @@ -493,6 +548,12 @@ private async Task ExecuteQueryInternalAsync(CancellationToken canc // Return query result - use 0 if row count is not available long rowCount = response.Manifest?.TotalRowCount ?? 0; + activity?.AddEvent(new ActivityEvent("statement.execute.completed", + tags: new ActivityTagsCollection + { + { "final_state", state ?? "(null)" }, + { "result.row_count", rowCount } + })); return new QueryResult(rowCount, reader); } @@ -538,6 +599,12 @@ private async Task PollWithTimeoutAsync(string stateme /// private async Task PollUntilCompleteAsync(string statementId, CancellationToken cancellationToken) { + // Capture the ambient execute span (opened by ExecuteQueryInternalAsync) so aggregate + // poll metrics can be stamped on it once a terminal state is reached. + Activity? executeActivity = Activity.Current; + int pollCount = 0; + long pollLatencyMs = 0; + while (true) { // Check for cancellation before each polling iteration @@ -549,8 +616,27 @@ private async Task PollUntilCompleteAsync(string state // Check for cancellation after delay cancellationToken.ThrowIfCancellationRequested(); - // Get statement status - var response = await _client.GetStatementAsync(statementId, cancellationToken).ConfigureAwait(false); + // Get statement status. Trace each poll (count + latency + state) so the SEA path + // exposes the same polling signal Thrift's DatabricksOperationStatusPoller does. + var response = await this.TraceActivityAsync(async pollActivity => + { + Stopwatch pollStopwatch = Stopwatch.StartNew(); + var r = await _client.GetStatementAsync(statementId, cancellationToken).ConfigureAwait(false); + pollStopwatch.Stop(); + + pollCount++; + pollLatencyMs += pollStopwatch.ElapsedMilliseconds; + + pollActivity?.SetTag(StatementExecutionEvent.PollCount, pollCount); + pollActivity?.AddEvent(new ActivityEvent("statement.poll_status", + tags: new ActivityTagsCollection + { + { StatementExecutionEvent.PollCount, pollCount }, + { StatementExecutionEvent.PollLatencyMs, pollStopwatch.ElapsedMilliseconds }, + { "operation_state", r.Status?.State ?? "(null)" } + })); + return r; + }, activityName: nameof(StatementExecutionStatement) + "." + nameof(PollUntilCompleteAsync)).ConfigureAwait(false); // Convert GetStatementResponse to ExecuteStatementResponse var executeResponse = new ExecuteStatementResponse @@ -568,6 +654,9 @@ private async Task PollUntilCompleteAsync(string state state == "CANCELED" || state == "CLOSED") { + // Stamp aggregate polling metrics on the execute span. + executeActivity?.SetTag(StatementExecutionEvent.PollCount, pollCount); + executeActivity?.SetTag(StatementExecutionEvent.PollLatencyMs, pollLatencyMs); return executeResponse; } @@ -730,6 +819,10 @@ public async Task ExecuteUpdateAsync(CancellationToken cancellatio throw new InvalidOperationException("SQL query is required"); } + // Refresh the CloudFetch cancellation source (see ExecuteQueryAsync) so a prior Cancel() + // doesn't poison the internal result read this update performs. + RefreshCloudFetchStatementCts(); + var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); lock (_cancelLock) { _executeCts = cts; } try @@ -747,7 +840,17 @@ public async Task ExecuteUpdateAsync(CancellationToken cancellatio } } - private async Task ExecuteUpdateInternalAsync(CancellationToken cancellationToken) + private Task ExecuteUpdateInternalAsync(CancellationToken cancellationToken) + { + // Named span for the SEA update path, mirroring ExecuteQueryInternalAsync — otherwise a + // DML/DDL statement's poll span (PollUntilCompleteAsync) would be orphaned with no parent + // execute span, and its aggregate poll.count/poll.latency_ms would land nowhere. + return this.TraceActivityAsync( + activity => ExecuteUpdateInternalCoreAsync(activity, cancellationToken), + activityName: nameof(StatementExecutionStatement) + "." + nameof(ExecuteUpdateInternalAsync)); + } + + private async Task ExecuteUpdateInternalCoreAsync(Activity? activity, CancellationToken cancellationToken) { // Scope the session to the caller's catalog first (see ExecuteQueryInternalAsync), so a // DML statement with a bare 2-level `schema`.`table` name resolves against it too. @@ -780,6 +883,11 @@ private async Task ExecuteUpdateInternalAsync(CancellationToken ca // Handle query status - poll until complete var state = response.Status?.State; + activity?.SetTag(StatementExecutionEvent.StatementId, response.StatementId); + activity?.SetTag(StatementExecutionEvent.ResultFormat, _resultFormat); + activity?.SetTag(StatementExecutionEvent.ResultCompressionEnabled, !string.IsNullOrEmpty(_resultCompression)); + activity?.AddEvent(new ActivityEvent("statement.execute.submitted", + tags: new ActivityTagsCollection { { "initial_state", state ?? "(null)" } })); if (state == "PENDING" || state == "RUNNING") { response = await PollWithTimeoutAsync(response.StatementId, cancellationToken).ConfigureAwait(false); @@ -818,7 +926,14 @@ private async Task ExecuteUpdateInternalAsync(CancellationToken ca // yields no batch; return -1 per the ADBC convention for "unknown or not // applicable", matching what the Thrift path does. using var reader = CreateReader(response, cancellationToken); - return new UpdateResult(await ReadNumAffectedRowsAsync(reader, cancellationToken).ConfigureAwait(false)); + long affectedRows = await ReadNumAffectedRowsAsync(reader, cancellationToken).ConfigureAwait(false); + activity?.AddEvent(new ActivityEvent("statement.execute.completed", + tags: new ActivityTagsCollection + { + { "final_state", state ?? "(null)" }, + { "num_affected_rows", affectedRows } + })); + return new UpdateResult(affectedRows); } private static async Task ReadNumAffectedRowsAsync(IArrowArrayStream reader, CancellationToken cancellationToken) @@ -850,6 +965,34 @@ private static async Task ReadNumAffectedRowsAsync(IArrowArrayStream reade /// public override void Dispose() { + // Cancel + release the CloudFetch statement CTS first, regardless of whether a statement + // ran, so in-flight downloads stop promptly and its linked registration on the connection + // shutdown token is always freed (this method early-returns below when nothing executed). + lock (_cancelLock) + { + // 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 the statement-close teardown below. + try { _cloudFetchStatementCts.Cancel(); } + catch (ObjectDisposedException) + { + // Expected on a repeated Dispose(): the source was already disposed below on the + // first pass. Dispose() has no idempotency guard, so this is a normal + // double-dispose, not an error — swallow silently (no error event), matching the + // Thrift DatabricksStatement.Dispose and CloudFetchStatementToken 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 } + })); + } + _cloudFetchStatementCts.Dispose(); + } + if (_currentStatementId == null) { return; @@ -897,12 +1040,54 @@ public override void Dispose() }, activityName: nameof(StatementExecutionStatement) + "." + nameof(Dispose)); } + /// + /// Recreates the statement-lifetime CloudFetch cancellation source (re-linked to the + /// connection's shutdown token) at the start of each execution. Cancel()/Dispose() cancel + /// this source permanently; without a refresh a cancel-then-reexecute would start the next + /// CloudFetch read with an already-cancelled token. Mirrors the Thrift DatabricksStatement. + /// + private void RefreshCloudFetchStatementCts() + { + lock (_cancelLock) + { + var previous = _cloudFetchStatementCts; + _cloudFetchStatementCts = CancellationTokenSource.CreateLinkedTokenSource(_connection.CloudFetchShutdownToken); + // Release the prior source's registration on the connection shutdown token. Under + // normal usage the previous result set is fully consumed/disposed before the next + // Execute, so no pipeline is still linked to the old token. Disposing a source does + // not cancel its already-created linked children, so an open prior reader is unaffected. + previous?.Dispose(); + } + } + public override void Cancel() { string? statementId; lock (_cancelLock) { _executeCts?.Cancel(); + // Also cancel the CloudFetch pipeline for this statement: _executeCts is already + // disposed once results are streaming, so without this a Cancel() during CloudFetch + // would leave the downloads running. + // Best-effort: don't let a faulting cancellation callback (surfaced as + // AggregateException, not ObjectDisposedException) escape and skip the remote + // CancelStatement RPC below. + try { _cloudFetchStatementCts.Cancel(); } + catch (ObjectDisposedException) + { + // Cancel() can race (or follow) a Dispose() that already disposed this source. + // A disposed-source Cancel() throwing ObjectDisposedException is benign teardown, + // not an error — swallow silently (no error event), matching Dispose's 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 } + })); + } statementId = _currentStatementId; } if (statementId != null) diff --git a/csharp/test/E2E/ConcurrencyStressTests.cs b/csharp/test/E2E/ConcurrencyStressTests.cs index 1123fc1f3..b89156d79 100644 --- a/csharp/test/E2E/ConcurrencyStressTests.cs +++ b/csharp/test/E2E/ConcurrencyStressTests.cs @@ -434,5 +434,85 @@ public async Task CancelStatement_FromAnotherThread_ShouldStopPromptly() OutputHelper?.WriteLine("Query completed before cancel took effect (race condition — acceptable)."); } } + + /// + /// CONCURRENT-006: Cancel a statement while it is actively fetching results via CloudFetch. + /// Distinct from CONCURRENT-005 (cancel during execution): here cancellation must tear down + /// the CloudFetch download pipeline promptly rather than let it keep downloading. Regression + /// for the connection ⊃ statement ⊃ cloudfetch cancel cascade — statement.Cancel() cancels the + /// statement-lifetime token linked into the pipeline, so the parked reader unblocks instead of + /// running to completion (the per-execute token is already disposed by the time results stream). + /// + [SkippableFact] + public async Task CancelStatement_DuringCloudFetch_ShouldStopPromptly() + { + const int timeoutMs = 30_000; + var queryStarted = new ManualResetEventSlim(false); + Exception? queryException = null; + int rowsRead = 0; + + using AdbcConnection connection = NewConnection(); + using var statement = connection.CreateStatement(); + + // Deliberately huge so a full read takes far longer than timeoutMs: without the fix, + // Cancel() is a no-op during streaming and the read would keep going well past the + // window (test fails); with the fix, Cancel() tears down the pipeline and the read ends + // in well under the window. We only ever read the first batch before cancelling. + statement.SqlQuery = "SELECT * FROM RANGE(1000000000)"; + + var queryTask = Task.Run(async () => + { + try + { + var result = statement.ExecuteQuery(); + using var reader = result.Stream; + + while (true) + { + var batch = await reader.ReadNextRecordBatchAsync(); + if (batch == null) break; + + Interlocked.Add(ref rowsRead, batch.Length); + + // Signal after the first batch so we know CloudFetch streaming is active. + if (!queryStarted.IsSet) + { + queryStarted.Set(); + } + } + } + catch (Exception ex) + { + queryException = ex; + } + }); + + bool started = queryStarted.Wait(30_000); + if (!started) + { + statement.Cancel(); + Assert.Fail("Query did not start producing results within 30s"); + } + + OutputHelper?.WriteLine($"Query started, read {rowsRead} rows so far. Cancelling statement..."); + + // Cancel the statement while results are streaming via CloudFetch. + statement.Cancel(); + + var completed = await Task.WhenAny(queryTask, Task.Delay(timeoutMs)); + + OutputHelper?.WriteLine($"Query task completed: {queryTask.IsCompleted}"); + OutputHelper?.WriteLine($"Rows read before cancel: {rowsRead}"); + + if (queryException != null) + { + // A cancellation exception is expected — the pipeline was torn down mid-read. + OutputHelper?.WriteLine($"Query exception (expected): {queryException.GetType().Name}: {queryException.Message}"); + } + + Assert.True(queryTask.IsCompleted, + $"Statement cancel during CloudFetch did not stop the read within {timeoutMs}ms. " + + $"Query completed: {queryTask.IsCompleted}, rows read: {rowsRead}."); + } } } diff --git a/csharp/test/Unit/DatabricksStatementUnitTests.cs b/csharp/test/Unit/DatabricksStatementUnitTests.cs index 574518bcc..ee2a35b5c 100644 --- a/csharp/test/Unit/DatabricksStatementUnitTests.cs +++ b/csharp/test/Unit/DatabricksStatementUnitTests.cs @@ -170,5 +170,32 @@ public void GetMetadataOperationType_IsCaseInsensitive(string command) { Assert.NotNull(DatabricksStatement.GetMetadataOperationType(command)); } + + /// + /// Regression: Cancel() cancels the statement-lifetime CloudFetch token permanently, but the + /// statement is reusable (repeated ExecuteQuery). Each execution refreshes the CloudFetch CTS, + /// so a cancel-then-reexecute must start the next CloudFetch read with a fresh, non-cancelled + /// token rather than the poisoned one. + /// + [Fact] + public void RefreshCloudFetchStatementCts_AfterCancel_YieldsFreshUncancelledToken() + { + using var statement = CreateStatement(); + + // Fresh statement: token is live. + Assert.False(statement.CloudFetchStatementToken.IsCancellationRequested); + + // Cancel() cancels the statement-lifetime CloudFetch source. + statement.Cancel(); + Assert.True(statement.CloudFetchStatementToken.IsCancellationRequested); + + // The next execution refreshes the source; the new token must not be born cancelled. + statement.RefreshCloudFetchStatementCts(); + Assert.False(statement.CloudFetchStatementToken.IsCancellationRequested); + + // And a subsequent Cancel() still cancels the refreshed source. + statement.Cancel(); + Assert.True(statement.CloudFetchStatementToken.IsCancellationRequested); + } } } diff --git a/csharp/test/Unit/Reader/CloudFetch/CloudFetchDownloadManagerTests.cs b/csharp/test/Unit/Reader/CloudFetch/CloudFetchDownloadManagerTests.cs index 0c4087cae..c17649c3e 100644 --- a/csharp/test/Unit/Reader/CloudFetch/CloudFetchDownloadManagerTests.cs +++ b/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(); } + + /// + /// 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. + /// + [Fact] + public async Task StartAsync_TokenCancelled_UnblocksReaderWaitingForNextFile() + { + // Arrange — a real downloader backed by a tracer-providing statement mock. + var mockStatement = new Mock(); + 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(); + mockFetcher.Setup(f => f.HasError).Returns(false); + mockFetcher.Setup(f => f.Error).Returns((Exception?)null); + mockFetcher.Setup(f => f.StartAsync(It.IsAny())).Returns(Task.CompletedTask); + mockFetcher.Setup(f => f.StopAsync()).Returns(Task.CompletedTask); + + var mockMemoryManager = new Mock(); + var downloadQueue = new BlockingCollection(new ConcurrentQueue(), 10); + var resultQueue = new BlockingCollection(new ConcurrentQueue(), 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(); + } } } diff --git a/csharp/test/Unit/Reader/CloudFetch/CloudFetchReaderCancellationTests.cs b/csharp/test/Unit/Reader/CloudFetch/CloudFetchReaderCancellationTests.cs new file mode 100644 index 000000000..f86661068 --- /dev/null +++ b/csharp/test/Unit/Reader/CloudFetch/CloudFetchReaderCancellationTests.cs @@ -0,0 +1,108 @@ +/* +* 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 AdbcDrivers.Databricks.Reader.CloudFetch; +using Xunit; + +namespace AdbcDrivers.Databricks.Tests.Reader.CloudFetch +{ + /// + /// Regression tests for , which + /// guarantees the reader unparks from an in-flight download wait when the statement is + /// cancelled / the connection is disposed, even if the download never completes on its own. + /// + public class CloudFetchReaderCancellationTests + { + [Fact] + public async Task AwaitWithCancellation_TokenCancelledWhileDownloadHangs_Throws() + { + // A download that never completes on its own (simulates a hung download that + // does not honor its own token). + var neverCompletes = new TaskCompletionSource(); + using var cts = new CancellationTokenSource(); + + var waitTask = CloudFetchReader.AwaitWithCancellationAsync(neverCompletes.Task, cts.Token); + Assert.False(waitTask.IsCompleted); + + // Cancelling the token must promptly unpark the wait. + cts.Cancel(); + + await Assert.ThrowsAnyAsync(() => waitTask); + } + + [Fact] + public async Task AwaitWithCancellation_DownloadCompletesFirst_Returns() + { + var tcs = new TaskCompletionSource(); + using var cts = new CancellationTokenSource(); + + var waitTask = CloudFetchReader.AwaitWithCancellationAsync(tcs.Task, cts.Token); + Assert.False(waitTask.IsCompleted); + + tcs.SetResult(true); + + // Completes normally without observing cancellation. + await waitTask; + Assert.False(cts.IsCancellationRequested); + } + + [Fact] + public async Task AwaitWithCancellation_AlreadyCompletedTask_ReturnsImmediately() + { + using var cts = new CancellationTokenSource(); + await CloudFetchReader.AwaitWithCancellationAsync(Task.CompletedTask, cts.Token); + } + + [Fact] + public async Task AwaitWithCancellation_FaultedDownload_PropagatesException() + { + var tcs = new TaskCompletionSource(); + tcs.SetException(new InvalidOperationException("download failed")); + using var cts = new CancellationTokenSource(); + + var ex = await Assert.ThrowsAsync( + () => CloudFetchReader.AwaitWithCancellationAsync(tcs.Task, cts.Token)); + Assert.Equal("download failed", ex.Message); + } + + [Fact] + public async Task AwaitWithCancellation_TokenWinsThenDownloadFaults_ObservesException() + { + // Reproduces the cancel/dispose teardown case: the token wins the race (so the + // reader abandons the wait with an OperationCanceledException) and the in-flight + // download subsequently fails against the torn-down HttpClient. The abandoned + // task's fault must be observed so it does not resurface via + // TaskScheduler.UnobservedTaskException. + var neverCompletes = new TaskCompletionSource(); + using var cts = new CancellationTokenSource(); + + var waitTask = CloudFetchReader.AwaitWithCancellationAsync(neverCompletes.Task, cts.Token); + cts.Cancel(); + await Assert.ThrowsAnyAsync(() => waitTask); + + // The download now fails after having been abandoned. + neverCompletes.SetException(new InvalidOperationException("download failed after cancel")); + + // The observing continuation runs synchronously on fault, so the exception is + // observed by the time SetException returns. + Assert.True(neverCompletes.Task.IsFaulted); + Assert.NotNull(neverCompletes.Task.Exception); + } + } +}