fix(csharp): tear down in-flight CloudFetch on connection dispose and statement cancel/dispose - #659
eric-wang-1990 wants to merge 33 commits into
Conversation
Closing a DatabricksConnection while a CloudFetch query was streaming could hang the reader. The CloudFetch pipeline ran on background tasks driven by a cancellation token that CloudFetchDownloadManager.StartAsync created itself and linked to nothing external. The only thing that cancelled it was disposing the reader — but the reader is owned by the query task, which is blocked awaiting the next chunk. connection.Dispose() disposed the shared HttpClient and closed the session, but never cancelled the pipeline or reached the active reader, so the in-flight download failed on the dead HttpClient and DownloadFileAsync retry-spun for up to the retry timeout (minutes) while the reader stayed parked on DownloadCompletedTask — blowing the 60s limit in CloseConnection_DuringCloudFetch_ShouldNotHang. Intermittent because it only fires when Dispose lands mid-download. Give connection shutdown a way to cancel the pipeline: add a connection-scoped CancellationTokenSource cancelled at the top of Dispose (before the HttpClient is torn down), and have CloudFetchDownloadManager.StartAsync link the caller's token into its own source. On dispose the download loop exits and completes the result queue (unblocking the reader's Take) and any in-flight download faults DownloadCompletedTask (unblocking the reader's await), so the read task ends promptly instead of spinning on retries. Wired on the Thrift path; SEA uses a separate connection and already passed, so its call keeps the default token. Add a deterministic unit test that cancels the StartAsync token and asserts a reader parked in GetNextDownloadedFileAsync unblocks within a short bound (times out without the fix, passes with it). Co-authored-by: Isaac
![peco-review-bot[bot]](https://avatars.githubusercontent.com/in/3801177?s=60&v=4){kind=small}
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — a well-scoped, correctly-implemented fix for the CloudFetch dispose hang. The linked-token approach is sound, the SEA default-token path stays behavior-equivalent (non-breaking), per-query registrations don't accumulate on the connection token (linked CTS is disposed in the manager's Stop/Dispose), and the new regression test is deterministic (the parked read cannot complete without cancellation, so its timing checks aren't flaky). Only one minor, non-blocking ordering note filed inline.
eric-wang-1990
added
the
engineer-bot
Addresses: - #3876787273 at csharp/src/DatabricksConnection.cs:1241 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
…cade) Extends the connection-level CloudFetch cancellation to the statement level, so the full connection ⊃ statement ⊃ cloudfetch cancel cascade holds: - DatabricksStatement gets a statement-lifetime CancellationTokenSource created linked to the connection's shutdown token, cancelled in Cancel() and Dispose() and disposed in Dispose(). It is distinct from the base HiveServer2Statement._executeTokenSource, which is disposed when ExecuteQuery() returns and so cannot cover the later CloudFetch result-fetch phase. - The Thrift reader factory passes this statement token into the pipeline instead of the raw connection token. Because the statement token is linked to the connection token, connection dispose still cancels every statement's downloads (preserving the prior fix), and cancelling/disposing a single statement now stops just its downloads. Also make the reader observe cancellation so a cancel stops the read promptly even while draining already-buffered chunks, not only when it next blocks for a download: CloudFetchDownloadManager exposes the pipeline token (PipelineToken), and CloudFetchReader links it with the caller's token and checks it at the top of each read-loop iteration. Without this, a statement cancel during a large stream would keep returning buffered rows until the in-memory buffer drained. Add E2E CancelStatement_DuringCloudFetch_ShouldStopPromptly: reads the first CloudFetch batch of a huge RANGE, cancels the statement, and asserts the read ends well within the window (it does not without the fix — the read keeps going until the timeout). Full unit suite (925) and full-read CloudFetch E2E pass. Co-authored-by: Isaac
eric-wang-1990
changed the title
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — a well-scoped fix for the CloudFetch shutdown hang; the connection⊃statement⊃pipeline linked-CTS cascade is sound, the PipelineToken getter handles the disposed-CTS race, and the reader now observes cancellation while draining buffered chunks. One low-severity consistency issue: the statement's CTS Dispose() is placed after (unguarded by) the telemetry emit, unlike the connection which the author deliberately guarded against exactly that throw.
Addresses: - #3877046084 at csharp/src/DatabricksStatement.cs:1497 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Verdict: 1 Medium · 1 Low
Solid, well-motivated fix for the CloudFetch teardown hang, with good tests. One medium concern: the statement-lifetime CloudFetch CTS is never recreated, so Cancel() permanently poisons the CloudFetch path and breaks statement reuse-after-cancel (F1). Also flagging that the connection-dispose cascade is Thrift-only — the SEA reader path is left untokenized (F2).
Addresses: - #3877068667 at csharp/src/DatabricksStatement.cs:1525 - #3877068674 at csharp/src/Reader/CloudFetch/CloudFetchReaderFactory.cs:119 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Verdict: 1 Medium · 1 Low
Solid, well-reasoned fix for the CloudFetch teardown hang — the linked-CTS cascade, ODE-guarded PipelineToken, and dispose ordering all check out, and the download manager disposes its linked source so no token registration leaks. Two concerns: Cancel() cancels the pipeline CTS after the throwing base.Cancel() RPC (medium — teardown is skipped in the very failure path the fix targets), and RefreshCloudFetchStatementCts disposing the previous CTS can sever a still-open reader's link to the connection shutdown token on re-execute (low).
Addresses: - #3877106059 at csharp/src/DatabricksStatement.cs:1553 - #3877106064 at csharp/src/DatabricksStatement.cs:170 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
CreateThriftReader is Thrift-only and already hard-casts statement.Connection to DatabricksConnection, so the statement is always a DatabricksStatement there. The `statement is DatabricksStatement ? ... : connection.CloudFetchShutdownToken` fallback was unreachable — replace it with a direct cast, matching the existing connection cast. Co-authored-by: Isaac
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — solid, well-documented fix for the CloudFetch teardown hang; the cancel cascade (connection ⊃ statement ⊃ pipeline), the per-execute CTS refresh, and the reader-level prompt-stop are all sound, and the two documented gaps (open-reader re-execute detaching from the connection cascade; SEA path left unlinked) are honestly called out in the code. One Low thread-safety note on the unsynchronized _cloudFetchStatementCts field swap racing a cross-thread Cancel() is filed inline.
Addresses: - #3882435357 at csharp/src/DatabricksStatement.cs:171 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — a clean, well-documented fix for the CloudFetch shutdown hang. The connection ⊃ statement ⊃ cloudfetch cancel cascade is coherent: I confirmed the downloader's retry loop (CloudFetchDownloader.cs:595/:771) actually observes the linked pipeline token, so the cancel truly breaks the retry-spin, and the reader/refresh/lock discipline all hold. Only one minor robustness note: the new CloudFetchStatementToken getter isn't defensive against ObjectDisposedException the way its sibling PipelineToken is (currently unreachable in valid usage).
Addresses: - #3882499934 at csharp/src/DatabricksStatement.cs:134 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — a solid, well-documented fix for the CloudFetch shutdown hang with the connection ⊃ statement ⊃ cloudfetch cancel cascade, good test coverage (unit + E2E), and correct linked-CTS lifecycle handling. One low-severity consistency gap: the new CloudFetchShutdownToken getter lacks the post-dispose ObjectDisposedException→None guard that its two sibling token getters intentionally have. The SEA path being left out of the cascade is clearly documented and reasonable to defer.
Addresses: - #3882538923 at csharp/src/DatabricksConnection.cs:137 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
The Thrift path got the connection ⊃ statement ⊃ cloudfetch cancel cascade; the SEA (StatementExecution) path had the same latent gap — statement.Cancel() during CloudFetch streaming was a no-op (its _executeCts is disposed once results stream, and the pipeline was started with no token), so a cancel left downloads running (verified: an E2E cancel test read 67M rows and never stopped on SEA). Mirror the Thrift wiring so both protocols behave identically: - StatementExecutionConnection gets a CloudFetchShutdownToken, cancelled at the top of Dispose() and disposed at the end. - StatementExecutionStatement gets a statement-lifetime CTS linked to that token, refreshed per-execute (ExecuteQueryAsync/ExecuteUpdateAsync) so a reused statement isn't poisoned by a prior Cancel(), cancelled in Cancel() and Dispose() (before the no-statement early-return so the linked registration is always freed). - CreateStatementExecutionReader passes the statement token into StartAsync, replacing the no-token call (and the comment that scoped SEA out). Verified on a real warehouse: SEA CancelStatement_DuringCloudFetch now stops promptly (~17s vs never), SEA CloseConnection still passes (~1s), SEA full-read CloudFetchE2ETest.TestCloudFetch (10 cases) unaffected; Thrift unchanged; unit suite 926 green. Co-authored-by: Isaac
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — a well-constructed connection ⊃ statement ⊃ cloudfetch cancel cascade with sound locking, per-execute CTS refresh, symmetric Thrift/SEA wiring, and defensive post-dispose token guards; the reader's linked-CTS + top-of-loop ThrowIfCancellationRequested correctly delivers the mid-buffer prompt-stop, and normal reads remain unaffected. One low-severity note inline about the narrow catch (ObjectDisposedException) around CancellationTokenSource.Cancel() in the Dispose teardown paths. Nit (not inline): the new E2E test CancelStatement_DuringCloudFetch_ShouldStopPromptly never disposes its ManualResetEventSlim queryStarted.
Addresses: - #3888803204 at csharp/src/DatabricksConnection.cs:1221 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — a well-constructed, defensively-documented fix for the CloudFetch shutdown hang. The connection ⊃ statement ⊃ pipeline linked-token cascade is correct: per-execute CTS refresh keeps statements reusable after cancel, teardown cancels are best-effort so HttpClient/session cleanup can't be skipped, read-after-dispose is guarded to CancellationToken.None, and the reader now observes the pipeline token plus passes it into GetNextDownloadedFileAsync (I confirmed the parked DownloadCompletedTask await still unblocks because the downloader faults the task on cancel). Only one minor nit: the SEA connection's shutdown-cancel telemetry tag shape diverges from the sibling paths (see inline).
Addresses: - #3888821600 at csharp/src/StatementExecution/StatementExecutionConnection.cs:1216 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — solid, well-documented fix for the CloudFetch shutdown hang. The connection ⊃ statement ⊃ cloudfetch linked-token cascade is correctly wired: the download manager's linked CTS is disposed in StopAsync/Dispose, token getters guard ObjectDisposedException, the per-execute refresh swaps under a lock, and both production StartAsync callers pass the statement token so the cascade is never silently dropped. One low-severity note inline about misleading error telemetry on repeated dispose. Nit (no anchor needed): StartAsync(CancellationToken = default) makes an unlinked pipeline the silent default — fine today since both callers pass a token, but a future caller that omits it would quietly regress the cascade; consider making the token required.
Addresses: - #3888832882 at csharp/src/DatabricksStatement.cs:1538 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Verdict: 1 Low
Solid, carefully-implemented fix — the connection ⊃ statement ⊃ cloudfetch cancel cascade is wired correctly (linked CTS in StartAsync, in-flight downloads faulted on cancel, guarded per-execute CTS refresh, unconditional disposal), and I found no correctness or resource-leak defects in the cancellation/disposal paths. One low-severity behavioral note: a statement.Cancel() while the reader is parked waiting for the next file surfaces as a silent clean EOF (partial results, no exception) rather than an OperationCanceledException, so cancellation is surfaced inconsistently depending on reader timing. Nit (not inline): the new E2E CancelStatement_DuringCloudFetch_ShouldStopPromptly has a fairly tight margin (reported 5–20s vs a 30s timeout) on the SEA path, though it's a live-workspace SkippableFact not run in ordinary CI.
…, drop reader helper The reader parked on IDownloadResult.DownloadCompletedTask for a chunk that was enqueued while its download is still in flight. On statement cancel / connection dispose the download task ends cancelled, but the downloader's ContinueWith left the result's completion source unset, so the reader hung — which is why the reader grew AwaitWithCancellationAsync (a Task.Delay(Timeout.Infinite) race plus unobserved-exception handling). Move the one responsibility to the layer that owns the download: complete the result with OperationCanceledException in the downloader's ContinueWith IsCanceled branch (mirrors the JDBC driver's future.completeExceptionally on shutdown). The reader then unblocks on a plain await this.currentDownloadResult.DownloadCompletedTask. - CloudFetchDownloader: +6 lines (IsCanceled branch). - CloudFetchReader: -43 lines (delete AwaitWithCancellationAsync; restore plain await). - Move the regression coverage to the downloader layer: replace the 5 helper unit tests with one downloader-level test (hanging HttpClient) asserting the parked wait unblocks with a cancellation within 5s. The connection ⊃ statement ⊃ pipeline linked-CTS cascade is unchanged. Co-authored-by: Isaac
There was a problem hiding this comment.
Verdict: 1 Low
Solid, well-reasoned fix — the connection ⊃ statement ⊃ cloudfetch cancel cascade is correctly wired (linked CTS in the download manager, per-execute refresh to keep reusable statements from being poisoned by a prior Cancel(), PipelineToken observed in the read loop, and the downloader completing the in-flight DownloadCompletedTask with OperationCanceledException on cancel). The token getters defensively swallow ObjectDisposedException, the lock discipline around the swap/cancel/dispose of the statement CTS is sound, and the reader's post-GetNextDownloadedFileAsync ThrowIfCancellationRequested() correctly avoids presenting a cancelled fetch as a clean EOF. New unit tests exercise both the parked-reader and in-flight-download cancel paths. Only one minor (Low) polish item inline. One non-blocking scope note: the SEA files also bundle a sizable, unrelated tracing/telemetry expansion (per-connection FileActivityListener wiring, `Ge
[...truncated to keep verdict scannable]
Addresses: - #3908994895 at csharp/src/Reader/CloudFetch/CloudFetchReaderFactory.cs:116 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — the connection ⊃ statement ⊃ cloudfetch cancel cascade is correctly wired on both protocols and well-tested. I verified the non-obvious risks and found them handled: the pipeline token only fires on explicit cancel/dispose (never at normal EOF, so healthy reads aren't truncated), DownloadResult.SetFailed is TrySet-based (no double-completion), straggler re-downloads remain inside DownloadFileAsync's retry loop so the new t.IsCanceled branch can't false-fail them, and the reader's post-GetNextDownloadedFileAsync ThrowIfCancellationRequested() stops a cancelled null being served as clean EOF. One non-blocking scope concern is filed inline (unrelated SEA tracing subsystem in a hotfix). Nit: RefreshCloudFetchStatementCts disposes the previous CTS outside the lock in the Thrift DatabricksStatement but inside the lock in the SEA StatementExecutionStatement — harmless, but worth aligning for consistency since the two are explicitly written to mirror each other.
…tchReader The per-iteration token.ThrowIfCancellationRequested() only made cancellation marginally more prompt while draining an already-downloaded in-memory chunk. It's not needed for correctness: the two blocking calls observe the token (ReadNextRecordBatchAsync throws on its next await), and the post-fetch check after GetNextDownloadedFileAsync — which IS load-bearing, since the downloader returns null rather than throwing on cancel — still guards against a cancelled read being mistaken for a clean end-of-results. Worst case now is the read finishes the current in-memory chunk (bounded, milliseconds) before observing the cancel. Co-authored-by: Isaac
There was a problem hiding this comment.
Verdict: 1 Low
Solid, well-documented fix — the connection ⊃ statement ⊃ cloudfetch cancel cascade is coherent and I found no correctness defect. Verified the key paths: the downloader's new t.IsCanceled → SetFailed(OCE) branch correctly unblocks a reader parked on DownloadCompletedTask (and straggler-cancel is handled internally, so IsCanceled only means real shutdown); the reader's token.ThrowIfCancellationRequested() after GetNextDownloadedFileAsync correctly distinguishes a cancelled null from a genuine EOF so a cancel can't surface a truncated-but-successful result; the per-execute CTS refresh with lock-guarded field swap is sound; and the factory cast guards match the codebase's existing is DatabricksStatement defensive pattern (statement is always the concrete type in practice). One low-severity note (inline) about unrelated SEA tracing/telemetry scope creep bundled into a bugfix. Nit (not inline): `CloudFetchReader.ReadNextRecordBatchAsync
[...truncated to keep verdict scannable]
…xplicitly Both CloudFetchReaderFactory methods took a broad statement interface (IHiveServer2Statement / ITracingStatement) and downcast it to the concrete type (DatabricksStatement / StatementExecutionStatement) solely to read the internal CloudFetchStatementToken, guarding the cast with a throw. Replace that with an explicit CancellationToken parameter on both methods and delete both guards; the callers supply the token directly (the SEA caller is the statement itself; the Thrift caller already downcasts to DatabricksStatement elsewhere). Behavior-preserving: same token, same connection ⊃ statement ⊃ pipeline cancel cascade. Removes the "factory reaches through an interface into a concrete type" smell from both paths and makes the token dependency explicit in each signature. Co-authored-by: Isaac
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — the connection ⊃ statement ⊃ cloudfetch cancel cascade is correct and well-covered by the new unit/E2E tests. The downloader's t.IsCanceled branch correctly completes DownloadCompletedTask with an OCE, straggler cancellation is internally distinguished from pipeline cancellation (effectiveToken.IsCancellationRequested && !cancellationToken.IsCancellationRequested) so it doesn't collide with the new branch, the manager disposes its linked CTS in StopAsync/Dispose (no token-registration leak), and the token getters read into locals to stay race-safe against concurrent dispose. My only substantive concern is scope: an undescribed, CI-untested SEA trace-file/observability feature is bundled into this hang fix (one low finding inline). Nit: SEA RefreshCloudFetchStatementCts disposes the previous CTS inside the lock while the Thrift variant disposes outside it — harmless but inconsistent.
…oll paths" This reverts commit 587b090.
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — a solid, well-documented fix for the CloudFetch teardown hang. The connection ⊃ statement ⊃ cloudfetch cancel cascade is wired symmetrically across both protocols, the downloader's new IsCanceled branch correctly completes the parked DownloadCompletedTask (faulted downloads are still handled by the prior if (t.IsFaulted) and are not miscounted as successful), and the statement-CTS swap/refresh is properly guarded by a lock against cross-thread Cancel()/Dispose(). Only one narrow, pre-existing cancellation gap noted inline (Low).
| /// 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() |
There was a problem hiding this comment.
This smells, why we need this?
There was a problem hiding this comment.
Clarification only — no code change warranted. The method recreates the statement-lifetime CloudFetch CTS at the start of each execution because AdbcStatement is reusable (settable SqlQuery + repeated Execute) while Cancel()/Dispose() cancel the source permanently and a CancellationTokenSource can never be un-cancelled. The CloudFetch pipeline downloads results after ExecuteQuery() returns (past the base _executeTokenSource's lifetime), so a statement-lifetime token is needed for Cancel/Dispose/connection-dispose to tear down in-flight downloads; without the per-execute refresh a cancel-then-reexecute would start the next read with an already-cancelled token. This mirrors the base HiveServer2Statement._executeTokenSource per-execute refresh. Offered to explore a reader-lifetime alternative if the reviewer prefers.
Addresses: - #3909725209 at csharp/src/Reader/CloudFetch/CloudFetchReader.cs:198 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Verdict: 1 Low
Solid, well-reasoned fix — the connection ⊃ statement ⊃ cloudfetch cancel cascade is correctly wired on both protocols and the core hang (a canceled download never completing DownloadCompletedTask) is fixed by SetFailed(OCE) in the downloader continuation, with the reader re-checking the pipeline token after GetNextDownloadedFileAsync to avoid silently truncating a result as clean EOF. Field swaps are lock-guarded against cross-thread Cancel(), per-execute CTS refresh keeps reusable statements usable, and the new unit + E2E tests exercise the key paths. Only one low-severity consistency nit (SEA Dispose cleanup ordering).
Addresses: - #3909767644 at csharp/src/StatementExecution/StatementExecutionConnection.cs:1264 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Verdict: 1 Low
Solid, carefully-documented fix — the connection⊃statement⊃cloudfetch cancel cascade is coherent and the new t.IsCanceled branch in the downloader correctly targets only pipeline teardown (straggler/body-timeout cancels are caught-and-retried inside DownloadFileAsync, so they never surface as a canceled task). The hard (DatabricksStatement)_statement cast in DatabricksCompositeReader.CreateCloudFetchReader is safe because DatabricksConnection.NewReader throws unless the statement is a DatabricksStatement (DatabricksConnection.cs:613-615). One low-severity note inline about an unobserved faulted task on the reader's cancel path. Nit: DatabricksStatement.RefreshCloudFetchStatementCts disposes the previous CTS outside the lock while the SEA counterpart does it inside — harmless but inconsistent.
Addresses: - #3909806356 at csharp/src/Reader/CloudFetch/CloudFetchReader.cs:268 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — a thorough, well-reasoned connection ⊃ statement ⊃ cloudfetch cancel cascade. The pipeline-token linkage, the reader's AwaitWithCancellation netstandard2.0 shim (with unobserved-exception guard), the token.ThrowIfCancellationRequested() after GetNextDownloadedFileAsync (prevents a cancel being mistaken for clean EOF), the downloader's new IsCanceled → SetFailed(OCE) branch, and per-execute CTS refresh under a lock all fit together correctly; the linked CTS is disposed in both StopAsync and Dispose. One Low resource-cleanup asymmetry noted inline. Nit (summary only): the new E2E CancelStatement_DuringCloudFetch_ShouldStopPromptly can wait up to ~60s total (30s queryStarted.Wait + 30s Task.Delay(timeoutMs)) — fine for a SkippableFact but worth tightening if it ever runs under a tighter harness budget.
Addresses: - #3909844457 at csharp/src/DatabricksConnection.cs:1243 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
The engineer-bot re-added a reader-side await helper (b3bfd8b/2d95ce60) to make the reader observe the caller's own token during the DownloadCompletedTask wait. That only unblocks the reader on a caller-token-only cancel; it does not stop the background downloads (the pipeline token is not fired), so it half-addresses that case at the cost of ~40 lines of TPL plumbing. Drop it and restore the plain await on this.currentDownloadResult.DownloadCompletedTask. Cancellation that matters is unaffected: statement Cancel()/Dispose() and connection dispose fire the pipeline token, the downloader completes the task with OperationCanceledException (CloudFetchDownloader IsCanceled branch), and the wait unblocks. A caller cancelling only its own read token unblocks at the next chunk boundary / on dispose. PBI cancels via statement.Cancel(), which fires the pipeline. Co-authored-by: Isaac
eric-wang-1990
removed
the
engineer-bot
…(JDBC/kernel semantics) Once results are streaming the query has already finished server-side, so tearing down the CloudFetch download is a close/dispose concern, not a query-cancel. Align with the JDBC and Rust-kernel drivers: statement.Cancel() now cancels execution only (server RPC via base.Cancel() / CancelStatementAsync + the per-execute token), and stopping in-flight downloads is done by dispose (reader/statement/connection). - statement.Cancel() no longer cancels the statement-lifetime CloudFetch CTS. - The CloudFetch CTS is created once (readonly), still linked to the connection's shutdown token, so connection dispose AND statement dispose still tear the pipeline down. Because Cancel() no longer poisons it, the per-execute RefreshCloudFetchStatementCts() and its lock are removed. - Net -150 lines across both protocols; no factory/caller changes. Behavioral change: Cancel() during the fetch phase no longer stops local downloads (nothing to cancel — the query is done); Dispose does. Execution-phase cancel and the connection-dispose teardown (the original CloseConnection_DuringCloudFetch hang) are unchanged. Tests: reframed CancelStatement_DuringCloudFetch -> DisposeStatement_DuringCloudFetch (dispose stops fetch); updated the unit test to assert Cancel() does not cancel the CloudFetch token. Validated live on a real warehouse, both Thrift and SEA: dispose-during-fetch stops promptly, CloseConnection no-hang still passes, execution cancel still stops, full-read 10/10 no regression, unit suite 983 green. Co-authored-by: Isaac
There was a problem hiding this comment.
Verdict: 1 Low
Solid, well-tested fix. The connection ⊃ statement ⊃ cloudfetch cancel cascade is correct: the downloader's new IsCanceled branch preserves the original success condition while faulting the in-flight DownloadCompletedTask on cancel; the reader's token.ThrowIfCancellationRequested() is genuinely needed because GetNextDownloadedFileAsync swallows OCE and returns null; straggler cancels are retried internally so they won't trip the cancel branch; and CTS cancel/dispose ordering plus double-dispose guards are handled on both protocols. One Low (telemetry: expected cancellation is now logged under an *_error event). Note: the PR description says the statement CloudFetch token is "refreshed per-execute," but the merged code (and its comments) create it once and deliberately do not refresh — the code is self-consistent, only the description is stale.
| 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. |
There was a problem hiding this comment.
🔵 Low — The new token.ThrowIfCancellationRequested() throws an OperationCanceledException on the expected teardown path (statement Dispose() / connection Dispose()). That exception propagates into the surrounding catch (Exception ex) block, which records an cloudfetch.get_next_file_error activity event and rethrows. So every normal dispose-mid-CloudFetch — the exact scenario this PR makes routine — now emits an error-named telemetry event for what is a benign, expected cancellation.
This is observability-only (no functional impact), but it will make dashboards/alerts that key off *_error events noisier and can mask genuine fetch errors. Consider special-casing OperationCanceledException here to emit a benign cloudfetch.reader_cancelled event (or skip the event entirely) before rethrowing, mirroring how the downloader treats a cancelled task distinctly from a faulted one.
Co-authored-by: Isaac
There was a problem hiding this comment.
Verdict: 1 Low
Solid, well-scoped fix — the connection/statement CloudFetch shutdown tokens are correctly linked into the download-manager pipeline CTS, the downloader now completes the in-flight DownloadCompletedTask with OperationCanceledException on teardown, and the reader's new token.ThrowIfCancellationRequested() after GetNextDownloadedFileAsync closes a real silent-truncation-on-cancel bug (a cancelled null previously looked like clean EOF). Token-getter disposal guards, dispose ordering (cancel → HttpClient/session teardown → CTS dispose in finally), and the _cancelLock usage all check out, and the added unit + E2E tests exercise the parked-reader and in-flight-download cancel paths. Only one minor consistency nit inline.
Nit (no anchor needed): in DatabricksConnection.Dispose, moving _cloudFetchHttpClient?.Dispose() inside the timed try means a (rare) HttpClient dispose throw is recorded as closeSessionError with a 0ms `closeSessio
[...truncated to keep verdict scannable]
| _httpClient, | ||
| _isLz4Compressed); | ||
| _isLz4Compressed, | ||
| ((DatabricksStatement)_statement).CloudFetchStatementToken); |
There was a problem hiding this comment.
🔵 Low — The new ((DatabricksStatement)_statement) is a hard cast, whereas every other access to _statement in this file (MaxBytes at ~L171, GetHeartbeatIntervalFromConnection, GetRequestTimeoutFromConnection) uses the defensive _statement is DatabricksStatement pattern. _statement is typed as IHiveServer2Statement, so if it is ever not a DatabricksStatement (a non-Databricks Thrift statement type, a test double), this throws InvalidCastException during reader creation rather than degrading gracefully.
In production the Thrift CloudFetch path always carries a DatabricksStatement, so this is not currently reachable — but the inconsistency with the surrounding defensive style is worth aligning. If the invariant is intended to be hard, consider Debug.Assert/a clearer cast-with-message; otherwise mirror the is guard used elsewhere.
## Summary Cuts the **v1.1.9** patch release of the C# Databricks ADBC driver. - Bumps `VersionPrefix` in `csharp/Directory.Build.props`: `1.1.8` → `1.1.9` - Adds the `[1.1.9]` section to `csharp/CHANGELOG.md` for the commits that landed on `main` since `v1.1.8` ## Release automation Merging this to `main` triggers `.github/workflows/csharp-sync-latest.yml`, which: - advances `release/csharp/v1.1.latest`, - tags the tip `csharp/v1.1.9`, and - force-updates `release/csharp/latest`. ## What's included (since v1.1.8) **Added** - Make the driver AOT/trim-safe under `net10.0` (adbc-drivers#509) - Apply server feature flags synchronously on connect (adbc-drivers#656) **Fixed** - Don't filter `GetCrossReference` by the seeded default catalog (adbc-drivers#660) - Align `GetCrossReference` Thrift and SEA paths on parent identifiers (adbc-drivers#645) - Default empty/null `TABLE_TYPE` to `TABLE` in `GetTables` (adbc-drivers#654) **Changed** - Update hiveserver2 submodule to downgrade ApacheThrift (adbc-drivers#649) The CloudFetch dispose/cancel work (adbc-drivers#659) is intentionally **not** included — it is not yet merged to `main`. This pull request and its description were written by Isaac. Co-authored-by: Isaac <no-reply@databricks.com>
1 participant
Problem
ConcurrencyStressTests.CloseConnection_DuringCloudFetch_ShouldNotHang(Thrift) flakes — closing a connection mid-CloudFetch hangs for 60s+. It bounced PR #654 out of the merge queue.Root cause
The CloudFetch pipeline runs on background tasks whose cancellation token is created inside
CloudFetchDownloadManager.StartAsync, linked to nothing external.connection.Dispose()disposes the sharedHttpClientand closes the session but never cancels that token, so the in-flight download fails on the deadHttpClientand the retry loop spins for the retry timeout (minutes) while the reader is parked onDownloadCompletedTask.Fix — disposing tears down the CloudFetch pipeline
DatabricksConnection(Thrift) andStatementExecutionConnection(SEA) each get aCloudFetchShutdownToken— aCancellationTokenSourcecancelled at the top ofDispose(), before theHttpClientis disposed.DownloadCompletedTaskwithOperationCanceledException, so a reader parked on it unblocks immediately instead of spinning on retries.Net: disposing the connection (or statement) during CloudFetch tears the pipeline down in milliseconds instead of hanging for minutes.
Scope (what this does not change)
statement.Cancel()is unchanged frommain— it cancels a still-executing query (server cancel RPC) but is a no-op once results are streaming, because by then the query has finished server-side. Stopping an in-flight CloudFetch download is a dispose/close concern, not a cancel — which matches the JDBC (close()) and Rust-kernel (drop) drivers. So this PR fixes the dispose hang only; it does not add cancel-stops-download behavior.The remaining gap — making
Cancel()also stop an in-flight CloudFetch download (so a consumer that cancels mid-fetch, e.g. Power BI Desktop, isn't left waiting on the server timeout) — is tracked separately in PECOBLR-4296. Verified live that cancel mid-CloudFetch is currently a no-op on both protocols.Testing
Unit (net8.0): full suite green — 983 passed.
E2E — live warehouse, both Thrift and SEA:
DisposeStatement_DuringCloudFetch_ShouldStopPromptly— disposing mid-stream stops the read (Thrift ~26s / SEA ~15s wall,OperationCanceledException).CloseConnection_DuringCloudFetch_ShouldNotHang— passes on both protocols (~14s / ~3s).CancelStatement_FromAnotherThread_ShouldStopPromptly— execution-phase cancel passes; separately verified live that cancelling a still-executing query showsCANCELEDin query history.CloudFetchE2ETest.TestCloudFetch(10 cases) passes on both protocols.This pull request and its description were written by Isaac.