Bound Cosmos diagnostics under retry storms

[NaluTripician]

🔗 Stacked on #4619 (nalutripician/diagnostics-capture-redesign) — this change depends on #4619's DiagnosticsSummary and its re-homed diagnostics::capture model. Please review/merge this AFTER #4619 lands. Because #4619 is not in main yet, the diff below currently includes #4619's commits (~122 files); it will shrink to just these 7 files once #4619 merges into main:

• sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/capture/model.rs
• sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/capture/mod.rs
• sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/mod.rs
• sdk/cosmos/azure_data_cosmos_driver/src/options/diagnostics_options.rs
• sdk/cosmos/azure_data_cosmos_driver/tests/live_storm_diagnostics.rs
• sdk/cosmos/azure_data_cosmos_driver/Cargo.toml
• sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md

Make Cosmos diagnostics storm-safe: bounded-size compaction + live fault-injection validation. Follow-up to #4619.

Why

With the cheap Threshold gate, a latency/failure spike trips the gate for a large fraction of requests at once, so the expensive diagnostics materialization becomes steady-state exactly when the system is already stressed. And a single operation that retries a hot partition grows its per-attempt list — and its detailed JSON — without bound. This implements the diagnostics contract's bounded-size guarantee and empirically validates the concern.

Deliverable 1 — Retry-storm compaction (always-on model)

• At operation finalization (DiagnosticsContextBuilder::complete), collapse runs of consecutive near-identical retries (same region / endpoint / status / sub-status / execution-context) into first + last + a count.
• Bounded by a new configurable DiagnosticsOptions::max_request_diagnostics (with_max_request_diagnostics, env AZURE_COSMOS_DIAGNOSTICS_MAX_REQUESTS, default 512, min 16).
• An order-robust global key-bucket fallback guarantees the retained count stays within the cap even under a region ping-pong (A→B→A→B), where consecutive run-length collapse alone would not.
• Lossless for the signal that matters: the summary (added by [Prototype] Diagnostics capture engine (Cosmos driver) #4619) is computed from the full attempt list before compaction, so the (status, sub-status) histogram, retry/throttle counts, total RU and regions stay exact. request_count() still reports the true total; new retained_request_count() and compaction() (returning CompactionInfo / CompactedRun) expose the bounded count and per-run rollup.
• Byte-identical when not triggered: detailed JSON gains a top-level compaction object only under a storm (skip_serializing_if). Regions stay normalized-lowercase. Additive and non-breaking.

Deliverable 2 — Threshold-storm validation

• Deterministic in-crate measurement (storm_materialization_cost_and_size, #[ignore]d): a 100k-retry storm produces ~48 MB / 3.6 s of detailed JSON uncompacted vs 1.8 KB / 421 µs compacted (~27,000× smaller), while the summary still reports the exact retry count.
• Env + feature-gated live test (tests/live_storm_diagnostics.rs, required-features = ["fault_injection"], #![cfg(feature = "reqwest")]) reads COSMOSDB_MULTI_REGION and injects latency + 429/503/410 storms, measuring gate-fire fraction and per-op materialization cost. Skips gracefully when the account is absent/unreachable; never prints secrets — so CI/playback never run it.
• Mitigations recommended back to the contract: structural compaction (this PR), configurable diagnostics depth under stress, sampling + max-events-per-interval cap, lazy/handle materialization.

Guardrails

• capture_engine untouched / still off by default; no Off mode; CosmosResponse::diagnostics() stays non-optional.
• Live tests are env + feature gated; no secrets committed.
• CHANGELOG updated for the additive public API (D1); D2 is internal tests only.

Testing

• cargo fmt clean; cargo clippy --all-features --all-targets -- -D warnings clean.
• cargo test -p azure_data_cosmos_driver --all-features: 1999 lib + all integration tests pass (incl. new compaction unit tests, options tests, and the live storm test against the real account).

[Copilot]

Pull request overview

This PR makes Cosmos driver diagnostics storm-safe by adding bounded-size compaction to the always-on DiagnosticsContext and a live fault-injection validation test. Under a 4xx/5xx retry storm a single operation could otherwise grow an unbounded per-attempt list (and unbounded detailed JSON). Compaction collapses runs of near-identical retries into first + last + a count at operation finalization, bounded by a new configurable DiagnosticsOptions::max_request_diagnostics (default 512, min 16, env AZURE_COSMOS_DIAGNOSTICS_MAX_REQUESTS), while the pre-compaction summary, request_count(), and total RU stay exact. Note: this PR is stacked on the unmerged #4619, so the diff currently also contains #4619's diagnostics-capture redesign (the diagnostics::capture module, gate wiring, encoding option, etc.); the standalone delta is the compaction and the live test.

Changes:

• Add retry-storm compaction (in model.rs, not shown in this diff) plus new public accessors retained_request_count() / compaction() and CompactionInfo/CompactedRun, re-exported from diagnostics and diagnostics::capture.
• Add DiagnosticsOptions::max_request_diagnostics with builder/env parsing, validation, and unit tests.
• Add an env + feature-gated live storm validation test and its Cargo target; document the additive API in the CHANGELOG.

Reviewed changes

Copilot reviewed 24 out of 25 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
src/options/diagnostics_options.rs	Adds max_request_diagnostics (default 512/min 16) + env parsing/tests; also #4619's DiagnosticsEncoding
tests/live_storm_diagnostics.rs	New env+feature-gated live storm test measuring gate-fire and compaction bounds
src/diagnostics/mod.rs	Re-exports CompactedRun/CompactionInfo (plus #4619 capture module re-homing)
src/diagnostics/capture/mod.rs	Re-exports compaction types from model; #4619 capture module root
CHANGELOG.md	Adds the compaction feature entry under "Features Added"
Cargo.toml	Registers the live_storm_diagnostics test target + #4619 bench/criterion dev-dep
sdk/cosmos/.cspell.json	Adds spelling dictionary terms (aggregatable, varint(s), Vecs, verbosities)
src/models/cosmos_response.rs, src/error/mod.rs, src/driver/cosmos_driver.rs, src/lib.rs, src/options/{mod,driver_options}.rs, src/diagnostics/capture/{recorder,pool,gate,event,encode,context}.rs, benches/diagnostics_capture.rs, tests/{live_capture_diagnostics,diagnostics_examples}.rs, DIAGNOSTICS-CAPTURE.md, Cargo.lock	Carried in from the stacked #4619 (capture engine, gate, encoding, docs); reviewed under #4619

[NaluTripician]

Reviewed this together with the contract doc (#4684) and the OTel spike (#4685). Overall the compaction design is solid — output stays byte-identical when compaction doesn't fire, the aggregate summary is computed before the lossy step so the histogram/counts stay exact, and the global-bucket fallback is deterministic (first-seen order, no HashMap iteration nondeterminism). A few things I'd want addressed before this merges:

Blocking — the bounded-size guarantee is only partial

global_bucket_compact truncates retained to cap, but never bounds runs:

if retained.len() > cap { retained.truncate(cap); }   // runs: Vec<CompactedRun> is left at one-per-distinct-key

CompactionInfo.runs is serialized into every detailed/summary output, and its length is the number of distinct CompactionKeys (region / endpoint / status / sub-status / execution-context). A 410/Gone fan-out across many physical-partition endpoints is high-cardinality exactly when the storm is worst, so the detailed JSON grows O(distinct keys) — bounded by topology, not by max_request_diagnostics. The two "many" tests don't catch this: region_ping_pong… uses 2 keys, and retry_storm_is_bounded_and_lossless asserts the 16 KB bound with a single key.

Suggest bounding runs too (top-N runs + an "other" rollup) and adding a distinct-key JSON-size test (e.g. 5000 unique endpoints → assert json.len() stays bounded).

Silent truncation contradicts the contract

The contract (#4684 §6/Q4) says "Truncation is marked, never silent" and "first+last-per-region always retained." The retained.truncate(cap) above silently drops later buckets with no marker in the retained list. Either surface a marker or amend the doc.

Eager summary on the completion path

When original_count > cap, complete() now computes the summary eagerly for every stormed op even when diagnostics are never read — this regresses the lazy design the (removed) comment described, and it fires precisely when the system is already under load. The #[ignore]d bench always reads the JSON, so it never measures the never-inspected case.

Peak memory mid-storm is still unbounded

Compaction runs only at complete(); self.requests grows one entry per attempt while the operation is in-flight. Worth stating explicitly that the guarantee is scoped to the final serialized artifact, not live memory.

request_count() semantic change

It now returns the pre-compaction total, so request_count() != requests().len() after a storm. Please audit azure_data_cosmos callers that size or iterate over requests() using request_count(). (The #4685 spike encodes exactly this pattern — see the comment there.)

Live test hardcodes the write region

error_rule(…, Region::new("West US 2"), …) while the account is supplied via COSMOSDB_MULTI_REGION. If the write region differs (the standard multi-region test account uses East US 2 / West US 3), the error rules never match, no storm is induced, compacted_ops stays 0, and the test passes green without exercising the compaction it exists to validate. Derive the region from a first probe's regions_contacted() (or document the account assumption).

Nit: the live test's std::env::set_var("AZURE_COSMOS_DIAGNOSTICS_MAX_REQUESTS", …) is process-global (and set_var is unsound under concurrency in newer Rust) — prefer setting the cap through the builder, or serialize the test.

[NaluTripician]

Addressed in 085f87f (plus a merge of upstream/main that resolves the conflict — the PR is now MERGEABLE):

Blocking — bounded-size guarantee was only partial. CompactionInfo.runs was O(distinct keys), so a high-cardinality 410 fan-out grew the detailed JSON by topology. runs is now bounded by bound_runs() (keep the largest runs by attempt count, deterministic tie-break by first-seen order, up to the cap) with the remainder rolled into explicit omitted_runs / omitted_request_count. Added distinct_key_fanout_runs_and_json_are_bounded (5000 unique endpoints) asserting the detailed JSON stays cap-bounded and lossless.

Silent truncation. CompactionInfo now carries total_runs and retained_truncated; the global-bucket retained.truncate(cap) sets the flag, so both retained-list truncation and run omission are explicit, never silent. Reconciled with the contract doc (#4684 §6/Q4).

Eager summary. Documented in complete() that the storm path must aggregate eagerly: compaction drops the source records, so the exact summary can't be recovered lazily, and retaining the full list to stay lazy would defeat the bounded-size guarantee. The fast-success (≤ cap) path stays fully lazy. Extended the deterministic bench to also measure the never-inspected complete() cost (ON eager-summary vs OFF lazy), addressing the "bench never measures the never-inspected case" note.

Peak mid-storm memory. Documented (code + CHANGELOG) that the guarantee is scoped to the finalized serialized artifact + retained list, not live mid-operation memory.

request_count() semantic change. Audited the azure_data_cosmos callers: they all use lower-bound checks (> 1, >= 2, >= 1) or iterate requests() directly, so they stay correct under compaction (first+last are always retained). Fixed the one cosmetic mismatch — assert_region_not_contacted now reports requests.len() (retained) as the denominator instead of request_count().

Live test hardcoded write region. Now derived from the baseline probe's regions_contacted() (skips gracefully if none), so a region mismatch can't silently induce no storm and pass green.

set_var nit. The cap is env-only at the driver today (DiagnosticsOptions::default() resolves AZURE_COSMOS_DIAGNOSTICS_MAX_REQUESTS per operation; there's no per-driver builder override yet), so the env is the only channel. Documented that this binary has exactly one #[tokio::test], so the write is serialized — and flagged a per-driver with_max_request_diagnostics knob as the preferred follow-up.

Merge conflict. Merged upstream/main; the only real conflicts were Cargo.lock (regenerated) and a driver_options.rs test block (kept both sides). Reconciled diagnostics_options.rs onto upstream's from_env + resolve_from_env pattern for max_request_diagnostics.

cargo fmt + cargo clippy --features fault_injection --all-targets clean; 1983 driver lib tests pass (incl. the new + existing compaction tests); the azure_data_cosmos test build is clean.

[NaluTripician]

Addressed the two round-2 review nits in fa3ff2804:

1. Phase-2 record/rollup coherence (the substantive one). Under the global key-bucket fallback, the retained per-attempt records and the per-run rollup were selected by different criteria — retained kept the earliest first-seen buckets, while the rollup kept the highest-count buckets. Under heterogeneous run counts those sets diverge, so a retained attempt could belong to a run omitted from the rollup, breaking the §7.3 span-emitter mapping. Fixed by folding the separate bound_runs step into global_bucket_compact and driving both the retained records and the rollup from one ranking (the cap-largest runs by attempt count), so every retained attempt now has a matching run in the rollup. total_runs / omitted_runs / omitted_request_count now flow through CompactionResult. Added phase2_heterogeneous_runs_keep_largest_and_stay_coherent (hot keys introduced last in first-seen order, interleaved with many cold single-attempt keys) — it fails on the old selection and passes now, asserting every retained record has a matching run + exact totals.

2. Live-test silent no-op guard. Replaced the weak max_request_count >= 1 assertion with a loud guard: once the storm batch reaches the service (and we only get there after deriving target_region from a region the baseline probe actually contacted), it must induce retries beyond the fault-free baseline. If the region-scoped faults silently fail to match (e.g. a region-name normalization mismatch), the test now fails with a clear message instead of passing green without exercising compaction.

cargo fmt + cargo clippy --all-features --all-targets -- -D warnings clean; all 61 model compaction unit tests pass.

[NaluTripician]

Superseded by #4789. The retry-storm bounding algorithm (max_request_diagnostics cap, run-collapse with exact aggregates, explicit truncation markers, order-robust key-bucket fallback) has been re-implemented over the existing upstream DiagnosticsContext/RequestDiagnostics (no dependency on the prototype capture engine) and is included in #4789 (WS6). Closing this first-pass prototype.