diff --git a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md index 2090ee46bd2..5b9b40f99eb 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md @@ -9,6 +9,7 @@ - Added a `.NET CosmosDiagnostics`-style top-level **`summary`** block to `DiagnosticsContext` (`DiagnosticsSummary`, reachable via `DiagnosticsContext::summary()`). It aggregates over the collected requests — `(status, sub-status)` histogram, total request charge, request/retry/throttle counts, regions contacted, final status, top error, total elapsed, and the client **`user_agent`** (SDK + runtime identity, via `DiagnosticsSummary::user_agent()`) — and is **computed lazily on first access** (never on the hot path) and cached via `OnceLock`, so a context whose diagnostics are never inspected pays nothing for it. Serialized as the first section of the diagnostics output; the `user_agent` is omitted when the driver did not supply one, so existing output is unchanged. - Added a **`DiagnosticsEncoding`** client option on `DriverOptions` (`DriverOptionsBuilder::with_diagnostics_encoding` / `DriverOptions::diagnostics_encoding`) controlling how diagnostics render to a string via the new `DiagnosticsContext::encode(encoding)` and `CosmosResponse::diagnostics_string(encoding)`: `Json` (pretty, **default** — output unchanged), `Compact` (minified JSON), and `Encoded` (base64 of the compact JSON, decodable, for size-sensitive logging). Additive and non-breaking. - Added serializable **wall-clock timestamps** (RFC 3339 / ISO 8601) to the diagnostics output. `DiagnosticsContext` (and its `summary`) now carry an operation `start_time`/`end_time`, and each `RequestDiagnostics` attempt carries a `start_time` plus an optional `end_time`, captured with `azure_core::time::OffsetDateTime` alongside the existing `Instant`-based durations. These absolute timestamps let diagnostics be correlated against server-side and external logs. Additive and non-breaking: the fields exist only when a context is built, durations are unchanged, and an in-flight attempt's `end_time` is omitted (`skip_serializing_if`). +- Added **retry-storm compaction** to the always-on `DiagnosticsContext`, giving a bounded-size guarantee under 4xx/5xx retry storms. When an operation accumulates more per-attempt records than the new configurable `DiagnosticsOptions::max_request_diagnostics` cap (`DiagnosticsOptionsBuilder::with_max_request_diagnostics`, env `AZURE_COSMOS_DIAGNOSTICS_MAX_REQUESTS`, default `512`, min `16`), runs of near-identical consecutive retries (same region / endpoint / status / sub-status / execution-context) are collapsed at operation finalization into first + last + a count, with an order-robust global key-bucket fallback that keeps the retained count within the cap even under a region ping-pong. The aggregate `summary` is computed from the full attempt list *before* compaction, so the `(status, sub-status)` histogram, retry/throttle counts, total RU and regions stay exact; `DiagnosticsContext::request_count()` still reports the true total, and new accessors `DiagnosticsContext::retained_request_count()` and `DiagnosticsContext::compaction()` (returning the new `CompactionInfo` / `CompactedRun`) expose the bounded retained count and a per-run rollup. Detailed JSON gains a top-level `compaction` object only when compaction fires (`skip_serializing_if`), so output stays byte-identical for normal operations. Additive and non-breaking. - Restructured the client / runtime options layering on the driver. Two new nested option groups, a per-client overrides surface on `DriverOptionsBuilder`, and a single canonical `AZURE_COSMOS_PPCB_*` namespace for partition-failover environment variables. The driver now consumes partition-failover configuration once at construction (`CosmosDriver::new` no longer fabricates an `OperationOptionsView` outside any operation context) ([#4588](https://github.com/Azure/azure-sdk-for-rust/pull/4588)): - Added new nested `OperationOptions::throughput_control` group (`ThroughputControlOptions` / `…Builder` / `…View`, mirroring the `ThrottlingRetryOptions` pattern). Exposes three layered fields ([#4588](https://github.com/Azure/azure-sdk-for-rust/pull/4588)): - `group_name: Option` — replaces the old top-level `OperationOptions::throughput_control_group`. diff --git a/sdk/cosmos/azure_data_cosmos_driver/Cargo.toml b/sdk/cosmos/azure_data_cosmos_driver/Cargo.toml index 7f302972205..daf146e2cfb 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/Cargo.toml +++ b/sdk/cosmos/azure_data_cosmos_driver/Cargo.toml @@ -136,6 +136,11 @@ name = "multi_region_failover" path = "tests/multi_region_failover.rs" required-features = ["fault_injection"] +[[test]] +name = "live_storm_diagnostics" +path = "tests/live_storm_diagnostics.rs" +required-features = ["fault_injection"] + [[test]] name = "gateway_query_plan_comparison" path = "tests/gateway_query_plan_comparison.rs" diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/capture/mod.rs b/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/capture/mod.rs index 95cb7cbd6a2..97635f01dd9 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/capture/mod.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/capture/mod.rs @@ -107,6 +107,7 @@ pub use recorder::{AttemptRecord, DiagnosticsRecorder, HedgeOutcome}; // `crate::diagnostics` so the public boundary (`diagnostics::DiagnosticsContext`, consumed by the // `azure_data_cosmos` SDK) is unchanged. pub(crate) use model::DiagnosticsContextBuilder; +pub use model::{CompactedRun, CompactionInfo}; pub use model::{ DiagnosticsContext, DiagnosticsSummary, ExecutionContext, FailedTransportShardDiagnostics, PipelineType, RequestDiagnostics, RequestEvent, RequestEventType, RequestHandle, diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/capture/model.rs b/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/capture/model.rs index 9df6212cd22..5da0ada6334 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/capture/model.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/capture/model.rs @@ -1334,6 +1334,10 @@ struct DiagnosticsOutput<'a> { system_usage: Option, #[serde(skip_serializing_if = "Option::is_none")] machine_id: Option<&'a str>, + /// Present only when the per-attempt list was compacted under a retry storm; + /// absent (and thus byte-identical to prior output) otherwise. + #[serde(skip_serializing_if = "Option::is_none")] + compaction: Option<&'a CompactionInfo>, #[serde(flatten)] payload: DiagnosticsPayload<'a>, } @@ -1865,19 +1869,49 @@ impl DiagnosticsContextBuilder { let duration = self.started_at.elapsed(); let start_time = self.start_time; let end_time = OffsetDateTime::now_utc(); - // The aggregatable summary is computed lazily on first access (see `DiagnosticsContext:: - // summary`), not here — a fast-success caller that never reads the summary or serializes the - // context pays nothing for it. The reduction inputs (requests, status, timing, user agent) - // are all retained on the context. + + let cap = self.options.max_request_diagnostics(); + let original_count = self.requests.len(); + + // Common path (attempts <= cap): the summary stays lazy (computed on first access, see + // `DiagnosticsContext::summary`) and the per-attempt list is retained verbatim, so a + // fast-success caller pays nothing and the serialized output is byte-identical to before. + // + // Storm path (attempts > cap): compute the summary eagerly from the FULL list first — so the + // status/sub-status histogram, retry/throttle counts, total RU and regions stay exact — then + // run-length compact the retained list and record a `CompactionInfo` marker. Compaction runs + // here at finalization, never mid-operation, so outstanding `RequestHandle` indices are never + // invalidated. + let (requests, summary, compaction) = if original_count > cap { + let summary = DiagnosticsSummary::compute( + &self.requests, + self.status, + duration, + start_time, + end_time, + self.user_agent.as_deref().map(str::to_string), + ); + let compacted = compact_requests(self.requests, cap); + let info = CompactionInfo { + original_request_count: original_count, + retained_request_count: compacted.retained.len(), + collapsed_runs: compacted.collapsed_runs, + runs: compacted.runs, + }; + (compacted.retained, OnceLock::from(summary), Some(info)) + } else { + (self.requests, OnceLock::new(), None) + }; + DiagnosticsContext { activity_id: self.activity_id, duration, start_time, end_time, - requests: Arc::new(self.requests), + requests: Arc::new(requests), status: self.status, user_agent: self.user_agent, - summary: OnceLock::new(), + summary, options: self.options, cpu_monitor: self.cpu_monitor, machine_id: self.machine_id, @@ -1886,6 +1920,7 @@ impl DiagnosticsContextBuilder { #[cfg(not(feature = "fault_injection"))] fault_injection_enabled: false, hedge_diagnostics: self.hedge_diagnostics, + compaction, #[cfg(test)] test_system_usage: self.test_system_usage, cached_json_detailed: OnceLock::new(), @@ -1982,6 +2017,12 @@ pub struct DiagnosticsContext { /// `Disabled`, eligibility check failed, etc.). hedge_diagnostics: Option, + /// Set only when a retry storm exceeded the configured + /// `max_request_diagnostics` cap and the per-attempt list was compacted. + /// `None` for normal operations, in which case the serialized output is + /// byte-identical to the pre-compaction behavior. + compaction: Option, + /// Test-only override for system usage snapshot, bypassing the CPU monitor. #[cfg(test)] test_system_usage: Option, @@ -2064,6 +2105,7 @@ impl DiagnosticsContext { machine_id: last.machine_id.clone(), fault_injection_enabled: sources.iter().any(|c| c.fault_injection_enabled), hedge_diagnostics: None, + compaction: None, #[cfg(test)] test_system_usage: last.test_system_usage.clone(), cached_json_detailed: OnceLock::new(), @@ -2150,17 +2192,60 @@ impl DiagnosticsContext { } /// Returns the total request charge (RU) across all requests. + /// + /// Under a compacted retry storm this is sourced from the (eagerly computed) + /// summary so it reflects the true total across all attempts, not just the + /// retained ones. pub fn total_request_charge(&self) -> RequestCharge { - self.requests.iter().map(|r| r.request_charge).sum() + if self.compaction.is_some() { + self.summary().total_request_charge() + } else { + self.requests.iter().map(|r| r.request_charge).sum() + } } /// Returns the number of requests made during this operation. + /// + /// This is always the **true** total number of attempts, even when the + /// per-attempt list was compacted under a retry storm. Use + /// [`retained_request_count`](Self::retained_request_count) for the number + /// of records actually retained in [`requests`](Self::requests). pub fn request_count(&self) -> usize { + self.compaction + .as_ref() + .map(|c| c.original_request_count) + .unwrap_or(self.requests.len()) + } + + /// Returns the number of per-attempt records retained after compaction. + /// + /// Equal to [`request_count`](Self::request_count) for normal operations; + /// bounded by the configured `max_request_diagnostics` cap under a storm. + pub fn retained_request_count(&self) -> usize { self.requests.len() } + /// Returns compaction metadata when a retry storm exceeded the configured + /// `max_request_diagnostics` cap and the per-attempt list was compacted. + /// + /// `None` for normal operations, where the retained list is the full, + /// unmodified set of attempts. + pub fn compaction(&self) -> Option<&CompactionInfo> { + self.compaction.as_ref() + } + /// Returns all regions contacted during this operation. pub fn regions_contacted(&self) -> Vec { + if self.compaction.is_some() { + // Source from the eager summary so every region is reported even + // though only first/last of each run are retained. + return self + .summary() + .regions_contacted() + .iter() + .map(|name| Region::new(name.clone())) + .collect(); + } let mut regions: Vec = self .requests .iter() @@ -2262,10 +2347,11 @@ impl DiagnosticsContext { start_time: self.start_time, end_time: self.end_time, total_duration_ms, - total_request_charge: self.requests.iter().map(|r| r.request_charge).sum(), - request_count: self.requests.len(), + total_request_charge: self.total_request_charge(), + request_count: self.request_count(), system_usage, machine_id: self.machine_id.as_ref().map(|s| s.as_str()), + compaction: self.compaction.as_ref(), payload: DiagnosticsPayload::Requests { requests: &self.requests, }, @@ -2301,10 +2387,11 @@ impl DiagnosticsContext { start_time: self.start_time, end_time: self.end_time, total_duration_ms, - total_request_charge: self.requests.iter().map(|r| r.request_charge).sum(), - request_count: self.requests.len(), + total_request_charge: self.total_request_charge(), + request_count: self.request_count(), system_usage: self.resolve_system_usage(), machine_id: self.machine_id.as_ref().map(|s| s.as_str()), + compaction: self.compaction.as_ref(), payload: DiagnosticsPayload::Summary { regions: region_summaries, }, @@ -2324,7 +2411,7 @@ impl DiagnosticsContext { start_time: self.start_time, end_time: self.end_time, total_duration_ms, - request_count: self.requests.len(), + request_count: self.request_count(), truncated: true, message: "Output truncated to fit size limit. Use Detailed verbosity for full diagnostics.", @@ -2358,6 +2445,7 @@ impl Clone for DiagnosticsContext { machine_id: self.machine_id.clone(), fault_injection_enabled: self.fault_injection_enabled, hedge_diagnostics: self.hedge_diagnostics.clone(), + compaction: self.compaction.clone(), #[cfg(test)] test_system_usage: self.test_system_usage.clone(), // OnceLock does not implement Clone, so we propagate any cached @@ -2510,6 +2598,213 @@ fn percentile_sorted(values: &[u64], p: u8) -> u64 { values[index.min(values.len() - 1)] } +/// Metadata describing how an operation's per-attempt diagnostics were compacted +/// under a retry storm. +/// +/// Present on a [`DiagnosticsContext`] only when the number of attempts exceeded +/// the configured `max_request_diagnostics` cap. It records the true attempt +/// count, how many records were retained, and a per-run rollup so the storm's +/// shape (which region/endpoint/status repeated, and how many times) is +/// preserved even though the middle of each run was dropped. +#[non_exhaustive] +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +pub struct CompactionInfo { + /// True total number of attempts before compaction. + pub original_request_count: usize, + /// Number of per-attempt records retained after compaction. + pub retained_request_count: usize, + /// Number of runs whose middle records were dropped (run length > 2). + pub collapsed_runs: usize, + /// Per-run rollup, in operation order (or first-seen order under the + /// global-bucket fallback). + pub runs: Vec, +} + +/// A single collapsed run of near-identical retries. +/// +/// Groups attempts that share the same region, endpoint, status (including +/// sub-status) and execution context. The first and last attempt of each run +/// are retained in full in [`DiagnosticsContext::requests`]; this rollup carries +/// the count and duration/charge statistics for the run so the elided middle is +/// still accounted for. +#[non_exhaustive] +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +pub struct CompactedRun { + /// Normalized region name for the run, if the attempts carried one. + #[serde(skip_serializing_if = "Option::is_none")] + pub region: Option, + /// Endpoint the run targeted. + pub endpoint: String, + /// HTTP status and Cosmos sub-status shared by the run. + #[serde(flatten)] + pub status: CosmosStatus, + /// Execution context shared by the run. + pub execution_context: ExecutionContext, + /// Number of attempts in the run. + pub count: usize, + /// Total request charge (RU) across the run. + pub total_request_charge: RequestCharge, + /// Minimum per-attempt duration in the run, in milliseconds. + pub min_duration_ms: u64, + /// Maximum per-attempt duration in the run, in milliseconds. + pub max_duration_ms: u64, + /// Median (P50) per-attempt duration in the run, in milliseconds. + pub p50_duration_ms: u64, +} + +/// The key that defines a "near-identical" run: same region, endpoint, status +/// (incl. sub-status) and execution context. +#[derive(Clone, PartialEq, Eq, Hash)] +struct CompactionKey { + region: Option, + endpoint: String, + status: CosmosStatus, + execution_context: ExecutionContext, +} + +impl CompactionKey { + fn of(req: &RequestDiagnostics) -> Self { + Self { + region: req.region.clone(), + endpoint: req.endpoint.clone(), + status: req.status, + execution_context: req.execution_context, + } + } +} + +/// The retained per-attempt records plus the per-run rollup produced by +/// compaction. +struct CompactionResult { + retained: Vec, + runs: Vec, + collapsed_runs: usize, +} + +/// Builds a [`CompactedRun`] rollup from a run/bucket of attempts. +fn compacted_run(reqs: &[&RequestDiagnostics]) -> CompactedRun { + let count = reqs.len(); + let mut durations: Vec = reqs.iter().map(|r| r.duration_ms).collect(); + durations.sort_unstable(); + let total_request_charge: RequestCharge = reqs.iter().map(|r| r.request_charge).sum(); + let first = reqs[0]; + CompactedRun { + region: first.region.as_ref().map(|r| r.as_str().to_string()), + endpoint: first.endpoint.clone(), + status: first.status, + execution_context: first.execution_context, + count, + total_request_charge, + min_duration_ms: durations.first().copied().unwrap_or(0), + max_duration_ms: durations.last().copied().unwrap_or(0), + p50_duration_ms: percentile_sorted(&durations, 50), + } +} + +/// Pushes the first and (when the run has more than one attempt) last record of +/// a run into `retained`. +fn push_first_last(retained: &mut Vec, run: &[&RequestDiagnostics]) { + if let Some(first) = run.first() { + retained.push((*first).clone()); + } + if run.len() > 1 { + if let Some(last) = run.last() { + retained.push((*last).clone()); + } + } +} + +/// Bounds an operation's per-attempt diagnostics to at most `cap` records. +/// +/// Phase 1 collapses runs of **consecutive** near-identical retries in order +/// (the common same-region storm). If that alone does not fit within `cap` +/// (e.g. a region ping-pong `A->B->A->B` where every consecutive run is length +/// one), Phase 2 falls back to a global key-bucket collapse that groups all +/// attempts by key regardless of order, bounding the result to at most two +/// records per distinct key. A final defensive truncation guarantees +/// `retained.len() <= cap` even with more than `cap / 2` distinct keys. +fn compact_requests(requests: Vec, cap: usize) -> CompactionResult { + let run_length = run_length_compact(&requests); + if run_length.retained.len() <= cap { + return run_length; + } + global_bucket_compact(&requests, cap) +} + +/// Phase 1: order-preserving run-length collapse of consecutive equal-key runs. +fn run_length_compact(requests: &[RequestDiagnostics]) -> CompactionResult { + let mut retained = Vec::new(); + let mut runs = Vec::new(); + let mut collapsed_runs = 0usize; + + let mut i = 0; + while i < requests.len() { + let key = CompactionKey::of(&requests[i]); + let mut j = i + 1; + while j < requests.len() && CompactionKey::of(&requests[j]) == key { + j += 1; + } + let run: Vec<&RequestDiagnostics> = requests[i..j].iter().collect(); + runs.push(compacted_run(&run)); + push_first_last(&mut retained, &run); + if run.len() > 2 { + collapsed_runs += 1; + } + i = j; + } + + CompactionResult { + retained, + runs, + collapsed_runs, + } +} + +/// Phase 2: order-robust global key-bucket collapse. +/// +/// Groups every attempt by [`CompactionKey`] preserving first-seen order, keeps +/// first + last of each bucket, and defensively truncates to `cap` if there are +/// somehow more than `cap / 2` distinct keys. +fn global_bucket_compact(requests: &[RequestDiagnostics], cap: usize) -> CompactionResult { + let mut order: Vec = Vec::new(); + let mut groups: HashMap> = HashMap::new(); + for req in requests { + let key = CompactionKey::of(req); + match groups.get_mut(&key) { + Some(bucket) => bucket.push(req), + None => { + order.push(key.clone()); + groups.insert(key, vec![req]); + } + } + } + + let mut retained = Vec::new(); + let mut runs = Vec::new(); + let mut collapsed_runs = 0usize; + for key in &order { + let bucket = &groups[key]; + runs.push(compacted_run(bucket)); + push_first_last(&mut retained, bucket); + if bucket.len() > 2 { + collapsed_runs += 1; + } + } + + // Defensive: with more distinct keys than cap/2, first+last per bucket could + // still exceed cap. Truncating keeps the earliest buckets' records; the true + // totals remain exact on the summary and `CompactionInfo`. + if retained.len() > cap { + retained.truncate(cap); + } + + CompactionResult { + retained, + runs, + collapsed_runs, + } +} + #[cfg(test)] mod tests { use super::*; @@ -3928,4 +4223,287 @@ mod tests { let ctx = builder.complete(); assert_eq!(ctx.machine_id(), None); } + + // ---- Compaction (retry-storm bounding) ------------------------------------------------ + + fn options_with_cap(cap: usize) -> Arc { + Arc::new( + DiagnosticsOptions::builder() + .with_max_request_diagnostics(cap) + .build() + .expect("valid cap"), + ) + } + + /// Records `count` identical attempts (same region/endpoint/status/exec-ctx). + fn record_run( + builder: &mut DiagnosticsContextBuilder, + exec: ExecutionContext, + region: &str, + endpoint: &str, + status: StatusCode, + count: usize, + ) { + for _ in 0..count { + let h = + builder.start_test_request(exec, Some(Region::new(region.to_string())), endpoint); + builder.complete_request(h, status, None); + } + } + + #[test] + fn retry_storm_is_bounded_and_lossless() { + let cap = 16; + let mut b = DiagnosticsContextBuilder::new( + ActivityId::from_string("storm".to_string()), + options_with_cap(cap), + ); + record_run( + &mut b, + ExecutionContext::Retry, + "East US", + "https://east/", + StatusCode::TooManyRequests, + 1000, + ); + b.set_operation_status(StatusCode::TooManyRequests, None); + let ctx = b.complete(); + + // True total is preserved; retained records are bounded by the cap. + assert_eq!(ctx.request_count(), 1000); + assert!( + ctx.retained_request_count() <= cap, + "retained {} exceeds cap {cap}", + ctx.retained_request_count() + ); + let info = ctx.compaction().expect("a storm past the cap must compact"); + assert_eq!(info.original_request_count, 1000); + assert_eq!(info.retained_request_count, ctx.retained_request_count()); + assert_eq!(info.runs.len(), 1); + assert_eq!(info.runs[0].count, 1000); + + // Aggregate signal stays exact (summary computed from the full list before compaction). + let s = ctx.summary(); + assert_eq!(s.request_count(), 1000); + assert_eq!(s.retry_count(), 1000); + + // Detailed output is bounded and carries the additive marker. + let json = ctx.to_json_string(Some(DiagnosticsVerbosity::Detailed)); + assert!(json.contains("\"compaction\"")); + assert!( + json.len() < 16 * 1024, + "detailed json {} bytes is not bounded", + json.len() + ); + + // First and last of the run are retained in full. + let requests = ctx.requests(); + assert_eq!( + u16::from(requests.first().unwrap().status().status_code()), + 429 + ); + assert_eq!( + u16::from(requests.last().unwrap().status().status_code()), + 429 + ); + } + + #[test] + fn mixed_runs_preserve_order_and_boundaries() { + let cap = 16; + let mut b = DiagnosticsContextBuilder::new( + ActivityId::from_string("mixed".to_string()), + options_with_cap(cap), + ); + record_run( + &mut b, + ExecutionContext::Retry, + "East US", + "https://east/", + StatusCode::TooManyRequests, + 100, + ); + record_run( + &mut b, + ExecutionContext::Retry, + "East US", + "https://east/", + StatusCode::ServiceUnavailable, + 50, + ); + record_run( + &mut b, + ExecutionContext::Retry, + "West US", + "https://west/", + StatusCode::Ok, + 1, + ); + b.set_operation_status(StatusCode::Ok, None); + let ctx = b.complete(); + + assert_eq!(ctx.request_count(), 151); + assert!(ctx.retained_request_count() <= cap); + let info = ctx.compaction().expect("compacted"); + assert_eq!(info.runs.len(), 3); + assert_eq!(info.runs[0].count, 100); + assert_eq!(info.runs[1].count, 50); + assert_eq!(info.runs[2].count, 1); + + // Onset (429) and terminal (200) boundaries retained; order preserved. + let requests = ctx.requests(); + assert_eq!( + u16::from(requests.first().unwrap().status().status_code()), + 429 + ); + assert_eq!( + u16::from(requests.last().unwrap().status().status_code()), + 200 + ); + } + + #[test] + fn region_ping_pong_is_bounded_via_global_fallback() { + let cap = 16; + let mut b = DiagnosticsContextBuilder::new( + ActivityId::from_string("pingpong".to_string()), + options_with_cap(cap), + ); + // Alternating regions: every consecutive run is length 1, defeating run-length collapse + // and forcing the global key-bucket fallback. + for _ in 0..200 { + let he = b.start_test_request( + ExecutionContext::Retry, + Some(Region::new("East US")), + "https://east/", + ); + b.complete_request(he, StatusCode::ServiceUnavailable, None); + let hw = b.start_test_request( + ExecutionContext::Retry, + Some(Region::new("West US")), + "https://west/", + ); + b.complete_request(hw, StatusCode::ServiceUnavailable, None); + } + b.set_operation_status(StatusCode::ServiceUnavailable, None); + let ctx = b.complete(); + + assert_eq!(ctx.request_count(), 400); + assert!( + ctx.retained_request_count() <= cap, + "retained {} exceeds cap", + ctx.retained_request_count() + ); + let info = ctx.compaction().expect("compacted"); + // Two distinct keys -> two runs, each covering 200 attempts. + assert_eq!(info.runs.len(), 2); + assert!(info.runs.iter().all(|r| r.count == 200)); + + // Regions are sourced from the summary, so both are still reported. + let mut regions: Vec = ctx + .regions_contacted() + .iter() + .map(|r| r.as_str().to_string()) + .collect(); + regions.sort(); + assert_eq!(regions, ["eastus".to_string(), "westus".to_string()]); + } + + #[test] + fn under_cap_is_not_compacted_and_output_has_no_marker() { + // Default cap (512) is far above a normal operation's attempts. + let ctx = make_context_with(ActivityId::from_string("normal".to_string()), |b| { + for _ in 0..3 { + let h = b.start_test_request( + ExecutionContext::Retry, + Some(Region::new("East US")), + "https://east/", + ); + b.complete_request(h, StatusCode::TooManyRequests, None); + } + b.set_operation_status(StatusCode::Ok, None); + }); + assert!(ctx.compaction().is_none()); + assert_eq!(ctx.request_count(), 3); + assert_eq!(ctx.retained_request_count(), 3); + let json = ctx.to_json_string(Some(DiagnosticsVerbosity::Detailed)); + assert!( + !json.contains("\"compaction\""), + "no compaction marker expected under the cap" + ); + } + + /// Deterministic companion measurement for the storm-validation deliverable. + /// + /// Quantifies the incremental materialization cost (struct build vs. building + /// + serializing detailed JSON) and the size bound that compaction provides, + /// comparing compaction ON (bounded cap) against OFF (cap above the attempt + /// count). `#[ignore]`d so it never runs in CI/playback; run manually with: + /// `cargo test -p azure_data_cosmos_driver --all-features --lib -- --ignored + /// --nocapture storm_materialization_cost_and_size`. + #[test] + #[ignore = "manual measurement; run with --ignored --nocapture"] + fn storm_materialization_cost_and_size() { + use std::time::Instant; + const N: usize = 100_000; + + fn build_storm(cap: usize, n: usize) -> DiagnosticsContextBuilder { + let mut b = DiagnosticsContextBuilder::new( + ActivityId::from_string("bench".to_string()), + options_with_cap(cap), + ); + record_run( + &mut b, + ExecutionContext::Retry, + "East US", + "https://east/", + StatusCode::TooManyRequests, + n, + ); + b.set_operation_status(StatusCode::TooManyRequests, None); + b + } + + // Compaction ON (bounded cap). + let b_on = build_storm(512, N); + let t0 = Instant::now(); + let ctx_on = b_on.complete(); + let build_on = t0.elapsed(); + let t1 = Instant::now(); + let json_on = ctx_on + .to_json_string(Some(DiagnosticsVerbosity::Detailed)) + .len(); + let json_on_t = t1.elapsed(); + + // Compaction OFF (cap above N, so nothing is collapsed). + let b_off = build_storm(N * 2, N); + let t2 = Instant::now(); + let ctx_off = b_off.complete(); + let build_off = t2.elapsed(); + let t3 = Instant::now(); + let json_off = ctx_off + .to_json_string(Some(DiagnosticsVerbosity::Detailed)) + .len(); + let json_off_t = t3.elapsed(); + + eprintln!("--- storm materialization (N={N} identical 429 retries) ---"); + eprintln!( + "compaction ON : build={build_on:?} json_build={json_on_t:?} detailed_json_bytes={json_on} retained={}", + ctx_on.retained_request_count() + ); + eprintln!( + "compaction OFF: build={build_off:?} json_build={json_off_t:?} detailed_json_bytes={json_off} retained={}", + ctx_off.retained_request_count() + ); + eprintln!( + "size reduction: {:.1}x", + json_off as f64 / json_on.max(1) as f64 + ); + + assert!(ctx_on.retained_request_count() <= 512); + assert!( + json_on < json_off / 10, + "compacted detailed json ({json_on} B) should be far smaller than uncompacted ({json_off} B)" + ); + } } diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/mod.rs b/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/mod.rs index 856546252c5..f329321e6c6 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/mod.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/mod.rs @@ -32,6 +32,7 @@ mod proxy_configuration; pub mod capture; pub(crate) use capture::DiagnosticsContextBuilder; +pub use capture::{CompactedRun, CompactionInfo}; pub use capture::{ DiagnosticsContext, DiagnosticsSummary, ExecutionContext, FailedTransportShardDiagnostics, PipelineType, RequestDiagnostics, RequestEvent, RequestEventType, RequestHandle, diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/options/diagnostics_options.rs b/sdk/cosmos/azure_data_cosmos_driver/src/options/diagnostics_options.rs index 6af4cd79c45..dee916ad978 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/options/diagnostics_options.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/options/diagnostics_options.rs @@ -11,6 +11,13 @@ const DEFAULT_MAX_SUMMARY_SIZE_BYTES: usize = 8 * 1024; /// Minimum allowed size for diagnostic summary output (4 KB). const MIN_MAX_SUMMARY_SIZE_BYTES: usize = 4 * 1024; +/// Default cap on the number of per-attempt request records retained in a +/// single operation's diagnostics before run-length compaction applies. +const DEFAULT_MAX_REQUEST_DIAGNOSTICS: usize = 512; + +/// Minimum allowed value for `max_request_diagnostics`. +const MIN_MAX_REQUEST_DIAGNOSTICS: usize = 16; + /// Controls the verbosity level of diagnostic output. /// /// Diagnostics can be output in different levels of detail depending on @@ -140,6 +147,17 @@ pub struct DiagnosticsOptions { /// /// Used when `to_json_string` is called with `None` for verbosity. pub(crate) default_verbosity: DiagnosticsVerbosity, + + /// Maximum number of per-attempt request records retained in a single + /// operation's diagnostics. + /// + /// When an operation accumulates more attempts than this (for example a + /// 4xx/5xx retry storm), runs of near-identical retries (same + /// region/endpoint/status/sub-status/execution-context) are run-length + /// compacted into first + last + a count, so the retained size stays + /// bounded. Default is 512; minimum is 16. Values at or above the number of + /// attempts a normal operation makes leave output byte-identical. + pub(crate) max_request_diagnostics: usize, } impl Default for DiagnosticsOptions { @@ -165,6 +183,12 @@ impl DiagnosticsOptions { pub fn default_verbosity(&self) -> DiagnosticsVerbosity { self.default_verbosity } + + /// Returns the maximum number of per-attempt request records retained in a + /// single operation's diagnostics before run-length compaction applies. + pub fn max_request_diagnostics(&self) -> usize { + self.max_request_diagnostics + } } /// Builder for [`DiagnosticsOptions`]. @@ -178,6 +202,9 @@ impl DiagnosticsOptions { /// summary mode output (default: `8192`, min: `4096`) /// - `AZURE_COSMOS_DIAGNOSTICS_DEFAULT_VERBOSITY`: Default verbosity level. /// Valid values: `default`, `summary`, `detailed` (default: `detailed`) +/// - `AZURE_COSMOS_DIAGNOSTICS_MAX_REQUESTS`: Maximum per-attempt request +/// records retained per operation before run-length compaction applies +/// (default: `512`, min: `16`) /// /// # Example /// @@ -195,6 +222,7 @@ impl DiagnosticsOptions { pub struct DiagnosticsOptionsBuilder { max_summary_size_bytes: Option, default_verbosity: Option, + max_request_diagnostics: Option, } impl DiagnosticsOptionsBuilder { @@ -220,6 +248,15 @@ impl DiagnosticsOptionsBuilder { self } + /// Sets the maximum number of per-attempt request records retained in a + /// single operation's diagnostics before run-length compaction applies. + /// + /// Must be at least 16. Default: 512. + pub fn with_max_request_diagnostics(mut self, max: usize) -> Self { + self.max_request_diagnostics = Some(max); + self + } + /// Builds the `DiagnosticsOptions` with configured values. /// /// Unset values are populated from environment variables or use sensible defaults. @@ -254,9 +291,17 @@ impl DiagnosticsOptionsBuilder { }, }; + let max_request_diagnostics = parse_from_env( + self.max_request_diagnostics, + "AZURE_COSMOS_DIAGNOSTICS_MAX_REQUESTS", + DEFAULT_MAX_REQUEST_DIAGNOSTICS, + ValidationBounds::min(MIN_MAX_REQUEST_DIAGNOSTICS), + )?; + Ok(DiagnosticsOptions { max_summary_size_bytes, default_verbosity, + max_request_diagnostics, }) } } @@ -270,6 +315,28 @@ mod tests { let options = DiagnosticsOptions::default(); assert_eq!(options.max_summary_size_bytes, 8 * 1024); assert_eq!(options.default_verbosity, DiagnosticsVerbosity::Detailed); + assert_eq!(options.max_request_diagnostics, 512); + } + + #[test] + fn custom_max_request_diagnostics() { + let options = DiagnosticsOptionsBuilder::new() + .with_max_request_diagnostics(64) + .build() + .unwrap(); + assert_eq!(options.max_request_diagnostics, 64); + } + + #[test] + fn max_request_diagnostics_too_small() { + let result = DiagnosticsOptionsBuilder::new() + .with_max_request_diagnostics(4) + .build(); + assert!(result.is_err()); + assert!(result + .unwrap_err() + .to_string() + .contains("must be at least 16")); } #[test] diff --git a/sdk/cosmos/azure_data_cosmos_driver/tests/live_storm_diagnostics.rs b/sdk/cosmos/azure_data_cosmos_driver/tests/live_storm_diagnostics.rs new file mode 100644 index 00000000000..3f9749d73a0 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos_driver/tests/live_storm_diagnostics.rs @@ -0,0 +1,309 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Live threshold-storm validation for diagnostics materialization cost + compaction. +//! +//! Deliverable 2 of the "storm-safe diagnostics" work item. Uses the +//! `fault_injection` feature to inject a latency + error (429/503/410) storm +//! against a LIVE multi-region account, then measures: +//! * the fraction of operations for which the `Threshold` gate fires +//! (`capture_diagnostics()` is `Some`) under a storm vs. a baseline batch, +//! * the incremental cost of materializing detailed diagnostics JSON per +//! operation, and +//! * that compaction keeps the retained per-attempt count bounded below the +//! true attempt count. +//! +//! **Env + feature gated.** Compiled only with `--features reqwest` (file-level +//! `cfg`) and `required-features = ["fault_injection"]` (Cargo), so CI/playback +//! never build or run it without both features. Reads `COSMOSDB_MULTI_REGION` +//! (a Cosmos connection string; master key) and **skips gracefully** — passing +//! without asserting — when the var is absent, does not parse, or the account +//! is unreachable. Secret values are never printed. +//! +//! The reproducible CPU/size numbers live in the in-crate deterministic +//! measurement +//! `diagnostics::capture::model::tests::storm_materialization_cost_and_size` +//! (run with `--ignored --nocapture`); this live test corroborates them against +//! real network latency and topology. + +#![cfg(feature = "reqwest")] + +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use azure_data_cosmos_driver::diagnostics::capture::DiagnosticsPolicy; +use azure_data_cosmos_driver::fault_injection::{ + FaultInjectionConditionBuilder, FaultInjectionErrorType, FaultInjectionResultBuilder, + FaultInjectionRule, FaultInjectionRuleBuilder, +}; +use azure_data_cosmos_driver::models::{ + AccountReference, ConnectionString, CosmosOperation, DatabaseReference, +}; +use azure_data_cosmos_driver::options::{DriverOptions, OperationOptions, Region}; +use azure_data_cosmos_driver::{ + CosmosDriver, CosmosDriverRuntime, CosmosError, CosmosResponse, DiagnosticsContext, + DiagnosticsVerbosity, +}; +use url::Url; + +/// Maximum time to wait for a live call before treating the account as unreachable. +const LIVE_TIMEOUT: Duration = Duration::from_secs(20); +/// Number of probe operations per batch. +const BATCH: usize = 12; +/// The gate trips when an operation exceeds this latency; the injected latency is above it. +const GATE_THRESHOLD: Duration = Duration::from_millis(5); +/// Per-request injected latency, comfortably above `GATE_THRESHOLD`. +const STORM_LATENCY: Duration = Duration::from_millis(50); +/// Small retained cap so a storm's bounded per-attempt count is observable. +const STORM_CAP: &str = "16"; + +type OpResult = std::result::Result; + +fn account_from_env() -> Option { + // Q4: COSMOSDB_MULTI_REGION is a Cosmos connection string (AccountEndpoint=...;AccountKey=...;). + let conn_str = std::env::var("COSMOSDB_MULTI_REGION").ok()?; + let conn: ConnectionString = conn_str.parse().ok()?; + let endpoint = Url::parse(conn.account_endpoint()).ok()?; + Some(AccountReference::with_master_key( + endpoint, + conn.account_key().clone(), + )) +} + +/// A `DiagnosticsContext` is available on both the success and error paths. +fn diagnostics_of(result: &OpResult) -> Option> { + match result { + Ok(resp) => Some(resp.diagnostics()), + Err(err) => err.diagnostics(), + } +} + +fn gate_fired(result: &OpResult) -> bool { + match result { + Ok(resp) => resp.capture_diagnostics().is_some(), + Err(err) => err.capture_diagnostics().is_some(), + } +} + +fn latency_rule(id: &str, delay: Duration) -> Arc { + let result = FaultInjectionResultBuilder::new() + .with_delay(delay) + .with_probability(1.0) + .build(); + Arc::new(FaultInjectionRuleBuilder::new(id, result).build()) +} + +fn error_rule( + id: &str, + err: FaultInjectionErrorType, + region: Region, + hit_limit: u32, +) -> Arc { + let condition = FaultInjectionConditionBuilder::new() + .with_region(region) + .build(); + let result = FaultInjectionResultBuilder::new() + .with_error(err) + .with_probability(1.0) + .build(); + Arc::new( + FaultInjectionRuleBuilder::new(id, result) + .with_condition(condition) + .with_hit_limit(hit_limit) + .build(), + ) +} + +#[derive(Default)] +struct BatchStats { + ops: usize, + reached: usize, + gate_fired: usize, + max_request_count: usize, + max_retained: usize, + compacted_ops: usize, + total_json_materialization: Duration, + json_samples: usize, +} + +impl BatchStats { + fn avg_json_materialization(&self) -> Duration { + if self.json_samples == 0 { + Duration::ZERO + } else { + self.total_json_materialization / self.json_samples as u32 + } + } +} + +async fn run_batch(driver: &CosmosDriver, account: &AccountReference, label: &str) -> BatchStats { + let mut stats = BatchStats::default(); + for i in 0..BATCH { + // Probe a database that almost certainly does not exist: it still exercises the full + // pipeline (and any injected faults/retries) and yields a real diagnostics context. + let db = DatabaseReference::from_name(account.clone(), "diag-storm-probe-nonexistent-db"); + let operation = CosmosOperation::read_database(db); + let outcome = tokio::time::timeout( + LIVE_TIMEOUT, + driver.execute_singleton_operation(operation, OperationOptions::default()), + ) + .await; + stats.ops += 1; + let result = match outcome { + Err(_elapsed) => { + eprintln!("[{label}] op {i} timed out (account unreachable / firewall-blocked)"); + continue; + } + Ok(result) => result, + }; + stats.reached += 1; + if gate_fired(&result) { + stats.gate_fired += 1; + } + if let Some(diag) = diagnostics_of(&result) { + stats.max_request_count = stats.max_request_count.max(diag.request_count()); + stats.max_retained = stats.max_retained.max(diag.retained_request_count()); + if diag.compaction().is_some() { + stats.compacted_ops += 1; + } + // Measure the incremental cost of materializing the detailed JSON (first call computes + // and caches it). This is the steady-state overhead paid when the gate fires broadly. + let started = Instant::now(); + let _ = diag.to_json_string(Some(DiagnosticsVerbosity::Detailed)); + stats.total_json_materialization += started.elapsed(); + stats.json_samples += 1; + } + } + stats +} + +#[tokio::test] +async fn live_storm_diagnostics_or_env_gated() { + let Some(account) = account_from_env() else { + eprintln!("live_storm env-gated: COSMOSDB_MULTI_REGION not set or unparseable; skipping"); + return; + }; + + // Best-effort: a small retained cap makes a storm's bounded output observable. Picked up when + // the driver builds its DiagnosticsOptions from the environment. + std::env::set_var("AZURE_COSMOS_DIAGNOSTICS_MAX_REQUESTS", STORM_CAP); + + let runtime = match CosmosDriverRuntime::builder().build().await { + Ok(r) => r, + Err(e) => { + eprintln!( + "live_storm env-gated: could not build runtime: {}", + e.status() + ); + return; + } + }; + + // Baseline batch: no faults, Threshold gate. + let baseline_options = DriverOptions::builder(account.clone()) + .with_capture_diagnostics_policy(DiagnosticsPolicy::threshold(GATE_THRESHOLD)) + .build(); + let baseline_driver = match runtime.create_driver(baseline_options).await { + Ok(d) => d, + Err(e) => { + eprintln!( + "live_storm env-gated: could not create baseline driver: {}", + e.status() + ); + return; + } + }; + let baseline = run_batch(&baseline_driver, &account, "baseline").await; + if baseline.reached == 0 { + eprintln!("live_storm env-gated: account unreachable; skipping without assertions"); + return; + } + + // Storm batch: per-request latency (trips the gate for a large fraction) plus 429/503/410 + // faults scoped to the write region (to induce retries/failover). + let write_region = Region::new("West US 2"); + let rules = vec![ + latency_rule("storm-latency", STORM_LATENCY), + error_rule( + "storm-429", + FaultInjectionErrorType::TooManyRequests, + write_region.clone(), + 32, + ), + error_rule( + "storm-503", + FaultInjectionErrorType::ServiceUnavailable, + write_region.clone(), + 32, + ), + error_rule( + "storm-410", + FaultInjectionErrorType::PartitionIsGone, + write_region, + 32, + ), + ]; + let storm_builder = match DriverOptions::builder(account.clone()) + .with_capture_diagnostics_policy(DiagnosticsPolicy::threshold(GATE_THRESHOLD)) + .with_fault_injection_rules(rules) + { + Ok(b) => b, + Err(e) => { + eprintln!( + "live_storm env-gated: fault rule install failed: {}", + e.status() + ); + return; + } + }; + let storm_driver = match runtime.create_driver(storm_builder.build()).await { + Ok(d) => d, + Err(e) => { + eprintln!( + "live_storm env-gated: could not create storm driver: {}", + e.status() + ); + return; + } + }; + let storm = run_batch(&storm_driver, &account, "storm").await; + + eprintln!("=== live storm diagnostics ==="); + eprintln!( + "baseline: reached={}/{} gate_fired={} max_requests={} avg_json_materialization={:?}", + baseline.reached, + baseline.ops, + baseline.gate_fired, + baseline.max_request_count, + baseline.avg_json_materialization() + ); + eprintln!( + "storm : reached={}/{} gate_fired={} max_requests={} max_retained={} compacted_ops={} avg_json_materialization={:?}", + storm.reached, + storm.ops, + storm.gate_fired, + storm.max_request_count, + storm.max_retained, + storm.compacted_ops, + storm.avg_json_materialization() + ); + + // Soft invariants (only assert when the storm batch actually reached the service, so the test + // stays green when the account is unreachable). + if storm.reached > 0 { + assert!( + storm.max_request_count >= 1, + "reached-service operations must carry diagnostics with at least one attempt" + ); + // Whenever compaction fired, the retained per-attempt count is strictly below the true + // total — the bounded-size guarantee holding under a live storm. + if storm.compacted_ops > 0 { + assert!( + storm.max_retained < storm.max_request_count, + "compaction must retain fewer records ({}) than the true attempt count ({})", + storm.max_retained, + storm.max_request_count + ); + } + } +}