Cosmos FFI Follow-Up: Cancellation & diagnostics

[kundadebdatta]

Summary

Establish the completion diagnostics surface: either make cancellation propagate partial diagnostics or remove handle-based cancellation, and expose success-path diagnostics via a cosmos_diagnostics_t.

Motivation

@FabianMeiswinkel (PR #4515 review, marked blocking, later approved as a follow-up):

"I don't think we should offer any cancellation mechanism preventing getting a (partial) diagnostics… Any cancellation / e2e timeout we allow should propagate partial diagnostics. rather not provide cancellation in first step than one without access to partial diagnostics."

The completion's diagnostics field is currently reserved and always NULL, so even successful operations expose no diagnostics.

Current state (verified)

• src/submit.rs races the driver future against cancel_notify in a tokio::select! and drops the future on cancel, discarding any partial DiagnosticsContext.
• The driver op entry points (execute_singleton_operation, execute_plan, plan_operation) take no cancellation token; diagnostics only escape via DiagnosticsContextBuilder::complete() on the Ok/Err return path.
• CosmosCompletion.diagnostics: *mut c_void (src/completion.rs) is "Reserved for a future diagnostics handle; always NULL for now."
• The handle-based cancellation surface: cosmos_operation_handle_cancel, the Cancelled outcome/state, was_cancel_requested, OperationInner's cancel_requested/cancel_notify, CosmosErrorCodeOperationCancelled = 4012.

Proposed approach

Pick a cancellation direction, then design the diagnostics surface once:

• Option A (wrapper-only, smaller): remove handle-based cancellation now (Fabian's endorsed fallback). Drop cosmos_operation_handle_cancel, the Cancelled outcome/state, was_cancel_requested, and the tokio::select! cancel arm (replace with let out = work.await;). Keep the end-to-end timeout (it resolves via Err, so it already carries diagnostics). Keep the 4012 slot reserved for ABI stability.
• Option B (cross-crate, larger): add a cancellation token to the driver op entry points + a contract that returns the partial DiagnosticsContext on cancel; the wrapper passes the token instead of dropping the future. File the driver change as its own azure_data_cosmos_driver issue.

Then, independent of A/B, expose a cosmos_diagnostics_t (regions contacted, retry timeline, per-attempt latency, RU) on completions — populated on success (and on cancel/timeout if Option B).

Scope

• In: cancellation surface decision + implementation; cosmos_diagnostics_t shape, population, ownership/free contract; header regeneration; tests; spec/docs.
• Out: the driver-side token/diagnostics contract if Option B is chosen — split into a separate azure_data_cosmos_driver issue.

Work breakdown

• Decide Option A vs Option B (record in DATA_MOVEMENT_MODEL.md).
• If A: remove the cancellation surface (see FFI-08 removal list); keep e2e timeout + reserved 4012.
• If B: file + implement the driver cancellation-with-diagnostics contract; wire the wrapper to it.
• Design cosmos_diagnostics_t + its free contract; populate on success.
• Regenerate header; update C tests, cancellation.c, and docs.

Acceptance criteria

• No cancellation API that silently discards diagnostics — either cancellation is gone or it always delivers partial diagnostics.
• Diagnostics reachable from a completion with a documented ownership/free contract.
• Full validation sweep + cspell clean; header regenerated.

References

FFI-08, FFI-13 in FFI_FOLLOWUPS.md (FFI-08 lists the full Option-A removal surface); src/submit.rs, src/completion.rs.

[NaluTripician]

Decision proposal: take Option A now (remove handle-based cancellation); make cancellation-with-diagnostics a driver-layer follow-up

Following the blocking rule from #4515 (thanks @FabianMeiswinkel): no cancellation API may silently discard diagnostics. After tracing the wrapper, the driver, and the in-flight diagnostics work (#4683, #4684, deferred engine #4619), I'm proposing we commit to Option A now and move the real cancellation-with-diagnostics story down into the driver as a gated follow-up. Reasoning below — pushback welcome.

The rule is a correctness bar, and removal clears it outright

The requirement is "don't lose diagnostics on cancel." Two ways to satisfy it: remove the losing path (A), or build a mechanism that preserves diagnostics through cancel (B). Since the bar is not losing data, A satisfies it completely — after A, every terminal state (Ok / Err / driver timeout / retry-budget exhaustion) runs DiagnosticsContextBuilder::complete(). B's extra value is a capability (proactive, host-initiated abort that still returns partial diagnostics), not a fix for the rule.

Verified current state

• submit.rs races the driver future against cancel_notify in tokio::select! and, on cancel, drops the future — complete() never runs, partial diagnostics are lost. This is the one path that breaks the rule.
• Driver op entry points (execute_singleton_operation, execute_plan, plan_operation) take no cancellation token.
• The wrapper has no timeout of its own; runaway ops are already bounded by driver-side request timeouts / retry budget, which resolve as Err carrying diagnostics.
• CosmosCompletion.diagnostics is void*, always NULL — success ops expose no diagnostics today either.

Option A — the change

Replace the select! cancel arm with let done = work.await; (keep the catch_unwind panic firewall); remove cosmos_operation_handle_cancel, OperationInner.{cancel_requested, cancel_notify}, PendingCompletion::cancelled, was_cancel_requested, and the Cancelled outcome/state. Regenerate the header; delete/repurpose cancellation.c; record in DATA_MOVEMENT_MODEL.md §3.6.

Why this is the right call now:

• Fixes the rule with a small, single-crate, easily-reviewed change — no coordination with the three in-flight diagnostics PRs.
• Deletes a class of races (cancel-won vs cancel-lost, the Notify permit ordering, the enqueue flag) and stops dropping mid-flight futures (today's drop can tear down in-flight transport / session / PK-range-cache state mid-write — A removes that hazard too).
• Fully reversible. Keep 4012 and the Cancelled numeric slots reserved, so reintroducing cancellation later is additive and needs no renumbering.
• Unblocks the work people actually want — the success-path cosmos_diagnostics_t — without entangling it in a cancellation rewrite.

Effects on the Rust SDK / driver (why this isn't just an FFI decision)

Three things I verified make this a driver-layer decision, not an FFI-local one:

1. The driver is becoming the shared execution core. On main, azure_data_cosmos depends only on azure_core. On the diagnostics branches it gains a path dependency on azure_data_cosmos_driver (v0.6.0). So the cancellation + diagnostics contract chosen here is the foundation under both the FFI wrapper and the public Rust SDK.
2. Rust cancellation is drop-based and equally lossy. There is no cancellation token / deadline in azure_core or typespec_client_core — idiomatic Rust cancel = drop the future (tokio::time::timeout, select!, runtime shutdown). In cosmos_driver.rs the DiagnosticsContextBuilder is attached only on the return-path match (Ok → with_capture_diagnostics, Err → …), so a dropped future never runs complete() and loses diagnostics. The bug this issue is about already exists for direct Rust SDK callers who wrap a call in a timeout — and Option A does not fix that side.
3. A token-based Option B is non-idiomatic for Rust and only half-solves it. With no ambient token, B would introduce a new cross-cutting primitive threaded execute_operation → pipeline → retry loop → transport, and a Rust caller who simply drops the future (the normal way) still never trips it.

Implication: the durable fix for cancellation-with-diagnostics belongs in the driver, via Drop-flush (RAII) — emit a partial complete() when the op's diagnostics guard is dropped, so both Rust drop-cancel and FFI drop-cancel preserve diagnostics with no token and no signature change. This aligns with #4683's already "Drop/panic-safe" recorder and must honor #4684's DiagnosticsLevel + bounded-size rules — so it's gated on #4683/#4684 landing. That's the proper shape of "Option B," and it's a stronger reason to take A at the FFI now: removing handle-cancel stops the wrapper from masking a driver-layer concern by dropping mid-flight futures.

Why the original Option B isn't truly needed now

1. The correctness bar is already met by A + the fact that the driver timeout / retry-budget path already returns Err with diagnostics. "Stop a runaway op and still see what happened" is already served; B only adds earlier, proactive, host-initiated abort.
2. No consumer depends on it. Completion diagnostics are NULL today and the native wrapper is pre-release, so nothing relies on handle-cancel returning a (partial) diagnostics object.
3. B has nowhere to deliver partial diagnostics yet. Its whole point is handing a partial DiagnosticsContext to the caller on cancel — but the FFI has no diagnostics surface to hand it into until cosmos_diagnostics_t exists (gated on Add Cosmos diagnostics contract design doc #4684 + Bound Cosmos diagnostics under retry storms #4683). Building it now is speculative.
4. It's an additive capability, not a one-way door — reserved 4012/Cancelled slots mean the driver-level fix slots back in cleanly later.

How the in-flight diagnostics work changes the calculus

The real fix lives in exactly the code being rewritten this cycle:

• Bound Cosmos diagnostics under retry storms #4683 (bound diagnostics under retry storms) rewrites the terminal-finalization path — the op-end gate, CosmosResponse::capture_diagnostics() attachment, error/mod.rs, and the requests() (bounded first+last-per-run) vs request_count() (true total) + max_request_diagnostics cap semantics. A cancel/drop path must hook into whatever finalization looks like after Bound Cosmos diagnostics under retry storms #4683, and Bound Cosmos diagnostics under retry storms #4683 is currently CONFLICTING, so parallel work would fight it directly.
• Add Cosmos diagnostics contract design doc #4684 (diagnostics contract) ratifies the FFI boundary (opaque Arc<DiagnosticsContext> handle + explicit materialize), the DiagnosticsLevel { Minimal, Standard, Full } gating, and the bounded-size guarantee. A cancel-partial context is subject to those same rules; designing it before Add Cosmos diagnostics contract design doc #4684 is settled risks contradicting it.
• Gate model + deferred engine ([Prototype] Diagnostics capture engine (Cosmos driver) #4619). With the gate configurable Off / Threshold / Always, a cancel is an op-end event that needs a gate classification (does a cancelled op always emit, like an error/slow, or is it threshold-gated?). And "partial at cancel" means the builder hasn't been complete()d — we must decide whether to run complete() (with Bound Cosmos diagnostics under retry storms #4683's compaction) at the drop point or snapshot raw. Both only make sense once the gate/finalization model has landed.

ABI sub-decision (recommended)

Since the wrapper is pre-release, remove the producing code and the was_cancel_requested field for a clean surface, but keep the numeric slots reserved — 4012, the Cancelled outcome (= 2) and handle state (= 3) — so Unknown and friends don't renumber and a future driver-level cancel can return additively. The diagnostics field stays regardless (it's the success-path handle).

Plan

1. Now: implement Option A (this issue).
2. Success-path cosmos_diagnostics_t: design against Add Cosmos diagnostics contract design doc #4684 §4 + Bound Cosmos diagnostics under retry storms #4683 bounds once those land (needed regardless of the cancel direction).
3. Follow-up (its own azure_data_cosmos_driver issue): cancellation-with-diagnostics via driver Drop-flush, after Bound Cosmos diagnostics under retry storms #4683 + Add Cosmos diagnostics contract design doc #4684 — and only if callers need earlier-than-timeout abort. This covers the Rust SDK and the FFI in one place.

Questions for reviewers

• Any objection to committing to Option A now, given the correctness bar is met and the driver timeout path already returns diagnostics?
• Is losing proactive in-flight abort acceptable for the first release, assuming a driver-side deadline reachable from FFI? (If callers need earlier-than-timeout abort, that's the trigger to prioritize the driver Drop-flush follow-up.)
• Agree that cancellation-with-diagnostics belongs in the driver (Drop-flush) rather than a token threaded from the FFI — so it fixes the Rust SDK's drop-cancel path too?
• OK to keep 4012 / Cancelled slots reserved-but-inert (clean removal, additive re-add path)?

[NaluTripician]

Aligning this with #4789 (client-side observability layer)

#4789 lands the SDK-side observability stack on top of the driver DiagnosticsContext, and it shifts the ground this issue stands on. Two things drive the plan:

1. The diagnostics foundation is now much richer and, crucially, bounded. Add Cosmos client-side observability layer #4789 adds retry-storm compaction (max_request_diagnostics, compact_requests, retained_request_count()/compaction()), operation-level predicates (is_failure, is_completed, is_threshold_violated), an operation_name seam, and exact total_request_charge(). It also formalizes a single once-per-op completion seam (ClientContext::complete_operation → DiagnosticsHandler chain) and a DIAGNOSTICS-CONTRACT.md.
2. Add Cosmos client-side observability layer #4789 does not add any cancellation to the driver op entry points. There is still no driver contract that returns a partial DiagnosticsContext on cancel.

So the approach I want to take keeps the FFI surface consistent with #4789 instead of inventing a parallel one:

Cancellation → Option A (remove handle-based cancellation now)

Since #4789 provides no driver cancel-with-diagnostics contract, the only way to honor "no cancellation that silently discards diagnostics" today is to not ship handle-based cancellation. Remove cosmos_operation_handle_cancel, the Cancelled outcome/state, was_cancel_requested, OperationInner's cancel_requested/cancel_notify, and the tokio::select! cancel arm (replace with let out = work.await;). Keep the e2e timeout (it resolves via Err, so it already flows diagnostics through the normal completion path) and keep 4012 reserved for ABI stability.

Option B (a driver cancellation token that returns the partial context) becomes a separate azure_data_cosmos_driver issue — and #4789 now gives it a clean design target: hand back a context with is_completed() == false, and run compaction on that partial context before handing it out so a cancelled retry storm stays bounded (compaction currently only runs at complete()).

Diagnostics → project the post-#4789 DiagnosticsContext, don't redefine it

cosmos_diagnostics_t should be the C-ABI projection of the exact same completed DiagnosticsContext that #4789's handler chain consumes — same seam, same data, same ownership story. Concretely, the struct has to respect compaction:

• Exact aggregates: request_count and total_request_charge (kept exact even under compaction).
• retained_request_count + a compaction flag/marker, so consumers know the per-attempt array is the retained/compacted set, not a raw 1:1 timeline.
• Per-attempt entries (region, endpoint, status incl. sub-status, latency, RU) for the retained attempts only.
• Contacted regions, plus the operation-level flags Add Cosmos client-side observability layer #4789 already computes: is_failure / is_completed / is_threshold_violated.
• Populated on success (and on Err/timeout, which already carry a completed context). With Option A there is no cancel path left to populate.
• Ownership/free contract documented against DIAGNOSTICS-CONTRACT.md rather than a parallel FFI-only doc.

Net

Take Option A + build cosmos_diagnostics_t on the #4789 foundation. This unblocks the FFI diagnostics surface now, keeps it consistent with the SDK handler chain's view of the same context, and leaves a well-specified Option-B driver issue (partial context + compaction on cancel) for later if we ever want cancellation back. I'll record the A-vs-B decision in DATA_MOVEMENT_MODEL.md.

[NaluTripician]

Refined plan (superseding the comment above): 3 phases

Folding the cancellation decision and the diagnostics surface into a sequenced 3-phase plan, anchored on #4683 (shipped bounds) and #4684 §4 (FFI diagnostics contract).

Phase 1 — Now (this issue): Option A + FFI-reachable e2e deadline

• Remove handle-based cancellation: cosmos_operation_handle_cancel, the Cancelled outcome/state, was_cancel_requested, OperationInner's cancel_requested/cancel_notify, and the tokio::select! cancel arm (→ let out = work.await;).
• Keep catch_unwind (panic firewall — orthogonal to cancel).
• Reserve 4012 / Cancelled for ABI stability.
• Regenerate the header, retire cancellation.c, record the decision in DATA_MOVEMENT_MODEL.md §3.6.
• Release criterion (guardrail): a driver-side end-to-end deadline that is reachable/configurable from the FFI and resolves via Err (so it carries diagnostics through the normal completion path). This is what makes removing handle-cancel safe — callers keep a bounded-runtime escape hatch instead of losing abort entirely.

Phase 2 — Next (independent of cancel direction): success-path cosmos_diagnostics_t

Designed against #4684 §4 (opaque handle + explicit-materialize FFI boundary; never force JSON or a single fixed format at the boundary) and #4683 bounds (per-attempt list capped by max_request_diagnostics, default 512 / min 16; first+last-per-run + exact aggregates; retained set).

• Opaque diagnostics handle, not an eagerly-populated flat struct.
• Cheap operation-level fields available eagerly: request_count, total_request_charge (exact under compaction), is_failure / is_completed / is_threshold_violated, contacted regions.
• Per-attempt detail (region, endpoint, status incl. sub-status, latency, RU) materialized on demand over the retained requests() set; FFI buffer sizing reads the configured cap.
• Documented ownership/free contract against DIAGNOSTICS-CONTRACT.md.
• Populated on success and on Err/timeout (both already carry a completed context).
• This is the surface a partial context would need anyway → it's the real unblock.

Phase 3 — Follow-up (its own azure_data_cosmos_driver issue): cancellation-with-diagnostics via Drop-flush

Only if callers need earlier-than-timeout abort. Reuses #4683's lock-free, Drop/panic-safe recorder: dropping the driver future on cancel flushes the partial captured log through the gate into a DiagnosticsContext. Covers the Rust SDK and the FFI together. Gated on #4683 + #4684. On cancel it must:

• force a gate build (so a Threshold/Off policy doesn't discard the partial),
• run compaction on the partial (keep a cancelled retry storm bounded),
• mark is_completed() == false.

Net: Option A now (with the deadline guardrail) → opaque-handle diagnostics surface → Drop-flush cancel only if callers actually need it.