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

[eric-wang-1990]

Problem

ConcurrencyStressTests.CloseConnection_DuringCloudFetch_ShouldNotHang (Thrift) flakes with:

Connection close during CloudFetch hung for more than 60000ms. Query completed: False, Dispose completed: True.

It bounced PR #654 out of the merge queue. The failure is unrelated to #654 — it's a pre-existing, intermittent hang in the CloudFetch shutdown path on main.

Root cause

The CloudFetch pipeline runs 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 — owned by the blocked query task, so it can't dispose until the read loop returns. connection.Dispose() disposes the shared HttpClient and closes the session but never cancels the pipeline, so the in-flight download fails on the dead HttpClient and the retry loop spins for the retry timeout (minutes) while the reader is parked on DownloadCompletedTask. Intermittent because it only fires when Dispose() lands mid-download.

Fix — full connection ⊃ statement ⊃ cloudfetch cancel cascade, both protocols

• Connection: DatabricksConnection (Thrift) and StatementExecutionConnection (SEA) each get a CloudFetchShutdownToken (a CancellationTokenSource cancelled at the top of Dispose()).
• Statement: DatabricksStatement and StatementExecutionStatement each get a statement-lifetime CTS linked to their connection's shutdown token, cancelled in Cancel() and Dispose(), and refreshed per-execute (a reusable statement mustn't be poisoned by a prior Cancel()). This is distinct from the per-execute token, which is disposed when execution returns and so can't cover the CloudFetch result-fetch phase.
• Pipeline: CloudFetchDownloadManager.StartAsync(CancellationToken) builds its source as a linked source; both reader factories pass the statement token in.
• Prompt stop even mid-buffer: CloudFetchDownloadManager exposes PipelineToken; CloudFetchReader links it with the caller token and checks it at the top of each read-loop iteration, so a cancel stops the read immediately instead of draining the in-memory buffer.

Net: closing a connection, or cancelling/disposing a statement, tears the pipeline down and the reader unblocks in milliseconds. A healthy read that never cancels is unaffected — the tokens only fire on explicit dispose/cancel (no timers). The only wall-clock bound on CloudFetch remains the pre-existing query timeout (ThriftResultFetcher / SEA, default ~3h).

Testing (real warehouse, both Thrift and SEA)

• New unit StartAsync_TokenCancelled_UnblocksReaderWaitingForNextFile: times out (5s) without the fix, ~0.2s with it.
• New E2E CancelStatement_DuringCloudFetch_ShouldStopPromptly: without the statement wiring it never stops (33s timeout on Thrift; read 67M rows on SEA); with it, ~5-20s. Verified as a real guard on both protocols.
• CloseConnection_DuringCloudFetch_ShouldNotHang: passes repeatedly on Thrift (~2-6s) and SEA (~1s).
• Full-read CloudFetchE2ETest.TestCloudFetch (10 cases) passes on both protocols — no regression to normal reads.
• Unit suite: 926 green.

This pull request and its description were written by Isaac.

[peco-review-bot]

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.

[peco-review-bot]

✅ No issues identified by the review bot.

[peco-review-bot]

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.

[peco-review-bot]

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).

[peco-review-bot]

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).

[peco-review-bot]

✅ No issues identified by the review bot.

[peco-review-bot]

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.

[peco-review-bot]

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).

[peco-review-bot]

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.

[peco-review-bot]

✅ No issues identified by the review bot.

[peco-review-bot]

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.

[peco-review-bot]

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).

[peco-review-bot]

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.

[peco-review-bot]

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.

[peco-review-bot]

Verdict: 1 Low

Looks good — a well-structured, heavily-documented fix for the CloudFetch shutdown hang, with a correct connection ⊃ statement ⊃ cloudfetch cancel cascade mirrored on both Thrift and SEA. I verified token-field locking is consistent per class, dispose ordering frees linked-token registrations unconditionally, the reader's cancelled-vs-EOF null distinction is sound, and the ICloudFetchDownloadManager change is backward-compatible (no hand-written test doubles). Only one low-severity note about a newly-introduced concrete cast whose safety rests on the current single caller.

[peco-review-bot]

Verdict: 1 Low

Looks good — a well-engineered, carefully-documented fix for the CloudFetch dispose/cancel hang, with one low-severity robustness note. The connection⊃statement⊃cloudfetch cancel cascade, cancel-before-teardown ordering, ObjectDisposedException-tolerant token getters, per-execute CTS refresh, and lock-guarded field swaps all check out on both Thrift and SEA paths; the StartAsync(CancellationToken = default) signature change is backward compatible and the new unit/E2E tests target the real regression.

[peco-review-bot]

Verdict: 1 Low

Looks good — a well-structured connection ⊃ statement ⊃ cloudfetch cancel cascade with careful locking around the per-statement CTS refresh/swap and correct linked-token teardown; AwaitWithCancellationAsync handles the completed/faulted/cancelled cases correctly and the manager disposes its linked CTS. One low-severity inconsistency: the connection dispose paths catch only generic Exception around the shutdown-CTS Cancel(), so a double-dispose emits a spurious error telemetry event, whereas the statement paths deliberately swallow ObjectDisposedException. Nit: CreateThriftReader casts (DatabricksStatement)statement unconditionally while CreateStatementExecutionReader guards the cast with a clear-message throw — harmless but asymmetric.

[peco-review-bot]

Verdict: 1 Low

Looks good — a solid, thoroughly-documented fix for the CloudFetch dispose/cancel hang. The connection⊃statement⊃cloudfetch token cascade is wired correctly on both protocols (both StartAsync callers pass the statement token; RefreshCloudFetchStatementCts runs on every execute path; token getters defensively map ObjectDisposedException→None), the reader's cancelled-null vs EOF distinction prevents silent truncation, and AwaitWithCancellationAsync correctly avoids leaking the delay timer. One low-severity note: a cross-thread Cancel() racing Dispose() emits a spurious error telemetry event for a benign ObjectDisposedException, since Cancel() (unlike Dispose()) doesn't special-case it.

[peco-review-bot]

Verdict: 1 Low

Solid, carefully-documented fix for the CloudFetch teardown hang — the connection ⊃ statement ⊃ cloudfetch cancel cascade, linked-CTS refresh-per-execute, and the AwaitWithCancellationAsync netstandard shim all look correct, and the lock scoping around the per-statement CTS swap is sound. One low-severity edge case: the cancel path now abandons the in-flight DownloadCompletedTask, which can leave a later download fault unobserved.

[peco-review-bot]

✅ No issues identified by the review bot.