From 4d3f53de09ea63f917a63a3ae635e62eff685c1a Mon Sep 17 00:00:00 2001 From: Nalu Tripician Date: Mon, 27 Jul 2026 11:25:52 -0700 Subject: [PATCH 01/10] Populate operation_name + active_instance metric Cosmos observability enrichment #3. Driver: add CosmosOperation::db_operation_name mapping to canonical OTel db.operation.name values, thread it through DiagnosticsContextBuilder (new field + setter + hedge propagation) and set it in the operation pipeline (execute_operation_direct) so DiagnosticsContext::operation_name is populated in production instead of always None. PATCH aggregates are stamped patch_item via DiagnosticsContext::with_operation_name so they no longer inherit the trailing Replace sub-op's name. This makes the point-vs non-point tail-sampling classification and the db.operation.name span/log attribute correct. SDK metrics: implement the previously-deferred azure.cosmosdb.client.active_instance.count as an opt-in up-down counter (MetricsOptions::with_active_instance_metric); CosmosMetricsHandler records +1 on construction and -1 on Drop so the value tracks live instrumented client instances. Request-scope per-request metrics remain a documented TODO (high cardinality; deferred per task scope). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- sdk/cosmos/azure_data_cosmos/CHANGELOG.md | 2 +- .../src/diagnostics/metrics/attributes.rs | 7 + .../src/diagnostics/metrics/handler.rs | 93 +++++++++++- .../src/diagnostics/metrics/instruments.rs | 19 ++- .../src/diagnostics/metrics/options.rs | 19 +++ .../azure_data_cosmos_driver/CHANGELOG.md | 2 + .../DIAGNOSTICS-CONTRACT.md | 2 +- .../src/diagnostics/diagnostics_context.rs | 89 ++++++++++- .../src/driver/cosmos_driver.rs | 11 +- .../src/driver/pipeline/patch_handler.rs | 42 +++++- .../src/models/cosmos_operation.rs | 140 ++++++++++++++++++ 11 files changed, 406 insertions(+), 20 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md index 9f7282ce222..d28d2b3a687 100644 --- a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md @@ -5,7 +5,7 @@ ### Features Added - Added a pluggable client-side diagnostics emission layer — the `DiagnosticsHandler` trait and ordered `DiagnosticsHandlerChain` (registered via `CosmosClientBuilder::with_diagnostics_handler`) — invoked once per operation (singleton and paginated, on success and failure) with the completed `DiagnosticsContext` plus an SDK-supplied `CosmosOperationContext`; the empty default chain is a zero-overhead no-op. ([#4789](https://github.com/Azure/azure-sdk-for-rust/pull/4789)) -- Added the `metrics`-gated `CosmosMetricsHandler` (with `MetricsOptions`), emitting the stable `db.client.operation.duration` histogram plus per-signal opt-in metrics (`with_request_charge_metric`, `with_returned_rows_metric`) and an opt-in extended attribute set (`with_extended_attributes`); a no-op when no meter provider is registered. ([#4789](https://github.com/Azure/azure-sdk-for-rust/pull/4789)) +- Added the `metrics`-gated `CosmosMetricsHandler` (with `MetricsOptions`), emitting the stable `db.client.operation.duration` histogram plus per-signal opt-in metrics (`with_request_charge_metric`, `with_returned_rows_metric`, and `with_active_instance_metric` — an `azure.cosmosdb.client.active_instance.count` up-down counter that is incremented on handler construction and decremented on drop to track live client instances) and an opt-in extended attribute set (`with_extended_attributes`); a no-op when no meter provider is registered. ([#4789](https://github.com/Azure/azure-sdk-for-rust/pull/4789)) - Added composable tail-sampled emission handlers — a `TracingLogHandler` leaf that writes a compact `tracing` line and a `SamplingLogHandler` wrapper (holding an `Arc`) that applies the sampling gate plus a shared per-window rate limit, defaulting to wrap a `TracingLogHandler` — and the `distributed_tracing`-gated `CosmosTracingHandler` (backdated span tree), also rate-limited so an error storm can't overwhelm exporters. All emit only for operations which fail or breach a configurable `DiagnosticsThresholds`, and stamp *why* they were sampled (a failure, or which threshold) on the emitted line and span. ([#4789](https://github.com/Azure/azure-sdk-for-rust/pull/4789)) ### Breaking Changes diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/attributes.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/attributes.rs index 6286e098e92..a076465d717 100644 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/attributes.rs +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/attributes.rs @@ -32,6 +32,10 @@ pub const METRIC_OPERATION_REQUEST_CHARGE: &str = "azure.cosmosdb.client.operati /// Optional histogram (rows): number of rows/items returned by an operation. pub const METRIC_RESPONSE_RETURNED_ROWS: &str = "db.client.response.returned_rows"; +/// Optional up-down counter (instances): number of live Cosmos client +/// instrumentation instances (one per active [`CosmosMetricsHandler`](super::CosmosMetricsHandler)). +pub const METRIC_ACTIVE_INSTANCE_COUNT: &str = "azure.cosmosdb.client.active_instance.count"; + // ========================================================================= // Instrument units // ========================================================================= @@ -49,6 +53,9 @@ pub const UNIT_REQUEST_UNIT: &str = "{request_unit}"; /// Unit for [`METRIC_RESPONSE_RETURNED_ROWS`] — rows. pub const UNIT_ROW: &str = "{row}"; +/// Unit for [`METRIC_ACTIVE_INSTANCE_COUNT`] — client instances. +pub const UNIT_INSTANCE: &str = "{instance}"; + // ========================================================================= // Stable attributes (always emitted; operation scope, low cardinality) // diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/handler.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/handler.rs index 1198782d289..b79a9cc5e8a 100644 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/handler.rs +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/handler.rs @@ -35,6 +35,12 @@ const METER_NAME: &str = "azure_data_cosmos"; /// histogram. The optional per-signal metrics and the extended attribute set are /// opt-in via [`MetricsOptions`] (see [`with_options`](CosmosMetricsHandler::with_options)). /// +/// When the active-instance metric is enabled +/// ([`MetricsOptions::with_active_instance_metric`]), the handler increments the +/// `azure.cosmosdb.client.active_instance.count` up-down counter on construction +/// and decrements it on [`Drop`], so the reported value tracks the number of +/// live instrumented client instances. +/// /// The handler captures a [`Meter`] from the globally-registered provider at /// construction. Install your meter provider **before** constructing the handler: /// a `Meter` obtained while the global provider is still the default no-op stays a @@ -76,10 +82,34 @@ impl CosmosMetricsHandler { } fn from_meter(meter: &Meter, options: MetricsOptions) -> Self { - Self { + let handler = Self { instruments: Instruments::new(meter), options, + }; + // Record the +1 half of the active-instance up-down counter at + // construction. One handler is created per instrumented client and + // held for that client's lifetime, so the handler's own lifecycle is + // a faithful proxy for a live client instance. The matching -1 is + // recorded in `Drop`. + if handler.options.active_instance_metric_enabled() { + handler + .instruments + .active_instance + .add(1, &Self::active_instance_attributes()); } + handler + } + + /// Low-cardinality attribute set for the active-instance up-down counter. + /// + /// Deliberately keyed on `db.system.name` alone so every client instance + /// aggregates into a single series whose value is the live instance count, + /// rather than fanning out per-instance. + fn active_instance_attributes() -> [KeyValue; 1] { + [KeyValue::new( + attributes::ATTR_DB_SYSTEM_NAME, + attributes::DB_SYSTEM_NAME_VALUE, + )] } /// Resolves `server.address`: the operation-context override if present, @@ -215,6 +245,18 @@ impl Default for CosmosMetricsHandler { } } +impl Drop for CosmosMetricsHandler { + fn drop(&mut self) { + // Record the -1 half of the active-instance up-down counter: the client + // instrumentation instance this handler represents is going away. + if self.options.active_instance_metric_enabled() { + self.instruments + .active_instance + .add(-1, &Self::active_instance_attributes()); + } + } +} + impl std::fmt::Debug for CosmosMetricsHandler { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { // Instruments and the seen-set are not meaningfully printable; surface @@ -353,6 +395,29 @@ mod tests { None } + /// Returns the summed value of the `active_instance.count` up-down counter + /// from the most recent export, or `None` if the metric was never emitted. + /// + /// The in-memory exporter accumulates one snapshot per `collect()` call, so + /// we scan every snapshot and keep the value from the last one — the current + /// cumulative count. + fn active_instance_value(metrics: &[ResourceMetrics]) -> Option { + let mut latest = None; + for rm in metrics { + for sm in rm.scope_metrics() { + for m in sm.metrics() { + if m.name() != attributes::METRIC_ACTIVE_INSTANCE_COUNT { + continue; + } + if let AggregatedMetrics::I64(MetricData::Sum(sum)) = m.data() { + latest = Some(sum.data_points().map(|point| point.value()).sum()); + } + } + } + } + latest + } + #[test] fn stable_duration_metric_carries_expected_attributes() { let harness = test_meter(); @@ -530,6 +595,32 @@ mod tests { handler.handle(&completed(200), &Context::new()); } + #[test] + fn active_instance_metric_off_by_default() { + // With default options the active-instance counter is never touched, so + // no such series is exported even across the handler's whole lifecycle. + let harness = test_meter(); + let handler = CosmosMetricsHandler::with_meter(harness.meter.clone()); + assert_eq!(active_instance_value(&harness.collect()), None); + drop(handler); + assert_eq!(active_instance_value(&harness.collect()), None); + } + + #[test] + fn active_instance_metric_tracks_handler_lifecycle() { + let harness = test_meter(); + let options = MetricsOptions::default().with_active_instance_metric(true); + + let handler = CosmosMetricsHandler::with_meter_and_options(harness.meter.clone(), options); + // +1 recorded at construction. + assert_eq!(active_instance_value(&harness.collect()), Some(1)); + + // Dropping the handler records the matching -1, so the up-down counter + // returns to zero — it reflects live instances, not a monotonic total. + drop(handler); + assert_eq!(active_instance_value(&harness.collect()), Some(0)); + } + #[test] fn host_of_extracts_host() { assert_eq!( diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/instruments.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/instruments.rs index 78e60c598e0..091847803c1 100644 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/instruments.rs +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/instruments.rs @@ -8,7 +8,7 @@ //! reference-counted). Building them eagerly keeps the per-operation hot path to //! just `record`/`add` calls with no allocation of instrument state. -use opentelemetry::metrics::{Histogram, Meter}; +use opentelemetry::metrics::{Histogram, Meter, UpDownCounter}; use crate::diagnostics::metrics::attributes; @@ -17,7 +17,8 @@ use crate::diagnostics::metrics::attributes; /// The stable operation-duration histogram is always recorded; the remaining /// per-signal instruments are recorded only when the matching /// [`MetricsOptions`](super::MetricsOptions) toggle -/// (`request_charge_metric_enabled` / `returned_rows_metric_enabled`) is set. +/// (`request_charge_metric_enabled` / `returned_rows_metric_enabled` / +/// `active_instance_metric_enabled`) is set. /// They are still created unconditionally because instrument creation is cheap /// and idempotent, and doing so keeps the handler's record path branch-free per /// instrument. @@ -31,6 +32,13 @@ pub(crate) struct Instruments { /// Development: `db.client.response.returned_rows` (rows). pub(crate) returned_rows: Histogram, + + /// Development: `azure.cosmosdb.client.active_instance.count` (instances). + /// + /// An up-down counter incremented when the handler is constructed and + /// decremented when it is dropped, so the reported value tracks the number + /// of live client instrumentation instances. + pub(crate) active_instance: UpDownCounter, } impl Instruments { @@ -54,10 +62,17 @@ impl Instruments { .with_description("Number of rows/items returned by a Cosmos DB operation.") .build(); + let active_instance = meter + .i64_up_down_counter(attributes::METRIC_ACTIVE_INSTANCE_COUNT) + .with_unit(attributes::UNIT_INSTANCE) + .with_description("Number of active Cosmos DB client instances.") + .build(); + Self { operation_duration, request_charge, returned_rows, + active_instance, } } } diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/options.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/options.rs index 2fddee6fc6a..eadc8038c20 100644 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/options.rs +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/options.rs @@ -33,15 +33,18 @@ /// let full = MetricsOptions::default() /// .with_request_charge_metric(true) /// .with_returned_rows_metric(true) +/// .with_active_instance_metric(true) /// .with_extended_attributes(true); /// assert!(full.request_charge_metric_enabled()); /// assert!(full.returned_rows_metric_enabled()); +/// assert!(full.active_instance_metric_enabled()); /// assert!(full.extended_attributes_enabled()); /// ``` #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub struct MetricsOptions { request_charge_metric: bool, returned_rows_metric: bool, + active_instance_metric: bool, extended_attributes: bool, } @@ -68,6 +71,17 @@ impl MetricsOptions { self } + /// Enables (or disables) the + /// `azure.cosmosdb.client.active_instance.count` up-down counter, which + /// tracks the number of live client instrumentation instances: it is + /// incremented by one when the [`CosmosMetricsHandler`](super::CosmosMetricsHandler) + /// is created and decremented by one when it is dropped. Off by default. + #[must_use] + pub fn with_active_instance_metric(mut self, enabled: bool) -> Self { + self.active_instance_metric = enabled; + self + } + /// Enables (or disables) the extended attribute set on every emitted metric: /// consistency level, contacted regions, sub-status code, and connection /// mode. These can be higher cardinality, so they are opt-in and off by @@ -88,6 +102,11 @@ impl MetricsOptions { self.returned_rows_metric } + /// Whether the active-instance up-down counter is emitted. + pub fn active_instance_metric_enabled(&self) -> bool { + self.active_instance_metric + } + /// Whether the extended attribute set is attached to emitted metrics. pub fn extended_attributes_enabled(&self) -> bool { self.extended_attributes diff --git a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md index e85982803d9..36c38f9a017 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md @@ -4,6 +4,8 @@ ### Features Added +- Added `CosmosOperation::db_operation_name`, returning the canonical OpenTelemetry `db.operation.name` (e.g. `read_item`, `query_items`, `execute_batch`) for an operation. The operation pipeline now populates `DiagnosticsContext::operation_name` from it in production (previously always `None`), so the emission layer's `db.operation.name` attribute and the point-vs.-non-point tail-sampling classification are accurate; PATCH aggregates report `patch_item` rather than the underlying Replace. ([#4789](https://github.com/Azure/azure-sdk-for-rust/pull/4789)) + ### Breaking Changes ### Bugs Fixed diff --git a/sdk/cosmos/azure_data_cosmos_driver/DIAGNOSTICS-CONTRACT.md b/sdk/cosmos/azure_data_cosmos_driver/DIAGNOSTICS-CONTRACT.md index 824cec0b3e9..f815836e1af 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/DIAGNOSTICS-CONTRACT.md +++ b/sdk/cosmos/azure_data_cosmos_driver/DIAGNOSTICS-CONTRACT.md @@ -351,7 +351,7 @@ attributes. This powers client-side Grafana dashboards (R7) with per-combination | `db.client.operation.duration` | stable | histogram | `s` | End-to-end operation duration — **the primary metric**. | | `db.client.response.returned_rows` | development | histogram | `{row}` | Rows returned in the result set. | | `azure.cosmosdb.client.operation.request_charge` | development | histogram | `{request_unit}` | Request units (RU) consumed by the operation. | -| `azure.cosmosdb.client.active_instance.count` | development | up-down counter | `{instance}` | Number of active Cosmos client instances. *(Deferred — not emitted yet; pending real client-lifecycle wiring, see [#4789](https://github.com/Azure/azure-sdk-for-rust/pull/4789).)* | +| `azure.cosmosdb.client.active_instance.count` | development | up-down counter | `{instance}` | Number of active Cosmos client instances. *(Opt-in via `MetricsOptions::with_active_instance_metric`; `CosmosMetricsHandler` records +1 on construction and −1 on `Drop`, so the value tracks live instrumented client instances — [#4789](https://github.com/Azure/azure-sdk-for-rust/pull/4789).)* | **Always-on metric attributes (low cardinality, D7):** `db.operation.name`, `db.response.status_code`, `db.collection.name`, `db.namespace`, `error.type`, `server.address`, diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/diagnostics_context.rs b/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/diagnostics_context.rs index 94cf02a77e8..b8a8b1a4350 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/diagnostics_context.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/diagnostics_context.rs @@ -1359,6 +1359,12 @@ pub(crate) struct DiagnosticsContextBuilder { /// Machine identifier (VM ID on Azure, generated UUID otherwise). machine_id: Option>, + /// Canonical `db.operation.name` for the operation (e.g. `read_item`), + /// when known. Set by the driver pipeline from + /// [`CosmosOperation::db_operation_name`](crate::models::CosmosOperation::db_operation_name) + /// and carried onto the finalized [`DiagnosticsContext::operation_name`]. + operation_name: Option>, + /// Whether fault injection is enabled for this operation's runtime. #[cfg(feature = "fault_injection")] fault_injection_enabled: bool, @@ -1383,6 +1389,7 @@ impl DiagnosticsContextBuilder { options, cpu_monitor: None, machine_id: None, + operation_name: None, #[cfg(feature = "fault_injection")] fault_injection_enabled: false, hedge_diagnostics: None, @@ -1401,6 +1408,13 @@ impl DiagnosticsContextBuilder { self.machine_id = Some(machine_id); } + /// Sets the canonical `db.operation.name` for this operation (e.g. + /// `read_item`). Carried onto the finalized + /// [`DiagnosticsContext::operation_name`]. + pub(crate) fn set_operation_name(&mut self, name: impl Into>) { + self.operation_name = Some(name.into()); + } + /// Sets the hedging diagnostics for this operation. pub(crate) fn set_hedge_diagnostics(&mut self, diagnostics: HedgeDiagnostics) { self.hedge_diagnostics = Some(diagnostics); @@ -1423,6 +1437,7 @@ impl DiagnosticsContextBuilder { options: Arc::clone(&self.options), cpu_monitor: self.cpu_monitor.clone(), machine_id: self.machine_id.clone(), + operation_name: self.operation_name.clone(), #[cfg(feature = "fault_injection")] fault_injection_enabled: self.fault_injection_enabled, hedge_diagnostics: None, @@ -1729,7 +1744,7 @@ impl DiagnosticsContextBuilder { options: self.options, cpu_monitor: self.cpu_monitor, machine_id: self.machine_id, - operation_name: None, + operation_name: self.operation_name, #[cfg(feature = "fault_injection")] fault_injection_enabled: self.fault_injection_enabled, #[cfg(not(feature = "fault_injection"))] @@ -1827,11 +1842,11 @@ pub struct DiagnosticsContext { /// Canonical `db.operation.name` for the operation (e.g. `read_item`, /// `query_items`), when known. /// - /// This is an optional seam for the emission layer: it feeds the - /// `db.operation.name` span/log attribute and lets + /// This feeds the `db.operation.name` span/log attribute and lets /// [`is_threshold_violated`](Self::is_threshold_violated) pick the point vs. - /// non-point latency threshold. It defaults to `None` — the driver pipeline - /// does not populate it yet, so today it is set only by test constructors. + /// non-point latency threshold. The driver pipeline populates it from + /// [`CosmosOperation::db_operation_name`](crate::models::CosmosOperation::db_operation_name); + /// it stays `None` for operations without a canonical name. operation_name: Option>, /// Whether fault injection was enabled when this operation executed. @@ -2224,12 +2239,27 @@ impl DiagnosticsContext { /// Returns the canonical `db.operation.name` for this operation, if known. /// /// Values are the semantic-convention operation names such as `read_item`, - /// `create_item`, or `query_items`. Returns `None` when the operation name - /// was not recorded (the common case today — see the field docs). + /// `create_item`, or `query_items`. The driver pipeline populates this from + /// [`CosmosOperation::db_operation_name`](crate::models::CosmosOperation::db_operation_name); + /// it is `None` only for operations without a canonical name (query plans, + /// partition-key-range reads, and other internal requests). pub fn operation_name(&self) -> Option<&str> { self.operation_name.as_deref() } + /// Returns this context with its canonical `db.operation.name` replaced. + /// + /// Used by aggregating callers (notably the PATCH handler) that build a + /// single operation-level context out of sub-operation contexts: the + /// aggregate would otherwise inherit the *last* sub-op's name (e.g. + /// `replace_item` for a PATCH's final Replace) instead of the virtual + /// operation's own name (`patch_item`). Consumes `self` before it is shared + /// via `Arc`, preserving the type's immutability contract. + pub(crate) fn with_operation_name(mut self, operation_name: Option>) -> Self { + self.operation_name = operation_name; + self + } + /// Returns `true` when this context represents a finished operation. /// /// A [`DiagnosticsContext`] is immutable and finalized at construction, so @@ -4556,6 +4586,51 @@ mod tests { assert_eq!(ctx.operation_name(), None); } + #[test] + fn set_operation_name_populates_completed_context() { + let ctx = make_context_with(ActivityId::new_uuid(), |b| { + b.set_operation_name("read_item"); + }); + assert_eq!(ctx.operation_name(), Some("read_item")); + } + + #[test] + fn with_operation_name_overrides_after_construction() { + let ctx = make_context_with(ActivityId::new_uuid(), |b| { + b.set_operation_name("replace_item"); + }); + assert_eq!(ctx.operation_name(), Some("replace_item")); + + let overridden = ctx.with_operation_name(Some(Arc::from("patch_item"))); + assert_eq!(overridden.operation_name(), Some("patch_item")); + + let cleared = overridden.with_operation_name(None); + assert_eq!(cleared.operation_name(), None); + } + + #[test] + fn aggregate_sub_operations_can_override_operation_name() { + // Mirrors the PATCH handler: a Read + Replace aggregate would inherit + // `replace_item` from the last source, but `with_operation_name` lets + // the caller stamp the virtual operation's own name. + let read = Arc::new(make_context_with(ActivityId::new_uuid(), |b| { + b.set_operation_name("read_item"); + })); + let replace = Arc::new(make_context_with(ActivityId::new_uuid(), |b| { + b.set_operation_name("replace_item"); + })); + + let inherited = + DiagnosticsContext::aggregate_sub_operations(&[read.clone(), replace.clone()]) + .expect("aggregation of two contexts yields Some"); + assert_eq!(inherited.operation_name(), Some("replace_item")); + + let stamped = DiagnosticsContext::aggregate_sub_operations(&[read, replace]) + .expect("aggregation of two contexts yields Some") + .with_operation_name(Some(Arc::from("patch_item"))); + assert_eq!(stamped.operation_name(), Some("patch_item")); + } + #[test] fn is_failure_reflects_operation_status() { let ok = make_context_with(ActivityId::new_uuid(), |b| { diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs index 3414323aceb..ef2c9cf5f02 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs @@ -2649,13 +2649,22 @@ impl CosmosDriver { false } }; - let (diagnostics_builder, transport_security) = Self::new_diagnostics_envelope( + let (mut diagnostics_builder, transport_security) = Self::new_diagnostics_envelope( &self.runtime, activity_id.clone(), &endpoint, fault_injection_enabled, ); + // Populate the canonical `db.operation.name` (e.g. `read_item`, + // `query_items`) so the finalized diagnostics carry it in production — + // feeding the emission layer's span/log attribute and the point-vs.- + // non-point tail-sampling classification. `None` for operations without + // a canonical name leaves it unset, matching prior behavior. + if let Some(operation_name) = operation.db_operation_name() { + diagnostics_builder.set_operation_name(operation_name); + } + let pipeline_type = if is_dataplane { PipelineType::DataPlane } else { diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/patch_handler.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/patch_handler.rs index df331af2dff..d9f8e06da3d 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/patch_handler.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/patch_handler.rs @@ -201,6 +201,12 @@ pub(crate) async fn execute_with_dispatcher( // Replace's. See `DiagnosticsContext::aggregate_sub_operations`. let mut sub_op_diagnostics: Vec> = Vec::with_capacity(2 * attempts as usize); + + // The aggregated context concatenates the Read + Replace sub-ops and would + // otherwise inherit the *last* sub-op's `db.operation.name` (`replace_item`). + // Stamp the virtual PATCH operation's own canonical name instead. + let operation_name: Option> = operation.db_operation_name().map(Arc::from); + for _ in 0..attempts { // Read the current item, propagating the freshest session token we // have observed so far (caller's on attempt 1; carried-forward on @@ -305,7 +311,7 @@ pub(crate) async fn execute_with_dispatcher( // (e.g. an empty source slice — which can't happen here, but // we keep the safe fallback for forward-compat). let diagnostics = DiagnosticsContext::aggregate_sub_operations(&sub_op_diagnostics) - .map(Arc::new) + .map(|ctx| Arc::new(ctx.with_operation_name(operation_name.clone()))) .unwrap_or_else(|| { sub_op_diagnostics .last() @@ -382,7 +388,12 @@ pub(crate) async fn execute_with_dispatcher( } } - Err(exhaustion_error(attempts, last_412, &sub_op_diagnostics)) + Err(exhaustion_error( + attempts, + last_412, + &sub_op_diagnostics, + operation_name, + )) } fn missing_body_error(msg: &'static str) -> crate::error::CosmosError { @@ -536,9 +547,11 @@ fn exhaustion_error( attempts: u8, last_412: Option, sub_op_diagnostics: &[Arc], + operation_name: Option>, ) -> crate::error::CosmosError { let message = format!("patch_item: ETag conflict after {attempts} attempts"); - let aggregated = DiagnosticsContext::aggregate_sub_operations(sub_op_diagnostics).map(Arc::new); + let aggregated = DiagnosticsContext::aggregate_sub_operations(sub_op_diagnostics) + .map(|ctx| Arc::new(ctx.with_operation_name(operation_name))); match last_412 { Some(source) => { let mut b = crate::error::CosmosErrorBuilder::from_error(source).with_context(message); @@ -915,7 +928,7 @@ mod tests { None, b"server-body", ); - let err = exhaustion_error(7, Some(underlying), &[]); + let err = exhaustion_error(7, Some(underlying), &[], Some(Arc::from("patch_item"))); // (a) Shape. assert_eq!( @@ -956,7 +969,7 @@ mod tests { // `attempts = 0` short-circuit), we still want the caller to see a // 412-shaped error so they can recognize "we gave up" the same way // they would for any other PATCH retry exhaustion. - let err = exhaustion_error(0, None, &[]); + let err = exhaustion_error(0, None, &[], Some(Arc::from("patch_item"))); assert_eq!(err.status().status_code(), StatusCode::PreconditionFailed); // No underlying service error was supplied, so the synthesized @@ -984,7 +997,7 @@ mod tests { Some("0:1#42"), b"{\"code\":\"PreconditionFailed\",\"message\":\"server: stale etag\"}", ); - let err = exhaustion_error(4, Some(underlying), &[]); + let err = exhaustion_error(4, Some(underlying), &[], Some(Arc::from("patch_item"))); assert_eq!(err.status().status_code(), StatusCode::PreconditionFailed); assert_eq!( @@ -1045,7 +1058,12 @@ mod tests { Arc::new(builder.complete()) }) .collect(); - let err = exhaustion_error(2, Some(underlying), &attempt_diags); + let err = exhaustion_error( + 2, + Some(underlying), + &attempt_diags, + Some(Arc::from("patch_item")), + ); let diag = err .diagnostics() @@ -1055,6 +1073,11 @@ mod tests { 4, "aggregated diagnostics must concatenate every per-attempt RequestDiagnostics", ); + assert_eq!( + diag.operation_name(), + Some("patch_item"), + "aggregated PATCH diagnostics must carry the virtual operation's own name", + ); // And critically, the attached diagnostics must be distinct from // every input Arc — the aggregator returns a fresh context. for input in &attempt_diags { @@ -1900,5 +1923,10 @@ mod tests { // The aggregated context inherits its activity_id from the LAST // source (the Replace), per `aggregate_sub_operations`'s contract. assert_eq!(returned.activity_id(), handed_out[1].activity_id()); + + // ...but its `db.operation.name` is the virtual PATCH operation's own + // name, not the Replace sub-op's, so telemetry labels the operation + // correctly. + assert_eq!(returned.operation_name(), Some("patch_item")); } } diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_operation.rs b/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_operation.rs index c3905213be9..e9f259043f5 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_operation.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_operation.rs @@ -149,6 +149,64 @@ impl CosmosOperation { self.resource_type } + /// Returns the canonical OpenTelemetry `db.operation.name` for this + /// operation, when it maps to a well-known Cosmos DB operation. + /// + /// The returned value uses the semantic-convention names the SDK surfaces + /// (`read_item`, `create_item`, `query_items`, `query_change_feed`, + /// `execute_batch`, `read_container`, …). It feeds + /// [`DiagnosticsContext::operation_name`](crate::diagnostics::DiagnosticsContext::operation_name) + /// so the emission layer can label spans/logs and + /// [`is_threshold_violated`](crate::diagnostics::DiagnosticsContext::is_threshold_violated) + /// can distinguish point from non-point operations for tail-based sampling. + /// + /// Operations without a canonical name (query plans, partition-key-range + /// reads, HEAD probes, stored procedures, triggers, UDFs, distributed + /// transactions) return `None`, which leaves the diagnostics + /// `operation_name` unset — identical to the pre-population behavior. + pub fn db_operation_name(&self) -> Option<&'static str> { + let name = match (self.operation_type, self.resource_type) { + // Data-plane item operations. + (OperationType::Create, ResourceType::Document) => "create_item", + (OperationType::Read, ResourceType::Document) => "read_item", + (OperationType::Replace, ResourceType::Document) => "replace_item", + (OperationType::Delete, ResourceType::Document) => "delete_item", + (OperationType::Upsert, ResourceType::Document) => "upsert_item", + (OperationType::Patch, ResourceType::Document) => "patch_item", + (OperationType::Batch, ResourceType::Document) => "execute_batch", + (OperationType::Query, ResourceType::Document) + | (OperationType::SqlQuery, ResourceType::Document) => "query_items", + (OperationType::ReadFeed, ResourceType::Document) => { + if self.is_change_feed { + "query_change_feed" + } else { + "read_all_items" + } + } + // Container (collection) management. + (OperationType::Create, ResourceType::DocumentCollection) => "create_container", + (OperationType::Read, ResourceType::DocumentCollection) => "read_container", + (OperationType::Replace, ResourceType::DocumentCollection) => "replace_container", + (OperationType::Delete, ResourceType::DocumentCollection) => "delete_container", + (OperationType::Query, ResourceType::DocumentCollection) + | (OperationType::SqlQuery, ResourceType::DocumentCollection) => "query_containers", + (OperationType::ReadFeed, ResourceType::DocumentCollection) => "read_all_containers", + // Database management. + (OperationType::Create, ResourceType::Database) => "create_database", + (OperationType::Read, ResourceType::Database) => "read_database", + (OperationType::Delete, ResourceType::Database) => "delete_database", + (OperationType::Query, ResourceType::Database) + | (OperationType::SqlQuery, ResourceType::Database) => "query_databases", + (OperationType::ReadFeed, ResourceType::Database) => "read_all_databases", + // Throughput (offer) management. + (OperationType::Read, ResourceType::Offer) => "read_throughput", + (OperationType::Replace, ResourceType::Offer) => "replace_throughput", + // Everything else has no canonical semconv name. + _ => return None, + }; + Some(name) + } + /// Returns a reference to the resource being operated on. pub(crate) fn resource_reference(&self) -> &CosmosResourceReference { &self.resource_reference @@ -1068,4 +1126,86 @@ mod tests { let resource_ref: CosmosResourceReference = item_ref.into(); let _op = CosmosOperation::new(OperationType::Create, resource_ref, None); } + + #[test] + fn db_operation_name_maps_item_operations() { + let item = + || ItemReference::from_name(&test_container(), PartitionKey::from("pk1"), "doc1"); + + assert_eq!( + CosmosOperation::create_item(item()).db_operation_name(), + Some("create_item") + ); + assert_eq!( + CosmosOperation::read_item(item()).db_operation_name(), + Some("read_item") + ); + assert_eq!( + CosmosOperation::replace_item(item()).db_operation_name(), + Some("replace_item") + ); + assert_eq!( + CosmosOperation::upsert_item(item()).db_operation_name(), + Some("upsert_item") + ); + assert_eq!( + CosmosOperation::delete_item(item()).db_operation_name(), + Some("delete_item") + ); + assert_eq!( + CosmosOperation::patch_item(item()).db_operation_name(), + Some("patch_item") + ); + } + + #[test] + fn db_operation_name_maps_feed_and_query_operations() { + assert_eq!( + CosmosOperation::query_items(test_container(), Some(FeedRange::full())) + .db_operation_name(), + Some("query_items") + ); + assert_eq!( + CosmosOperation::change_feed(test_container(), Some(FeedRange::full())) + .db_operation_name(), + Some("query_change_feed") + ); + assert_eq!( + CosmosOperation::read_all_items_cross_partition(test_container()).db_operation_name(), + Some("read_all_items") + ); + assert_eq!( + CosmosOperation::batch(test_container(), PartitionKey::from("pk1")).db_operation_name(), + Some("execute_batch") + ); + } + + #[test] + fn db_operation_name_maps_metadata_operations() { + let db = DatabaseReference::from_name(test_account(), "testdb"); + + assert_eq!( + CosmosOperation::read_container(test_container()).db_operation_name(), + Some("read_container") + ); + assert_eq!( + CosmosOperation::create_container(db.clone()).db_operation_name(), + Some("create_container") + ); + assert_eq!( + CosmosOperation::read_database(db.clone()).db_operation_name(), + Some("read_database") + ); + assert_eq!( + CosmosOperation::query_databases(test_account()).db_operation_name(), + Some("query_databases") + ); + } + + #[test] + fn db_operation_name_none_for_unmapped_operations() { + // Query plans have no canonical semconv operation name. + let op = CosmosOperation::query_plan(test_container(), std::borrow::Cow::Borrowed("")); + assert_eq!(op.db_operation_name(), None); + } } From 3c4bdd6b7d7d666afe48c1fa07eb437cf3e0a1f6 Mon Sep 17 00:00:00 2001 From: Nalu Tripician Date: Mon, 27 Jul 2026 13:09:19 -0700 Subject: [PATCH 02/10] Address PR review: consistent operation_name span label + throughput mapping - span_builder: prefer the caller-facing CosmosOperationContext operation name over the driver DiagnosticsContext (falling back to the driver value), so the tracing span label agrees with the db.operation.name metric and the tail-sampling classifier (both read the op context). Fixes a PATCH that fails during its internal Read being labeled read_item instead of patch_item. - cosmos_operation: map (Query|SqlQuery, Offer) => read_throughput so the user-facing throughput read (located via a query on the offers feed) is labeled, not just the poller's internal by-RID re-read. - Clarify active_instance docs across options/instruments/handler/attributes/ DIAGNOSTICS-CONTRACT: the counter tracks live CosmosMetricsHandler instances (one per client under the intended one-handler-per-client registration). - Note the intentional read_all_items/query_* naming divergence from .NET. - Fix CHANGELOG PR references to #4874; add throughput + span-label tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- sdk/cosmos/azure_data_cosmos/CHANGELOG.md | 2 +- .../src/diagnostics/metrics/attributes.rs | 5 ++- .../src/diagnostics/metrics/handler.rs | 3 +- .../src/diagnostics/metrics/instruments.rs | 4 +- .../src/diagnostics/metrics/options.rs | 8 ++-- .../src/diagnostics/tracing/mod.rs | 44 +++++++++++++++++-- .../src/diagnostics/tracing/span_builder.rs | 17 ++++--- .../azure_data_cosmos_driver/CHANGELOG.md | 2 +- .../DIAGNOSTICS-CONTRACT.md | 2 +- .../src/models/cosmos_operation.rs | 37 +++++++++++++++- 10 files changed, 104 insertions(+), 20 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md index d28d2b3a687..c37d583b15a 100644 --- a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md @@ -5,7 +5,7 @@ ### Features Added - Added a pluggable client-side diagnostics emission layer — the `DiagnosticsHandler` trait and ordered `DiagnosticsHandlerChain` (registered via `CosmosClientBuilder::with_diagnostics_handler`) — invoked once per operation (singleton and paginated, on success and failure) with the completed `DiagnosticsContext` plus an SDK-supplied `CosmosOperationContext`; the empty default chain is a zero-overhead no-op. ([#4789](https://github.com/Azure/azure-sdk-for-rust/pull/4789)) -- Added the `metrics`-gated `CosmosMetricsHandler` (with `MetricsOptions`), emitting the stable `db.client.operation.duration` histogram plus per-signal opt-in metrics (`with_request_charge_metric`, `with_returned_rows_metric`, and `with_active_instance_metric` — an `azure.cosmosdb.client.active_instance.count` up-down counter that is incremented on handler construction and decremented on drop to track live client instances) and an opt-in extended attribute set (`with_extended_attributes`); a no-op when no meter provider is registered. ([#4789](https://github.com/Azure/azure-sdk-for-rust/pull/4789)) +- Added the `metrics`-gated `CosmosMetricsHandler` (with `MetricsOptions`), emitting the stable `db.client.operation.duration` histogram plus per-signal opt-in metrics (`with_request_charge_metric`, `with_returned_rows_metric`, and `with_active_instance_metric` — an `azure.cosmosdb.client.active_instance.count` up-down counter that is incremented on handler construction and decremented on drop to track live client instances) and an opt-in extended attribute set (`with_extended_attributes`); a no-op when no meter provider is registered. ([#4789](https://github.com/Azure/azure-sdk-for-rust/pull/4789), [#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874)) - Added composable tail-sampled emission handlers — a `TracingLogHandler` leaf that writes a compact `tracing` line and a `SamplingLogHandler` wrapper (holding an `Arc`) that applies the sampling gate plus a shared per-window rate limit, defaulting to wrap a `TracingLogHandler` — and the `distributed_tracing`-gated `CosmosTracingHandler` (backdated span tree), also rate-limited so an error storm can't overwhelm exporters. All emit only for operations which fail or breach a configurable `DiagnosticsThresholds`, and stamp *why* they were sampled (a failure, or which threshold) on the emitted line and span. ([#4789](https://github.com/Azure/azure-sdk-for-rust/pull/4789)) ### Breaking Changes diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/attributes.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/attributes.rs index a076465d717..57b9e3839d8 100644 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/attributes.rs +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/attributes.rs @@ -32,8 +32,9 @@ pub const METRIC_OPERATION_REQUEST_CHARGE: &str = "azure.cosmosdb.client.operati /// Optional histogram (rows): number of rows/items returned by an operation. pub const METRIC_RESPONSE_RETURNED_ROWS: &str = "db.client.response.returned_rows"; -/// Optional up-down counter (instances): number of live Cosmos client -/// instrumentation instances (one per active [`CosmosMetricsHandler`](super::CosmosMetricsHandler)). +/// Optional up-down counter (instances): number of live +/// [`CosmosMetricsHandler`](super::CosmosMetricsHandler) instances (one per +/// instrumented client, under the intended one-handler-per-client registration). pub const METRIC_ACTIVE_INSTANCE_COUNT: &str = "azure.cosmosdb.client.active_instance.count"; // ========================================================================= diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/handler.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/handler.rs index b79a9cc5e8a..4d39406df6e 100644 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/handler.rs +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/handler.rs @@ -39,7 +39,8 @@ const METER_NAME: &str = "azure_data_cosmos"; /// ([`MetricsOptions::with_active_instance_metric`]), the handler increments the /// `azure.cosmosdb.client.active_instance.count` up-down counter on construction /// and decrements it on [`Drop`], so the reported value tracks the number of -/// live instrumented client instances. +/// live handler instances (one per instrumented client, under the intended +/// one-handler-per-client registration). /// /// The handler captures a [`Meter`] from the globally-registered provider at /// construction. Install your meter provider **before** constructing the handler: diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/instruments.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/instruments.rs index 091847803c1..8a7d4b0bd44 100644 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/instruments.rs +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/instruments.rs @@ -37,7 +37,9 @@ pub(crate) struct Instruments { /// /// An up-down counter incremented when the handler is constructed and /// decremented when it is dropped, so the reported value tracks the number - /// of live client instrumentation instances. + /// of live [`CosmosMetricsHandler`](super::CosmosMetricsHandler) instances + /// (one per instrumented client, under the intended one-handler-per-client + /// registration). pub(crate) active_instance: UpDownCounter, } diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/options.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/options.rs index eadc8038c20..5f1dfd1be79 100644 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/options.rs +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/options.rs @@ -73,9 +73,11 @@ impl MetricsOptions { /// Enables (or disables) the /// `azure.cosmosdb.client.active_instance.count` up-down counter, which - /// tracks the number of live client instrumentation instances: it is - /// incremented by one when the [`CosmosMetricsHandler`](super::CosmosMetricsHandler) - /// is created and decremented by one when it is dropped. Off by default. + /// tracks the number of live [`CosmosMetricsHandler`](super::CosmosMetricsHandler) + /// instances: it is incremented by one when the handler is created and + /// decremented by one when it is dropped. With the intended one-handler-per-client + /// registration this equals the number of live instrumented clients; sharing + /// a single handler across several clients reports one. Off by default. #[must_use] pub fn with_active_instance_metric(mut self, enabled: bool) -> Self { self.active_instance_metric = enabled; diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/mod.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/mod.rs index f678868c2d9..7d724c0ba50 100644 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/mod.rs +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/mod.rs @@ -208,9 +208,9 @@ mod tests { let now_instant = Instant::now(); let now_system = SystemTime::now(); - // A slow operation with NO driver-side operation name — the production - // case, since the driver never records one. The SDK-supplied - // CosmosOperationContext carries the operation identity instead. + // A slow operation with NO driver-side operation name — exercising the + // op-context fallback path. The SDK-supplied CosmosOperationContext + // carries the operation identity instead. let ctx = context( Duration::from_millis(1500), Some(CosmosStatus::new(StatusCode::Ok)), @@ -244,6 +244,44 @@ mod tests { })); } + #[test] + fn op_context_operation_name_wins_over_driver_context() { + // When both the driver context and the SDK operation context carry a + // name, the span must use the caller-facing op-context identity — so a + // PATCH whose surfaced sub-op is a Replace is still labeled `patch_item`, + // consistent with the `db.operation.name` metric and the tail-sampling + // classifier (which both read the op context). + let (provider, exporter) = exportable(); + let tracer = provider.tracer("test"); + + let now_instant = Instant::now(); + let now_system = SystemTime::now(); + let ctx = context( + Duration::from_millis(1500), + Some(CosmosStatus::new(StatusCode::Ok)), + Some("replace_item"), // driver-side sub-op name + &[(1500, 1500, CosmosStatus::new(StatusCode::Ok))], + now_instant, + ); + let op = CosmosOperationContext::new().with_operation_name("patch_item"); + + emit_backdated_span_tree(&tracer, &ctx, Some(&op), None, now_instant, now_system); + provider.force_flush().unwrap(); + + let spans = exporter.get_finished_spans().unwrap(); + let root = spans + .iter() + .find(|s| s.name == "patch_item") + .expect("root span named from op context, not the driver sub-op"); + assert!(root.attributes.iter().any(|kv| { + kv.key.as_str() == attributes::DB_OPERATION_NAME && kv.value.as_str() == "patch_item" + })); + assert!( + !spans.iter().any(|s| s.name == "replace_item"), + "the driver sub-op name must not label the operation span" + ); + } + #[test] fn incomplete_context_is_not_sampled_even_when_slow() { // A finalized context with neither a status nor any attempts does not diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/span_builder.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/span_builder.rs index f8152ec7a30..85673608480 100644 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/span_builder.rs +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/span_builder.rs @@ -101,11 +101,18 @@ pub(crate) fn emit_backdated_span_tree( let op_failed = diagnostics.is_failure(); // --- Operation (root) span --- - // Prefer the driver context's operation name; fall back to the SDK-supplied - // operation identity when the driver did not record one. - let op_name_ref = diagnostics - .operation_name() - .or_else(|| op.and_then(CosmosOperationContext::operation_name)); + // Prefer the SDK-supplied operation identity (the caller-facing operation, + // e.g. `patch_item`) so the span label agrees with the metric + // `db.operation.name` and the tail-sampling classifier, which both read the + // same `CosmosOperationContext`. Fall back to the driver context's operation + // name for operations not surfaced through the SDK wrapper (which therefore + // carry no `CosmosOperationContext`). Preferring the driver value here would + // mislabel an aggregate whose surfaced sub-op differs from the operation — + // e.g. a PATCH that fails during its internal Read would report `read_item` + // on the span while the metric reports `patch_item`. + let op_name_ref = op + .and_then(CosmosOperationContext::operation_name) + .or_else(|| diagnostics.operation_name()); let op_name = op_name_ref .unwrap_or(DEFAULT_OPERATION_SPAN_NAME) .to_string(); diff --git a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md index 36c38f9a017..94cb510ca36 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md @@ -4,7 +4,7 @@ ### Features Added -- Added `CosmosOperation::db_operation_name`, returning the canonical OpenTelemetry `db.operation.name` (e.g. `read_item`, `query_items`, `execute_batch`) for an operation. The operation pipeline now populates `DiagnosticsContext::operation_name` from it in production (previously always `None`), so the emission layer's `db.operation.name` attribute and the point-vs.-non-point tail-sampling classification are accurate; PATCH aggregates report `patch_item` rather than the underlying Replace. ([#4789](https://github.com/Azure/azure-sdk-for-rust/pull/4789)) +- Added `CosmosOperation::db_operation_name`, returning the canonical OpenTelemetry `db.operation.name` (e.g. `read_item`, `query_items`, `execute_batch`) for an operation. The operation pipeline now populates `DiagnosticsContext::operation_name` from it in production (previously always `None`), so `DiagnosticsContext::operation_name()` is populated for callers inspecting diagnostics and supplies the operation name to the emission layer's tail-sampling classifier and tracing span when no SDK-supplied `CosmosOperationContext` is present. The tracing span's operation label prefers the caller-facing `CosmosOperationContext` identity (consistent with the `db.operation.name` metric), and PATCH aggregates report `patch_item` rather than the underlying Replace. ([#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874)) ### Breaking Changes diff --git a/sdk/cosmos/azure_data_cosmos_driver/DIAGNOSTICS-CONTRACT.md b/sdk/cosmos/azure_data_cosmos_driver/DIAGNOSTICS-CONTRACT.md index f815836e1af..a13e2a20606 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/DIAGNOSTICS-CONTRACT.md +++ b/sdk/cosmos/azure_data_cosmos_driver/DIAGNOSTICS-CONTRACT.md @@ -351,7 +351,7 @@ attributes. This powers client-side Grafana dashboards (R7) with per-combination | `db.client.operation.duration` | stable | histogram | `s` | End-to-end operation duration — **the primary metric**. | | `db.client.response.returned_rows` | development | histogram | `{row}` | Rows returned in the result set. | | `azure.cosmosdb.client.operation.request_charge` | development | histogram | `{request_unit}` | Request units (RU) consumed by the operation. | -| `azure.cosmosdb.client.active_instance.count` | development | up-down counter | `{instance}` | Number of active Cosmos client instances. *(Opt-in via `MetricsOptions::with_active_instance_metric`; `CosmosMetricsHandler` records +1 on construction and −1 on `Drop`, so the value tracks live instrumented client instances — [#4789](https://github.com/Azure/azure-sdk-for-rust/pull/4789).)* | +| `azure.cosmosdb.client.active_instance.count` | development | up-down counter | `{instance}` | Number of active Cosmos client instances. *(Opt-in via `MetricsOptions::with_active_instance_metric`; `CosmosMetricsHandler` records +1 on construction and −1 on `Drop`, so the value tracks live `CosmosMetricsHandler` instances — one per instrumented client under the intended one-handler-per-client registration — [#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874).)* | **Always-on metric attributes (low cardinality, D7):** `db.operation.name`, `db.response.status_code`, `db.collection.name`, `db.namespace`, `error.type`, `server.address`, diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_operation.rs b/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_operation.rs index e9f259043f5..263bf102026 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_operation.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_operation.rs @@ -176,6 +176,14 @@ impl CosmosOperation { (OperationType::Batch, ResourceType::Document) => "execute_batch", (OperationType::Query, ResourceType::Document) | (OperationType::SqlQuery, ResourceType::Document) => "query_items", + // NOTE: `read_all_items` (and, below, `read_all_containers` / + // `read_all_databases` / the granular `query_containers` / + // `query_databases`) are this SDK's canonical values. They + // intentionally diverge from the .NET SDK, which emits + // `read_feed_ranges` for feed reads and funnels container/database + // queries through the generic `query_items`. Keep them aligned with + // this crate's own `read_all_*` / `query_*` public API, not with + // .NET. See DIAGNOSTICS-CONTRACT.md. (OperationType::ReadFeed, ResourceType::Document) => { if self.is_change_feed { "query_change_feed" @@ -198,8 +206,14 @@ impl CosmosOperation { (OperationType::Query, ResourceType::Database) | (OperationType::SqlQuery, ResourceType::Database) => "query_databases", (OperationType::ReadFeed, ResourceType::Database) => "read_all_databases", - // Throughput (offer) management. - (OperationType::Read, ResourceType::Offer) => "read_throughput", + // Throughput (offer) management. The user-facing throughput *read* + // locates the offer by querying the offers feed + // (`ContainerClient::read_throughput` -> `find_offer` -> + // `query_offers`), so its wire op is `(Query, Offer)`; `(Read, Offer)` + // is the throughput poller's internal by-RID re-read after a replace. + (OperationType::Read, ResourceType::Offer) + | (OperationType::Query, ResourceType::Offer) + | (OperationType::SqlQuery, ResourceType::Offer) => "read_throughput", (OperationType::Replace, ResourceType::Offer) => "replace_throughput", // Everything else has no canonical semconv name. _ => return None, @@ -1202,6 +1216,25 @@ mod tests { ); } + #[test] + fn db_operation_name_maps_throughput_operations() { + // Throughput reads locate the offer by querying the offers feed, so the + // user-facing read path is `(Query, Offer)`. + assert_eq!( + CosmosOperation::query_offers(test_account()).db_operation_name(), + Some("read_throughput") + ); + // `(Read, Offer)` is the throughput poller's internal by-RID re-read. + assert_eq!( + CosmosOperation::read_offer(test_account(), "offer-rid").db_operation_name(), + Some("read_throughput") + ); + assert_eq!( + CosmosOperation::replace_offer(test_account(), "offer-rid").db_operation_name(), + Some("replace_throughput") + ); + } + #[test] fn db_operation_name_none_for_unmapped_operations() { // Query plans have no canonical semconv operation name. From 9b5c5d1ec09ca63eb8c0ddaadf14ee052d76cc32 Mon Sep 17 00:00:00 2001 From: Nalu Tripician Date: Mon, 27 Jul 2026 15:05:34 -0700 Subject: [PATCH 03/10] Box::pin large emulator query futures to fix clippy The operation_name plumbing added in this PR grew the diagnostics context builder and driver operation state, tipping six emulator metadata/item query futures and two handler-propagation query futures just over clippy's `large_futures` deny threshold (16384 bytes) on the `--all-features --all-targets` CI Analyze job. Wrap the eight flagged `query_databases`/`query_containers`/ `query_items` futures in `Box::pin` (clippy's own suggestion), heap-allocating them so the awaited state is a pointer. Behavior is unchanged; only the two in_memory_emulator test files are touched. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../in_memory_emulator_tests/end_to_end.rs | 96 ++++++++++--------- .../handler_propagation.rs | 40 ++++---- 2 files changed, 70 insertions(+), 66 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/end_to_end.rs b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/end_to_end.rs index e4d19c51dce..7379c26ea51 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/end_to_end.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/end_to_end.rs @@ -546,25 +546,27 @@ async fn sdk_query_metadata_databases_and_containers() { let db_query = Query::from("SELECT * FROM c WHERE c.id = @id") .with_parameter("@id", db_name.as_str()) .unwrap(); - let emu_databases: Vec = backend - .emulator_client - .query_databases(db_query.clone(), None) - .await - .unwrap() - .try_collect() - .await - .unwrap(); + let emu_databases: Vec = Box::pin( + backend + .emulator_client + .query_databases(db_query.clone(), None), + ) + .await + .unwrap() + .try_collect() + .await + .unwrap(); assert_eq!(emu_databases.len(), 1); assert_eq!(emu_databases[0].id.as_deref(), Some(db_name.as_str())); if let Some(ref real_client) = backend.real_client { - let real_databases: Vec = real_client - .query_databases(db_query, None) - .await - .unwrap() - .try_collect() - .await - .unwrap(); + let real_databases: Vec = + Box::pin(real_client.query_databases(db_query, None)) + .await + .unwrap() + .try_collect() + .await + .unwrap(); assert_eq!(real_databases.len(), emu_databases.len()); assert_eq!(real_databases[0].id, emu_databases[0].id); } @@ -572,27 +574,31 @@ async fn sdk_query_metadata_databases_and_containers() { let container_query = Query::from("SELECT * FROM c WHERE c.id = @id") .with_parameter("@id", "testcoll") .unwrap(); - let emu_containers: Vec = backend - .emulator_client - .database_client(&db_name) - .query_containers(container_query.clone(), None) + let emu_containers: Vec = Box::pin( + backend + .emulator_client + .database_client(&db_name) + .query_containers(container_query.clone(), None), + ) + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(emu_containers.len(), 1); + assert_eq!(emu_containers[0].id, "testcoll"); + + if let Some(ref real_client) = backend.real_client { + let real_containers: Vec = Box::pin( + real_client + .database_client(&db_name) + .query_containers(container_query, None), + ) .await .unwrap() .try_collect() .await .unwrap(); - assert_eq!(emu_containers.len(), 1); - assert_eq!(emu_containers[0].id, "testcoll"); - - if let Some(ref real_client) = backend.real_client { - let real_containers: Vec = real_client - .database_client(&db_name) - .query_containers(container_query, None) - .await - .unwrap() - .try_collect() - .await - .unwrap(); assert_eq!(real_containers.len(), emu_containers.len()); assert_eq!(real_containers[0].id, emu_containers[0].id); } @@ -1085,25 +1091,25 @@ async fn sdk_query_items_with_filter_and_projection() { .unwrap() } - let emu_items: Vec = emu_container - .query_items(query(), FeedScope::partition("pk1"), None) - .await - .unwrap() - .try_collect() - .await - .unwrap(); - assert_eq!(emu_items.len(), 2); - assert_eq!(emu_items[0].id, "query-1"); - assert_eq!(emu_items[1].id, "query-2"); - - if let Some(ref real) = real_container { - let real_items: Vec = real - .query_items(query(), FeedScope::partition("pk1"), None) + let emu_items: Vec = + Box::pin(emu_container.query_items(query(), FeedScope::partition("pk1"), None)) .await .unwrap() .try_collect() .await .unwrap(); + assert_eq!(emu_items.len(), 2); + assert_eq!(emu_items[0].id, "query-1"); + assert_eq!(emu_items[1].id, "query-2"); + + if let Some(ref real) = real_container { + let real_items: Vec = + Box::pin(real.query_items(query(), FeedScope::partition("pk1"), None)) + .await + .unwrap() + .try_collect() + .await + .unwrap(); assert_eq!(real_items.len(), emu_items.len()); assert_eq!( real_items.iter().map(|i| &i.id).collect::>(), diff --git a/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/handler_propagation.rs b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/handler_propagation.rs index 319b8149c2f..b6dc56af179 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/handler_propagation.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/handler_propagation.rs @@ -234,17 +234,16 @@ async fn handler_receives_paginated_success() { } let before = handler.total(); - let items: Vec = c - .query_items( - Query::from("SELECT * FROM c"), - FeedScope::partition("pkQ"), - None, - ) - .await - .unwrap() - .try_collect() - .await - .unwrap(); + let items: Vec = Box::pin(c.query_items( + Query::from("SELECT * FROM c"), + FeedScope::partition("pkQ"), + None, + )) + .await + .unwrap() + .try_collect() + .await + .unwrap(); assert_eq!(items.len(), 3); assert_eq!( @@ -282,16 +281,15 @@ async fn handler_receives_paginated_failure() { // A syntactically invalid query is rejected by the emulator with a terminal // (non-retryable) 400 BadRequest, so the first page fetch errors instead of // returning a page — exercising the iterator's failure dispatch branch. - let result: Result, _> = c - .query_items( - Query::from("SELECT * FROM c WHERE"), - FeedScope::partition("pkFail"), - None, - ) - .await - .unwrap() - .try_collect() - .await; + let result: Result, _> = Box::pin(c.query_items( + Query::from("SELECT * FROM c WHERE"), + FeedScope::partition("pkFail"), + None, + )) + .await + .unwrap() + .try_collect() + .await; assert!( result.is_err(), "the invalid query must surface as a terminal page-fetch error" From fb80ac4d5fffdfcf22fc4cb97460bd71a938322f Mon Sep 17 00:00:00 2001 From: Nalu Tripician Date: Tue, 28 Jul 2026 11:54:46 -0700 Subject: [PATCH 04/10] Address PR review feedback on operation naming and metrics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Resolves the automated review comments on #4874. Operation naming (`cosmos_operation.rs`): a read feed scoped to a single logical partition now maps to the canonical `read_all_items_of_logical_partition` rather than `read_all_items`, and offer (throughput) operations are no longer mapped at all. Semconv has no unscoped `read_throughput`/`replace_throughput` — only the database- and container-scoped variants — and the driver's offer request carries no discriminator to tell them apart, so the scoped names are supplied by the SDK's operation contexts instead (`read_database_throughput`, `replace_database_throughput`, `read_container_throughput`, `replace_container_throughput`). PATCH identity (`patch_handler.rs`): every error exit now stamps `patch_item` on the diagnostics it propagates, so a failed read-modify-write is attributed to the caller-facing operation instead of the underlying Replace. Prior attempts are aggregated when more than one exists; the single-attempt case uses a new `DiagnosticsContext::clone_with_operation_name` so hedge diagnostics survive (`aggregate_sub_operations` drops them). Active-instance metric (`metrics/handler.rs`): the up-down counter was incremented in the handler constructor and decremented in its `Drop`, so it counted handler objects, not clients — one handler shared across N clients reported 1, and a handler built but never registered reported a phantom instance. The +1/-1 now ride a `ClientLifetimeToken` handed out by a new defaulted `DiagnosticsHandler::on_client_created` hook and owned by `ClientContext`, so the count follows live `CosmosClient` instances. Per semconv the counter is keyed on `server.address` (plus `server.port` for a non-default port), not `azure.client.id`; `DIAGNOSTICS-CONTRACT.md` D10 is corrected accordingly. Docs: rewrote the stale `is_threshold_violated_for` doc comment, and split the oversized changelog entries in both crates into focused single-purpose bullets, restoring the #4789 entry the active-instance note had been folded into. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 611a0d9f-360b-4c07-b45f-7dba858f7a76 --- sdk/cosmos/azure_data_cosmos/CHANGELOG.md | 4 +- .../src/clients/container_client.rs | 4 +- .../src/clients/cosmos_client_builder.rs | 11 +- .../src/clients/database_client.rs | 4 +- .../azure_data_cosmos/src/clients/mod.rs | 30 ++- .../src/diagnostics/attributes.rs | 6 + .../src/diagnostics/handler.rs | 118 +++++++++++ .../src/diagnostics/metrics/attributes.rs | 6 + .../src/diagnostics/metrics/handler.rs | 174 +++++++++++---- .../src/diagnostics/metrics/instruments.rs | 9 +- .../src/diagnostics/metrics/options.rs | 12 +- .../azure_data_cosmos/src/diagnostics/mod.rs | 4 +- .../azure_data_cosmos_driver/CHANGELOG.md | 6 +- .../DIAGNOSTICS-CONTRACT.md | 27 ++- .../src/diagnostics/diagnostics_context.rs | 47 ++++- .../src/driver/pipeline/patch_handler.rs | 199 +++++++++++++++--- .../src/models/cosmos_operation.rs | 69 ++++-- 17 files changed, 602 insertions(+), 128 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md index c37d583b15a..5e431dfcb1f 100644 --- a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md @@ -5,7 +5,9 @@ ### Features Added - Added a pluggable client-side diagnostics emission layer — the `DiagnosticsHandler` trait and ordered `DiagnosticsHandlerChain` (registered via `CosmosClientBuilder::with_diagnostics_handler`) — invoked once per operation (singleton and paginated, on success and failure) with the completed `DiagnosticsContext` plus an SDK-supplied `CosmosOperationContext`; the empty default chain is a zero-overhead no-op. ([#4789](https://github.com/Azure/azure-sdk-for-rust/pull/4789)) -- Added the `metrics`-gated `CosmosMetricsHandler` (with `MetricsOptions`), emitting the stable `db.client.operation.duration` histogram plus per-signal opt-in metrics (`with_request_charge_metric`, `with_returned_rows_metric`, and `with_active_instance_metric` — an `azure.cosmosdb.client.active_instance.count` up-down counter that is incremented on handler construction and decremented on drop to track live client instances) and an opt-in extended attribute set (`with_extended_attributes`); a no-op when no meter provider is registered. ([#4789](https://github.com/Azure/azure-sdk-for-rust/pull/4789), [#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874)) +- Added the `metrics`-gated `CosmosMetricsHandler` (with `MetricsOptions`), emitting the stable `db.client.operation.duration` histogram plus per-signal opt-in metrics (`with_request_charge_metric`, `with_returned_rows_metric`) and an opt-in extended attribute set (`with_extended_attributes`); a no-op when no meter provider is registered. ([#4789](https://github.com/Azure/azure-sdk-for-rust/pull/4789)) +- Added `MetricsOptions::with_active_instance_metric`, an opt-in `azure.cosmosdb.client.active_instance.count` up-down counter reporting the number of live `CosmosClient` instances per account endpoint, keyed on `server.address` (plus `server.port` for a non-default port). ([#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874)) +- Added `DiagnosticsHandler::on_client_created`, a defaulted hook that lets a handler observe client construction (`CosmosClientInfo`) and return a `ClientLifetimeToken` dropped with the client, for handlers that need to track client lifetime. ([#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874)) - Added composable tail-sampled emission handlers — a `TracingLogHandler` leaf that writes a compact `tracing` line and a `SamplingLogHandler` wrapper (holding an `Arc`) that applies the sampling gate plus a shared per-window rate limit, defaulting to wrap a `TracingLogHandler` — and the `distributed_tracing`-gated `CosmosTracingHandler` (backdated span tree), also rate-limited so an error storm can't overwhelm exporters. All emit only for operations which fail or breach a configurable `DiagnosticsThresholds`, and stamp *why* they were sampled (a failure, or which threshold) on the emitted line and span. ([#4789](https://github.com/Azure/azure-sdk-for-rust/pull/4789)) ### Breaking Changes diff --git a/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs b/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs index d3d4c0a4e9c..3d5eb446af7 100644 --- a/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs +++ b/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs @@ -179,7 +179,7 @@ impl ContainerClient { self.container_ref.account(), self.container_ref.rid(), options.operation, - self.operation_context("read_throughput"), + self.operation_context("read_container_throughput"), ) .await } @@ -222,7 +222,7 @@ impl ContainerClient { self.container_ref.rid(), throughput, options.operation, - self.operation_context("replace_throughput"), + self.operation_context("replace_container_throughput"), ) .await } diff --git a/sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client_builder.rs b/sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client_builder.rs index 21f7a86da37..9ae187ed030 100644 --- a/sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client_builder.rs +++ b/sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client_builder.rs @@ -7,7 +7,7 @@ use std::sync::Arc; use crate::{ clients::ClientContext, - diagnostics::DiagnosticsHandler, + diagnostics::{CosmosClientInfo, DiagnosticsHandler}, options::{ CosmosClientOptions, OperationOptions, PartitionFailoverOptions, ThroughputControlGroupOptions, UserAgentSuffix, @@ -282,6 +282,10 @@ impl CosmosClientBuilder { let (account_endpoint, credential) = account.into_parts(); let endpoint = account_endpoint.into_url(); + // Capture the account coordinates for client-scoped diagnostics before + // the endpoint is moved into the driver account. + let client_info = CosmosClientInfo::from_endpoint(&endpoint); + // Clone credential for the driver before the SDK consumes it for auth policy. let driver_credential = credential.clone(); @@ -305,10 +309,7 @@ impl CosmosClientBuilder { let driver = runtime.into_inner().create_driver(driver_options).await?; Ok(CosmosClient { - context: ClientContext { - driver, - diagnostics_handlers: self.options.diagnostics_handlers, - }, + context: ClientContext::new(driver, self.options.diagnostics_handlers, &client_info), }) } } diff --git a/sdk/cosmos/azure_data_cosmos/src/clients/database_client.rs b/sdk/cosmos/azure_data_cosmos/src/clients/database_client.rs index 4f9d851a82d..ccc44d400e7 100644 --- a/sdk/cosmos/azure_data_cosmos/src/clients/database_client.rs +++ b/sdk/cosmos/azure_data_cosmos/src/clients/database_client.rs @@ -246,7 +246,7 @@ impl DatabaseClient { self.context.driver.account(), &resource_id, options.operation, - self.operation_context("read_throughput"), + self.operation_context("read_database_throughput"), ) .await } @@ -292,7 +292,7 @@ impl DatabaseClient { &resource_id, throughput, options.operation, - self.operation_context("replace_throughput"), + self.operation_context("replace_database_throughput"), ) .await } diff --git a/sdk/cosmos/azure_data_cosmos/src/clients/mod.rs b/sdk/cosmos/azure_data_cosmos/src/clients/mod.rs index ffbe25d5e4a..86f0c13d619 100644 --- a/sdk/cosmos/azure_data_cosmos/src/clients/mod.rs +++ b/sdk/cosmos/azure_data_cosmos/src/clients/mod.rs @@ -40,7 +40,10 @@ use std::sync::Arc; use azure_data_cosmos_driver::CosmosDriver; -use crate::diagnostics::{CosmosOperationContext, DiagnosticsContext, DiagnosticsHandlerChain}; +use crate::diagnostics::{ + ClientLifetimeToken, CosmosClientInfo, CosmosOperationContext, DiagnosticsContext, + DiagnosticsHandlerChain, +}; use crate::models::CosmosResponse; /// Shared infrastructure threaded from [`CosmosClient`](super::CosmosClient) @@ -57,9 +60,34 @@ pub(crate) struct ClientContext { /// Empty by default, in which case the completion path does nothing beyond /// checking whether a handler is present. pub(crate) diagnostics_handlers: DiagnosticsHandlerChain, + /// Lifetime tokens handed back by the handlers when this client was built. + /// + /// Never read; held solely so the tokens' `Drop` runs when the last client + /// derived from this context goes away. Because the context is cloned down + /// into every `DatabaseClient`/`ContainerClient`, the shared `Arc` keeps the + /// tokens alive for as long as *any* of those clients is reachable, which is + /// the lifetime handlers are meant to observe. + _client_tokens: Arc<[ClientLifetimeToken]>, } impl ClientContext { + /// Builds the shared context for a newly constructed + /// [`CosmosClient`](super::CosmosClient), notifying every registered handler + /// that a client came online and taking ownership of the lifetime tokens + /// they hand back. + pub(crate) fn new( + driver: Arc, + diagnostics_handlers: DiagnosticsHandlerChain, + client_info: &CosmosClientInfo, + ) -> Self { + let client_tokens = diagnostics_handlers.dispatch_client_created(client_info); + Self { + driver, + diagnostics_handlers, + _client_tokens: client_tokens, + } + } + /// Converts a completed driver response into the SDK /// [`CosmosResponse`](crate::models::CosmosResponse) and invokes the /// diagnostics handler chain for the operation. diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/attributes.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/attributes.rs index b2f8c789716..0e54f25785a 100644 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics/attributes.rs +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/attributes.rs @@ -36,6 +36,12 @@ pub(crate) const DB_RESPONSE_STATUS_CODE: &str = "db.response.status_code"; /// `server.address` — the host contacted for the request. pub(crate) const SERVER_ADDRESS: &str = "server.address"; +/// `server.port` — the port contacted for the request. +/// +/// Per semantic conventions this is emitted only when the port differs from the +/// scheme's default (443 for HTTPS). +pub(crate) const SERVER_PORT: &str = "server.port"; + /// `error.type` — a low-cardinality identifier of the error (the status code). pub(crate) const ERROR_TYPE: &str = "error.type"; diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/handler.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/handler.rs index 78564ee2654..4f2071eb60c 100644 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics/handler.rs +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/handler.rs @@ -21,6 +21,91 @@ use azure_core::http::Context; use crate::diagnostics::DiagnosticsContext; +/// Identity of a Cosmos client instance, handed to +/// [`DiagnosticsHandler::on_client_created`] when a +/// [`CosmosClient`](crate::CosmosClient) is constructed. +/// +/// Carries only the account-level coordinates a handler needs to key +/// client-scoped telemetry; it deliberately exposes no credential material. +#[derive(Clone, Debug)] +pub struct CosmosClientInfo { + server_address: Option, + server_port: Option, +} + +impl CosmosClientInfo { + /// Builds the client identity from the account endpoint. + /// + /// `server_port` is populated only when the endpoint specifies a port other + /// than the scheme default, matching the `server.port` semantic convention. + pub(crate) fn from_endpoint(endpoint: &azure_core::http::Url) -> Self { + Self { + server_address: endpoint.host_str().map(str::to_owned), + server_port: endpoint.port(), + } + } + + /// The account endpoint's host, if the endpoint had one. + /// + /// Maps to the `server.address` semantic-convention attribute. + pub fn server_address(&self) -> Option<&str> { + self.server_address.as_deref() + } + + /// The account endpoint's port, when it is not the scheme default. + /// + /// Maps to the `server.port` semantic-convention attribute, which is only + /// emitted for non-default ports. + pub fn server_port(&self) -> Option { + self.server_port + } +} + +/// An opaque handle that represents one live client's registration with a +/// [`DiagnosticsHandler`]. +/// +/// A handler returns a token from +/// [`on_client_created`](DiagnosticsHandler::on_client_created) when it needs to +/// observe the end of that client's lifetime. The SDK stores the token in the +/// client's shared state, so it is dropped once the [`CosmosClient`](crate::CosmosClient) +/// and every client derived from it (database, container) have been dropped. +/// +/// Handlers use this to keep client-scoped state — such as the +/// `azure.cosmosdb.client.active_instance.count` up-down counter — balanced +/// without tying that state to the handler object's own lifetime (a single +/// handler may be registered on many clients, or on none). +pub struct ClientLifetimeToken { + on_drop: Option>, +} + +impl ClientLifetimeToken { + /// Creates a token that runs `on_drop` when the client it is attached to is + /// released. + /// + /// `on_drop` runs on whichever thread drops the last client handle, so it + /// must be cheap and non-blocking, and it must not panic. + pub fn new(on_drop: impl FnOnce() + Send + Sync + 'static) -> Self { + Self { + on_drop: Some(Box::new(on_drop)), + } + } +} + +impl Drop for ClientLifetimeToken { + fn drop(&mut self) { + if let Some(on_drop) = self.on_drop.take() { + on_drop(); + } + } +} + +impl fmt::Debug for ClientLifetimeToken { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ClientLifetimeToken") + .finish_non_exhaustive() + } +} + /// A sink that consumes a completed [`DiagnosticsContext`] for a single Cosmos /// operation. /// @@ -67,6 +152,24 @@ pub trait DiagnosticsHandler: Send + Sync { /// the caller's pipeline/trace context, so read it for operation metadata /// rather than for trace-context correlation. fn handle(&self, diagnostics: &DiagnosticsContext, cx: &Context<'_>); + + /// Notifies the handler that a new [`CosmosClient`](crate::CosmosClient) was + /// constructed with this handler registered. + /// + /// Called exactly once per client, at construction. Return a + /// [`ClientLifetimeToken`] to be notified when that client — and every + /// database/container client derived from it — has been dropped; return + /// `None` (the default) when the handler does not track client lifetimes. + /// + /// This is the seam for client-scoped telemetry. A handler object may be + /// shared across several clients or registered on none, so its own lifetime + /// is not a proxy for a live client; this hook and the returned token are. + /// + /// * `client` - Account-level identity of the newly created client. + fn on_client_created(&self, client: &CosmosClientInfo) -> Option { + let _ = client; + None + } } /// An ordered, cheaply cloneable chain of [`DiagnosticsHandler`]s. @@ -143,6 +246,21 @@ impl DiagnosticsHandlerChain { handler.handle(diagnostics, cx); } } + + /// Notifies every handler that a client was created, collecting the lifetime + /// tokens they hand back. + /// + /// The returned tokens must be stored for the client's lifetime; dropping + /// them is what signals client teardown to the handlers. + pub(crate) fn dispatch_client_created( + &self, + client: &CosmosClientInfo, + ) -> Arc<[ClientLifetimeToken]> { + self.handlers + .iter() + .filter_map(|handler| handler.on_client_created(client)) + .collect() + } } impl Default for DiagnosticsHandlerChain { diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/attributes.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/attributes.rs index 57b9e3839d8..223a4c26042 100644 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/attributes.rs +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/attributes.rs @@ -88,6 +88,12 @@ pub const ATTR_ERROR_TYPE: &str = attributes::ERROR_TYPE; /// `server.address` — host of the contacted endpoint. pub const ATTR_SERVER_ADDRESS: &str = attributes::SERVER_ADDRESS; +/// `server.port` — port of the contacted endpoint. +/// +/// Conditionally required: emitted only when the endpoint uses a non-default +/// port (i.e. anything other than 443 for HTTPS). +pub const ATTR_SERVER_PORT: &str = attributes::SERVER_PORT; + /// Fallback value for [`ATTR_ERROR_TYPE`] when the error is otherwise unknown /// (per semantic conventions). pub const ERROR_TYPE_OTHER: &str = attributes::ERROR_TYPE_OTHER; diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/handler.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/handler.rs index 4d39406df6e..1cc92a2188a 100644 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/handler.rs +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/handler.rs @@ -11,7 +11,10 @@ use opentelemetry::{global, Array, KeyValue, StringValue, Value}; use crate::diagnostics::metrics::attributes; use crate::diagnostics::metrics::instruments::Instruments; use crate::diagnostics::metrics::MetricsOptions; -use crate::diagnostics::{CosmosOperationContext, DiagnosticsContext, DiagnosticsHandler}; +use crate::diagnostics::{ + ClientLifetimeToken, CosmosClientInfo, CosmosOperationContext, DiagnosticsContext, + DiagnosticsHandler, +}; /// Instrumentation scope name used for the Cosmos [`Meter`]. const METER_NAME: &str = "azure_data_cosmos"; @@ -37,10 +40,12 @@ const METER_NAME: &str = "azure_data_cosmos"; /// /// When the active-instance metric is enabled /// ([`MetricsOptions::with_active_instance_metric`]), the handler increments the -/// `azure.cosmosdb.client.active_instance.count` up-down counter on construction -/// and decrements it on [`Drop`], so the reported value tracks the number of -/// live handler instances (one per instrumented client, under the intended -/// one-handler-per-client registration). +/// `azure.cosmosdb.client.active_instance.count` up-down counter each time a +/// [`CosmosClient`](crate::CosmosClient) is built with it registered, and +/// decrements it when that client (and every database/container client derived +/// from it) is dropped. The reported value is therefore the number of live +/// client instances per account endpoint, independent of how many handler +/// objects exist. /// /// The handler captures a [`Meter`] from the globally-registered provider at /// construction. Install your meter provider **before** constructing the handler: @@ -83,34 +88,34 @@ impl CosmosMetricsHandler { } fn from_meter(meter: &Meter, options: MetricsOptions) -> Self { - let handler = Self { + Self { instruments: Instruments::new(meter), options, - }; - // Record the +1 half of the active-instance up-down counter at - // construction. One handler is created per instrumented client and - // held for that client's lifetime, so the handler's own lifecycle is - // a faithful proxy for a live client instance. The matching -1 is - // recorded in `Drop`. - if handler.options.active_instance_metric_enabled() { - handler - .instruments - .active_instance - .add(1, &Self::active_instance_attributes()); } - handler } - /// Low-cardinality attribute set for the active-instance up-down counter. + /// Attribute set for the active-instance up-down counter. /// - /// Deliberately keyed on `db.system.name` alone so every client instance - /// aggregates into a single series whose value is the live instance count, - /// rather than fanning out per-instance. - fn active_instance_attributes() -> [KeyValue; 1] { - [KeyValue::new( + /// Per the `azure.cosmosdb.client.active_instance.count` semantic + /// convention, the counter is keyed on the account endpoint + /// (`server.address`, plus `server.port` only when the endpoint uses a + /// non-default port), so the value reads as "live clients per account". + fn active_instance_attributes(client: &CosmosClientInfo) -> Vec { + let mut attrs = Vec::with_capacity(3); + attrs.push(KeyValue::new( attributes::ATTR_DB_SYSTEM_NAME, attributes::DB_SYSTEM_NAME_VALUE, - )] + )); + if let Some(address) = client.server_address() { + attrs.push(KeyValue::new( + attributes::ATTR_SERVER_ADDRESS, + address.to_string(), + )); + } + if let Some(port) = client.server_port() { + attrs.push(KeyValue::new(attributes::ATTR_SERVER_PORT, i64::from(port))); + } + attrs } /// Resolves `server.address`: the operation-context override if present, @@ -246,18 +251,6 @@ impl Default for CosmosMetricsHandler { } } -impl Drop for CosmosMetricsHandler { - fn drop(&mut self) { - // Record the -1 half of the active-instance up-down counter: the client - // instrumentation instance this handler represents is going away. - if self.options.active_instance_metric_enabled() { - self.instruments - .active_instance - .add(-1, &Self::active_instance_attributes()); - } - } -} - impl std::fmt::Debug for CosmosMetricsHandler { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { // Instruments and the seen-set are not meaningfully printable; surface @@ -292,6 +285,23 @@ impl DiagnosticsHandler for CosmosMetricsHandler { } } } + + fn on_client_created(&self, client: &CosmosClientInfo) -> Option { + if !self.options.active_instance_metric_enabled() { + return None; + } + + // Record the +1 half of the up-down counter now, and hand back a token + // whose `Drop` records the matching -1. The token rides on the client's + // shared state, so the counter tracks live *clients* rather than live + // handler objects — a single handler may be registered on many clients. + let attributes = Self::active_instance_attributes(client); + let counter = self.instruments.active_instance.clone(); + counter.add(1, &attributes); + Some(ClientLifetimeToken::new(move || { + counter.add(-1, &attributes); + })) + } } /// Extracts the host portion of an endpoint URI for `server.address`. @@ -396,8 +406,18 @@ mod tests { None } - /// Returns the summed value of the `active_instance.count` up-down counter - /// from the most recent export, or `None` if the metric was never emitted. + /// Builds a [`CosmosClientInfo`] for `host`, optionally on a non-default + /// port, the way `CosmosClientBuilder::build` would from an account + /// endpoint. + fn test_client_info(host: &str, port: Option) -> CosmosClientInfo { + let url = match port { + Some(port) => format!("https://{host}:{port}/"), + None => format!("https://{host}/"), + }; + CosmosClientInfo::from_endpoint(&url::Url::parse(&url).expect("valid test endpoint")) + } + + /// Returns the summed value of the `active_instance.count` up-down counter /// from the most recent export, or `None` if the metric was never emitted. /// /// The in-memory exporter accumulates one snapshot per `collect()` call, so /// we scan every snapshot and keep the value from the last one — the current @@ -599,29 +619,93 @@ mod tests { #[test] fn active_instance_metric_off_by_default() { // With default options the active-instance counter is never touched, so - // no such series is exported even across the handler's whole lifecycle. + // no such series is exported even across a full client lifecycle. let harness = test_meter(); let handler = CosmosMetricsHandler::with_meter(harness.meter.clone()); + let token = handler.on_client_created(&test_client_info("acct.documents.azure.com", None)); + assert!(token.is_none(), "disabled metric must not take a token"); assert_eq!(active_instance_value(&harness.collect()), None); - drop(handler); + drop(token); assert_eq!(active_instance_value(&harness.collect()), None); } #[test] - fn active_instance_metric_tracks_handler_lifecycle() { + fn active_instance_metric_tracks_client_lifecycle_not_handler_lifecycle() { let harness = test_meter(); let options = MetricsOptions::default().with_active_instance_metric(true); let handler = CosmosMetricsHandler::with_meter_and_options(harness.meter.clone(), options); - // +1 recorded at construction. + // Constructing the handler alone records nothing: a handler is not a + // client, and one handler may be registered on many clients or none. + assert_eq!(active_instance_value(&harness.collect()), None); + + let info = test_client_info("acct.documents.azure.com", None); + let first = handler + .on_client_created(&info) + .expect("enabled metric must take a lifetime token"); + assert_eq!(active_instance_value(&harness.collect()), Some(1)); + + // The same handler registered on a second client counts twice — the + // regression this replaces counted handler objects, so it reported 1. + let second = handler + .on_client_created(&info) + .expect("enabled metric must take a lifetime token"); + assert_eq!(active_instance_value(&harness.collect()), Some(2)); + + drop(first); assert_eq!(active_instance_value(&harness.collect()), Some(1)); - // Dropping the handler records the matching -1, so the up-down counter - // returns to zero — it reflects live instances, not a monotonic total. + // Dropping the handler while a client token is still alive must not + // decrement: the client, not the handler, owns the count. drop(handler); + assert_eq!(active_instance_value(&harness.collect()), Some(1)); + + drop(second); assert_eq!(active_instance_value(&harness.collect()), Some(0)); } + #[test] + fn active_instance_metric_is_keyed_on_account_endpoint() { + // Per semconv the counter carries `server.address`, and `server.port` + // only when the endpoint uses a non-default port. + let default_port = CosmosMetricsHandler::active_instance_attributes(&test_client_info( + "acct.documents.azure.com", + None, + )); + let by_key: HashMap<_, _> = default_port + .iter() + .map(|kv| (kv.key.as_str().to_string(), kv.value.as_str().to_string())) + .collect(); + assert_eq!( + by_key + .get(attributes::ATTR_DB_SYSTEM_NAME) + .map(String::as_str), + Some(attributes::DB_SYSTEM_NAME_VALUE) + ); + assert_eq!( + by_key + .get(attributes::ATTR_SERVER_ADDRESS) + .map(String::as_str), + Some("acct.documents.azure.com") + ); + assert!( + !by_key.contains_key(attributes::ATTR_SERVER_PORT), + "default port must be omitted; got {by_key:?}" + ); + + let custom_port = CosmosMetricsHandler::active_instance_attributes(&test_client_info( + "localhost", + Some(8081), + )); + assert!( + custom_port + .iter() + .any(|kv| kv.key.as_str() == attributes::ATTR_SERVER_PORT + && kv.value.as_str() == "8081"), + "non-default port must be emitted; got {custom_port:?}" + ); + } + #[test] fn host_of_extracts_host() { assert_eq!( diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/instruments.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/instruments.rs index 8a7d4b0bd44..1645351c801 100644 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/instruments.rs +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/instruments.rs @@ -35,11 +35,10 @@ pub(crate) struct Instruments { /// Development: `azure.cosmosdb.client.active_instance.count` (instances). /// - /// An up-down counter incremented when the handler is constructed and - /// decremented when it is dropped, so the reported value tracks the number - /// of live [`CosmosMetricsHandler`](super::CosmosMetricsHandler) instances - /// (one per instrumented client, under the intended one-handler-per-client - /// registration). + /// An up-down counter incremented when a + /// [`CosmosClient`](crate::CosmosClient) is created with the handler + /// registered, and decremented when that client is dropped, so the reported + /// value tracks the number of live client instances per account endpoint. pub(crate) active_instance: UpDownCounter, } diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/options.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/options.rs index 5f1dfd1be79..adbe5ee0b78 100644 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/options.rs +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/options.rs @@ -73,11 +73,13 @@ impl MetricsOptions { /// Enables (or disables) the /// `azure.cosmosdb.client.active_instance.count` up-down counter, which - /// tracks the number of live [`CosmosMetricsHandler`](super::CosmosMetricsHandler) - /// instances: it is incremented by one when the handler is created and - /// decremented by one when it is dropped. With the intended one-handler-per-client - /// registration this equals the number of live instrumented clients; sharing - /// a single handler across several clients reports one. Off by default. + /// tracks the number of live [`CosmosClient`](crate::CosmosClient) + /// instances: it is incremented by one when a client is built with this + /// handler registered and decremented by one when that client — and every + /// database/container client derived from it — has been dropped. The counter + /// is keyed on the account endpoint (`server.address`, plus `server.port` + /// for non-default ports), so sharing one handler across several clients + /// still reports each client. Off by default. #[must_use] pub fn with_active_instance_metric(mut self, enabled: bool) -> Self { self.active_instance_metric = enabled; diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/mod.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/mod.rs index e785d31cfdc..fe754d13136 100644 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics/mod.rs +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/mod.rs @@ -46,7 +46,9 @@ pub use azure_data_cosmos_driver::diagnostics::{ }; #[doc(inline)] pub use azure_data_cosmos_driver::DiagnosticsThresholds; -pub use handler::{DiagnosticsHandler, DiagnosticsHandlerChain}; +pub use handler::{ + ClientLifetimeToken, CosmosClientInfo, DiagnosticsHandler, DiagnosticsHandlerChain, +}; pub use logging::{SamplingLogHandler, TracingLogHandler}; pub use operation_context::CosmosOperationContext; pub use rate_limiter::RateLimiterConfig; diff --git a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md index 94cb510ca36..3cf7f919097 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md @@ -4,12 +4,16 @@ ### Features Added -- Added `CosmosOperation::db_operation_name`, returning the canonical OpenTelemetry `db.operation.name` (e.g. `read_item`, `query_items`, `execute_batch`) for an operation. The operation pipeline now populates `DiagnosticsContext::operation_name` from it in production (previously always `None`), so `DiagnosticsContext::operation_name()` is populated for callers inspecting diagnostics and supplies the operation name to the emission layer's tail-sampling classifier and tracing span when no SDK-supplied `CosmosOperationContext` is present. The tracing span's operation label prefers the caller-facing `CosmosOperationContext` identity (consistent with the `db.operation.name` metric), and PATCH aggregates report `patch_item` rather than the underlying Replace. ([#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874)) +- Added `CosmosOperation::db_operation_name`, returning the canonical OpenTelemetry `db.operation.name` (e.g. `read_item`, `query_items`, `execute_batch`, and `read_all_items_of_logical_partition` for a read feed scoped to one logical partition) for an operation. ([#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874)) ### Breaking Changes ### Bugs Fixed +- `DiagnosticsContext::operation_name()` is now populated in production (previously always `None`): the operation pipeline sets it from `CosmosOperation::db_operation_name`, so tail-sampling classification and the tracing span have an operation name even when no SDK-supplied `CosmosOperationContext` is present. ([#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874)) +- The tracing span's operation label now prefers the caller-facing `CosmosOperationContext` identity, matching how the `db.operation.name` metric attribute is resolved. ([#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874)) +- PATCH operations now report `patch_item` rather than the underlying Replace, on both the aggregated success path and every error path (including read, deserialize, patch-evaluation, serialize, and non-412 replace failures). ([#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874)) + ### Other Changes ## 0.6.1 (2026-07-23) diff --git a/sdk/cosmos/azure_data_cosmos_driver/DIAGNOSTICS-CONTRACT.md b/sdk/cosmos/azure_data_cosmos_driver/DIAGNOSTICS-CONTRACT.md index a13e2a20606..58a2bbc58fc 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/DIAGNOSTICS-CONTRACT.md +++ b/sdk/cosmos/azure_data_cosmos_driver/DIAGNOSTICS-CONTRACT.md @@ -351,7 +351,7 @@ attributes. This powers client-side Grafana dashboards (R7) with per-combination | `db.client.operation.duration` | stable | histogram | `s` | End-to-end operation duration — **the primary metric**. | | `db.client.response.returned_rows` | development | histogram | `{row}` | Rows returned in the result set. | | `azure.cosmosdb.client.operation.request_charge` | development | histogram | `{request_unit}` | Request units (RU) consumed by the operation. | -| `azure.cosmosdb.client.active_instance.count` | development | up-down counter | `{instance}` | Number of active Cosmos client instances. *(Opt-in via `MetricsOptions::with_active_instance_metric`; `CosmosMetricsHandler` records +1 on construction and −1 on `Drop`, so the value tracks live `CosmosMetricsHandler` instances — one per instrumented client under the intended one-handler-per-client registration — [#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874).)* | +| `azure.cosmosdb.client.active_instance.count` | development | up-down counter | `{instance}` | Number of live Cosmos client instances per account endpoint. *(Opt-in via `MetricsOptions::with_active_instance_metric`; `CosmosMetricsHandler` records +1 when a `CosmosClient` is built with it registered and −1 when that client and every client derived from it is dropped, so the value is independent of how many handler objects exist — [#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874).)* | **Always-on metric attributes (low cardinality, D7):** `db.operation.name`, `db.response.status_code`, `db.collection.name`, `db.namespace`, `error.type`, `server.address`, @@ -399,18 +399,31 @@ a metric dimension to control time-series cardinality (D7). | `azure.client.id` | Stable per-client instance id (see D10). | | `azure.resource_provider.namespace` | `Microsoft.DocumentDB`. | -> **Client instance id (D10).** `azure.client.id` and the active-instance metric use a stable -> per-client id. Prefer `vmId`; when VM metadata is unreachable, fall back to a **static GUID** -> so two requests can be attributed to the same `CosmosClient`/driver instance. **Check whether -> [`DiagnosticsContext`][ctx] already carries this before adding it.** +> **Client instance id (D10).** `azure.client.id` is a stable per-client instance id. Prefer +> `vmId`; when VM metadata is unreachable, fall back to a **static GUID** so two requests can be +> attributed to the same `CosmosClient`/driver instance. **Check whether +> [`DiagnosticsContext`][ctx] already carries this before adding it.** Note that +> `azure.cosmosdb.client.active_instance.count` is *not* keyed on `azure.client.id`: semconv +> defines its attribute set as `server.address` plus `server.port` (the latter only for a +> non-default port), so the counter reads as "live clients per account endpoint" +> ([#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874)). #### 10.3.1 Canonical `db.operation.name` values Use the semconv well-known values verbatim — e.g. `read_item`, `create_item`, `upsert_item`, -`replace_item`, `patch_item`, `delete_item`, `query_items`, `read_all_items`, `read_many_items`, -`execute_batch`, `execute_bulk`, `query_change_feed`, `read_container`, `create_container`, … . +`replace_item`, `patch_item`, `delete_item`, `query_items`, `read_all_items`, +`read_all_items_of_logical_partition`, `read_many_items`, `execute_batch`, `execute_bulk`, +`query_change_feed`, `read_container`, `create_container`, … . If none applies, use a language-agnostic snake_case method name. +Some canonical values encode a scope the driver cannot see. Throughput (offer) operations are the +notable case: semconv distinguishes `read_database_throughput` / `replace_database_throughput` from +`read_container_throughput` / `replace_container_throughput`, but the driver's offer requests carry +no database-vs-container discriminator, and there is no unscoped `read_throughput`. Those names are +therefore supplied by the SDK's `CosmosOperationContext`, and the driver's mapping leaves offer +operations unnamed rather than emitting a non-canonical value +([#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874)). + ### 10.4 Traces (span tree) — reconstructed & backdated Emitted by `CosmosTracingHandler` [WS4], **only when tail-sampling (§5.2) says so**. diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/diagnostics_context.rs b/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/diagnostics_context.rs index b8a8b1a4350..39505657c68 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/diagnostics_context.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/diagnostics_context.rs @@ -2260,6 +2260,39 @@ impl DiagnosticsContext { self } + /// Returns a copy of this context with its canonical `db.operation.name` + /// replaced, leaving every other field — including status, hedging + /// diagnostics, and compaction metadata — intact. + /// + /// [`with_operation_name`](Self::with_operation_name) consumes `self`, which + /// works when the caller still owns a freshly aggregated context. Error + /// paths instead hold an `Arc` that a deeper layer + /// already attached to a [`CosmosError`](crate::error::CosmosError), so they + /// need to re-stamp the identity without taking ownership. The JSON caches + /// are intentionally not carried over: they may already have been rendered + /// with the old name. + pub(crate) fn clone_with_operation_name(&self, operation_name: Option>) -> Self { + DiagnosticsContext { + activity_id: self.activity_id.clone(), + duration: self.duration, + requests: Arc::clone(&self.requests), + total_request_charge: self.total_request_charge, + regions_contacted: self.regions_contacted.clone(), + status: self.status, + options: Arc::clone(&self.options), + cpu_monitor: self.cpu_monitor.clone(), + machine_id: self.machine_id.clone(), + operation_name, + fault_injection_enabled: self.fault_injection_enabled, + hedge_diagnostics: self.hedge_diagnostics.clone(), + #[cfg(test)] + test_system_usage: self.test_system_usage.clone(), + compaction: self.compaction.clone(), + cached_json_detailed: OnceLock::new(), + cached_json_summary: OnceLock::new(), + } + } + /// Returns `true` when this context represents a finished operation. /// /// A [`DiagnosticsContext`] is immutable and finalized at construction, so @@ -2308,12 +2341,14 @@ impl DiagnosticsContext { /// Like [`is_threshold_violated`](Self::is_threshold_violated), but takes an /// explicit operation name for point/non-point latency classification. /// - /// Production `DiagnosticsContext`s do not carry an operation name, so the - /// SDK's emission handlers pass the caller-facing name from the - /// `CosmosOperationContext` here; otherwise every operation would be - /// classified with the stricter 1s point-operation threshold. When - /// `operation_name` is `None` this falls back to - /// [`operation_name`](Self::operation_name), then to the point threshold. + /// The driver stamps its own canonical name onto the context, so the + /// explicit argument is an *override*: the SDK's emission handlers pass the + /// caller-facing name from the `CosmosOperationContext` so classification + /// matches the name the caller sees, which also covers operations the driver + /// leaves unmapped (throughput reads, for instance, whose canonical name is + /// scope-dependent). When `operation_name` is `None` this falls back to + /// [`operation_name`](Self::operation_name), then to the stricter point + /// threshold. pub fn is_threshold_violated_for( &self, thresholds: &DiagnosticsThresholds, diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/patch_handler.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/patch_handler.rs index d9f8e06da3d..20ce172c3c4 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/patch_handler.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/patch_handler.rs @@ -215,22 +215,34 @@ pub(crate) async fn execute_with_dispatcher( // Any non-2xx Read response is mapped by the driver pipeline into // `Err(ErrorKind::HttpResponse { .. })` (see retry_evaluation.rs's - // `build_http_error`). Propagating with `?` is sufficient — the - // caller wants the original error verbatim, complete with - // `raw_response` and diagnostics — and there is nothing useful the - // PATCH handler can do on a Read failure. + // `build_http_error`). The caller wants that error verbatim, complete + // with `raw_response`, status, and source — there is nothing useful the + // PATCH handler can do on a Read failure — but the diagnostics riding on + // it still describe the *sub-op* (`read_item`). Re-stamp the virtual + // PATCH operation's identity so the failure reports the same + // `db.operation.name` as its success and retry-exhaustion counterparts. let read_resp = dispatcher .execute_operation(read_op, options.clone()) - .await?; + .await + .map_err(|err| { + stamp_patch_identity(err, operation_name.clone(), &sub_op_diagnostics) + })?; sub_op_diagnostics.push(read_resp.diagnostics()); - let etag = read_resp.headers().etag.clone().ok_or_else(|| { - crate::error::CosmosError::builder() - .with_status(crate::error::CosmosStatus::new( - azure_core::http::StatusCode::BadRequest, - )) - .with_message("PATCH cannot proceed: the Read response did not include an ETag") - .build() - })?; + let etag = read_resp + .headers() + .etag + .clone() + .ok_or_else(|| { + crate::error::CosmosError::builder() + .with_status(crate::error::CosmosStatus::new( + azure_core::http::StatusCode::BadRequest, + )) + .with_message("PATCH cannot proceed: the Read response did not include an ETag") + .build() + }) + .map_err(|err| { + stamp_patch_identity(err, operation_name.clone(), &sub_op_diagnostics) + })?; // R3-DRIVER: forward the session token returned by the Read on the // Replace, so the write commits against the same replica view we // just read from. This is what mitigates SE-004 (session token @@ -243,16 +255,25 @@ pub(crate) async fn execute_with_dispatcher( effective_session_token = Some(token); } - // Locally apply the patch ops. - let read_body_bytes = read_resp.into_body().single().map_err(|err| { - crate::error::CosmosError::builder() - .with_status(crate::error::CosmosStatus::SERIALIZATION_RESPONSE_BODY_INVALID) - .with_message("PATCH could not extract Read response body") - .with_source(err) - .build() - })?; - let mut value: serde_json::Value = - serde_json::from_slice(&read_body_bytes).map_err(|err| { + // Locally apply the patch ops. These failures are synthesized here + // rather than returned by the pipeline, so they carry no diagnostics of + // their own; hand them the PATCH-identified aggregate of the sub-ops + // issued so far. + let read_body_bytes = read_resp + .into_body() + .single() + .map_err(|err| { + crate::error::CosmosError::builder() + .with_status(crate::error::CosmosStatus::SERIALIZATION_RESPONSE_BODY_INVALID) + .with_message("PATCH could not extract Read response body") + .with_source(err) + .build() + }) + .map_err(|err| { + stamp_patch_identity(err, operation_name.clone(), &sub_op_diagnostics) + })?; + let mut value: serde_json::Value = serde_json::from_slice(&read_body_bytes) + .map_err(|err| { crate::error::CosmosError::builder() .with_status(crate::error::CosmosStatus::SERIALIZATION_RESPONSE_BODY_INVALID) .with_message(format!( @@ -260,15 +281,24 @@ pub(crate) async fn execute_with_dispatcher( )) .with_source(err) .build() + }) + .map_err(|err| { + stamp_patch_identity(err, operation_name.clone(), &sub_op_diagnostics) })?; - apply_patch_ops(&mut value, &spec.operations)?; - let merged_bytes = serde_json::to_vec(&value).map_err(|err| { - crate::error::CosmosError::builder() - .with_status(crate::error::CosmosStatus::SERIALIZATION_RESPONSE_BODY_INVALID) - .with_message("PATCH could not serialize merged item") - .with_source(err) - .build() + apply_patch_ops(&mut value, &spec.operations).map_err(|err| { + stamp_patch_identity(err.into(), operation_name.clone(), &sub_op_diagnostics) })?; + let merged_bytes = serde_json::to_vec(&value) + .map_err(|err| { + crate::error::CosmosError::builder() + .with_status(crate::error::CosmosStatus::SERIALIZATION_RESPONSE_BODY_INVALID) + .with_message("PATCH could not serialize merged item") + .with_source(err) + .build() + }) + .map_err(|err| { + stamp_patch_identity(err, operation_name.clone(), &sub_op_diagnostics) + })?; // Issue the ETag-guarded Replace, forwarding the Read response's // session token (overriding any caller-supplied value). @@ -384,7 +414,13 @@ pub(crate) async fn execute_with_dispatcher( last_412 = Some(err); continue; } - Err(err) => return Err(err), + Err(err) => { + return Err(stamp_patch_identity( + err, + operation_name.clone(), + &sub_op_diagnostics, + )) + } } } @@ -396,6 +432,49 @@ pub(crate) async fn execute_with_dispatcher( )) } +/// Re-stamps the virtual PATCH operation's canonical `db.operation.name` onto +/// the diagnostics attached to a failure escaping the RMW loop. +/// +/// The handler executes 2+ real sub-operations (`read_item` + `replace_item`), +/// so a failure surfaced verbatim from a sub-op would report that sub-op's +/// identity while the matching success and retry-exhaustion paths report +/// `patch_item`. This keeps "one PATCH operation = one `DiagnosticsContext`" +/// true on every exit. +/// +/// The wire error itself flows through untouched — status, sub-status, raw +/// response, and source are carried forward by +/// [`CosmosErrorBuilder::from_error`] — only the diagnostics are replaced. +/// `prior_sub_ops` are the contexts accumulated before the failure; when there +/// are any, the failing sub-op's context is aggregated with them so the error +/// carries the whole PATCH attempt history. With a single context there is +/// nothing to aggregate, so it is copied verbatim (preserving hedging +/// diagnostics and compaction metadata) with only the name rewritten. Errors +/// with no diagnostics anywhere are returned unchanged; the operation pipeline +/// grafts the operation-level context onto them on the way out. +fn stamp_patch_identity( + err: crate::error::CosmosError, + operation_name: Option>, + prior_sub_ops: &[Arc], +) -> crate::error::CosmosError { + let mut sources: Vec> = prior_sub_ops.to_vec(); + if let Some(failed) = err.diagnostics() { + sources.push(failed); + } + let stamped = match sources.as_slice() { + [] => return err, + [only] => Arc::new(only.clone_with_operation_name(operation_name)), + many => match DiagnosticsContext::aggregate_sub_operations(many) { + Some(ctx) => Arc::new(ctx.with_operation_name(operation_name)), + // Unreachable: `many` is non-empty. Keep the error intact rather + // than panicking if that ever changes. + None => return err, + }, + }; + crate::error::CosmosErrorBuilder::from_error(err) + .with_diagnostics(stamped) + .build() +} + fn missing_body_error(msg: &'static str) -> crate::error::CosmosError { crate::error::CosmosError::builder() .with_status(crate::error::CosmosStatus::new( @@ -1427,6 +1506,16 @@ mod tests { "non-412 must propagate verbatim; got {:?}", err.status() ); + // The wire failure keeps its own status/response, but its diagnostics + // must be labeled with the virtual PATCH operation rather than the + // `replace_item` sub-op that actually failed. + assert_eq!( + err.diagnostics() + .as_deref() + .and_then(DiagnosticsContext::operation_name), + Some("patch_item"), + "non-412 Replace failure must carry the PATCH operation identity" + ); // Single Read + single Replace — no retry. assert_eq!(dispatcher.calls().len(), 2); } @@ -1456,6 +1545,16 @@ mod tests { "PATCH on missing item must surface the Read's 404 verbatim; got {:?}", err.status() ); + // The Read's own diagnostics ride along on the error, but they must be + // re-labeled with the virtual PATCH operation's name so a failed PATCH + // is never reported as a `read_item`. + assert_eq!( + err.diagnostics() + .as_deref() + .and_then(DiagnosticsContext::operation_name), + Some("patch_item"), + "Read failure must carry the PATCH operation identity" + ); // Exactly one sub-op was issued: the Read. No Replace. let calls = dispatcher.calls(); assert_eq!(calls.len(), 1, "no Replace must be issued on Read failure"); @@ -1486,6 +1585,44 @@ mod tests { assert_eq!(calls[0].op_type, OperationType::Read); } + #[tokio::test] + async fn rmw_read_error_on_retry_aggregates_prior_attempts() { + // A Read failure on attempt 2 must still be labeled `patch_item` and + // must fold in attempt 1's sub-op diagnostics, so the error reports the + // whole PATCH — not just the sub-op that happened to fail. + let read_failure = http_error(StatusCode::ServiceUnavailable, "read down"); + let failure_diagnostics = read_failure + .diagnostics() + .expect("fixture error carries diagnostics"); + let dispatcher = ScriptedDispatcher::new(vec![ + ScriptedReply::ok( + br#"{"id":"doc1","pk":"pk1","visits":0}"#.to_vec(), + Some("\"v1\""), + StatusCode::Ok, + ), + ScriptedReply::Err(http_error(StatusCode::PreconditionFailed, "etag conflict")), + ScriptedReply::Err(read_failure), + ]); + + let err = execute_with_dispatcher( + &dispatcher, + canonical_patch_op(), + OperationOptions::default(), + NonZeroU8::new(3), + ) + .await + .expect_err("Read failure on retry must abort the loop"); + + assert_eq!(err.status().status_code(), StatusCode::ServiceUnavailable); + let diagnostics = err.diagnostics().expect("error must carry diagnostics"); + assert_eq!(diagnostics.operation_name(), Some("patch_item")); + assert!( + !Arc::ptr_eq(&diagnostics, &failure_diagnostics), + "with prior sub-ops in flight the error's diagnostics must be an aggregate, \ + not the failing sub-op's context verbatim" + ); + } + #[tokio::test] async fn pk_guard_rejection_issues_no_sub_operations() { // Gap #4 closure: when the PK guard fires, the handler MUST return diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_operation.rs b/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_operation.rs index 263bf102026..8e09ff6f14d 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_operation.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_operation.rs @@ -164,6 +164,11 @@ impl CosmosOperation { /// reads, HEAD probes, stored procedures, triggers, UDFs, distributed /// transactions) return `None`, which leaves the diagnostics /// `operation_name` unset — identical to the pre-population behavior. + /// Throughput (offer) operations are also unmapped: the canonical names are + /// scope-specific (`read_container_throughput` vs. `read_database_throughput`) + /// but an offer operation carries only the account and the offer ID, so the + /// scope is not recoverable here. The SDK, which knows whether the caller + /// addressed a container or a database, supplies those names instead. pub fn db_operation_name(&self) -> Option<&'static str> { let name = match (self.operation_type, self.resource_type) { // Data-plane item operations. @@ -187,6 +192,11 @@ impl CosmosOperation { (OperationType::ReadFeed, ResourceType::Document) => { if self.is_change_feed { "query_change_feed" + } else if self.targets_logical_partition() { + // `read_all_items(container, partition_key)` narrows the + // feed to one logical partition, which semconv names + // distinctly from the cross-partition read. + "read_all_items_of_logical_partition" } else { "read_all_items" } @@ -206,21 +216,27 @@ impl CosmosOperation { (OperationType::Query, ResourceType::Database) | (OperationType::SqlQuery, ResourceType::Database) => "query_databases", (OperationType::ReadFeed, ResourceType::Database) => "read_all_databases", - // Throughput (offer) management. The user-facing throughput *read* - // locates the offer by querying the offers feed - // (`ContainerClient::read_throughput` -> `find_offer` -> - // `query_offers`), so its wire op is `(Query, Offer)`; `(Read, Offer)` - // is the throughput poller's internal by-RID re-read after a replace. - (OperationType::Read, ResourceType::Offer) - | (OperationType::Query, ResourceType::Offer) - | (OperationType::SqlQuery, ResourceType::Offer) => "read_throughput", - (OperationType::Replace, ResourceType::Offer) => "replace_throughput", + // Throughput (offer) management has no driver-layer mapping: the + // canonical names are scope-specific (`read_container_throughput` / + // `read_database_throughput` and their `replace_` variants), but an + // offer operation is addressed by account + offer ID only, so this + // layer cannot tell a container offer from a database offer. The + // SDK stamps the scoped name via `CosmosOperationContext`. // Everything else has no canonical semconv name. _ => return None, }; Some(name) } + /// Returns `true` when this operation targets exactly one logical partition + /// (or a hierarchical-partition-key prefix), as opposed to an EPK range or + /// the whole container. + fn targets_logical_partition(&self) -> bool { + self.target + .as_ref() + .is_some_and(FeedRange::is_logical_partition) + } + /// Returns a reference to the resource being operated on. pub(crate) fn resource_reference(&self) -> &CosmosResourceReference { &self.resource_reference @@ -1188,12 +1204,32 @@ mod tests { CosmosOperation::read_all_items_cross_partition(test_container()).db_operation_name(), Some("read_all_items") ); + assert_eq!( + CosmosOperation::read_all_items(test_container(), PartitionKey::from("pk1")) + .db_operation_name(), + Some("read_all_items_of_logical_partition") + ); assert_eq!( CosmosOperation::batch(test_container(), PartitionKey::from("pk1")).db_operation_name(), Some("execute_batch") ); } + #[test] + fn db_operation_name_change_feed_ignores_logical_partition_scope() { + // A change feed scoped to one logical partition is still + // `query_change_feed`; semconv has no partition-scoped variant for it. + let container = test_container(); + let range = FeedRange::for_partition( + PartitionKey::from("pk1"), + container.partition_key_definition(), + ); + assert_eq!( + CosmosOperation::change_feed(container, Some(range)).db_operation_name(), + Some("query_change_feed") + ); + } + #[test] fn db_operation_name_maps_metadata_operations() { let db = DatabaseReference::from_name(test_account(), "testdb"); @@ -1217,21 +1253,22 @@ mod tests { } #[test] - fn db_operation_name_maps_throughput_operations() { - // Throughput reads locate the offer by querying the offers feed, so the - // user-facing read path is `(Query, Offer)`. + fn db_operation_name_none_for_throughput_operations() { + // Offer operations carry no database/container scope, and semconv only + // defines scoped throughput names, so the driver leaves them unmapped + // and the SDK supplies `read_container_throughput` / + // `read_database_throughput` (and their `replace_` variants). assert_eq!( CosmosOperation::query_offers(test_account()).db_operation_name(), - Some("read_throughput") + None ); - // `(Read, Offer)` is the throughput poller's internal by-RID re-read. assert_eq!( CosmosOperation::read_offer(test_account(), "offer-rid").db_operation_name(), - Some("read_throughput") + None ); assert_eq!( CosmosOperation::replace_offer(test_account(), "offer-rid").db_operation_name(), - Some("replace_throughput") + None ); } From ac2001ceed6d401ef7d3d81c0e3c5581a1389e8e Mon Sep 17 00:00:00 2001 From: Nalu Tripician Date: Wed, 29 Jul 2026 12:48:37 -0700 Subject: [PATCH 05/10] Address remaining review feedback on operation naming metrics Picks up the three review points that were raised but left open on #4874. `DiagnosticsHandlerChain::dispatch_client_created` visited every chain entry, so registering the same `Arc` twice (the chain is additive, so that is reachable) recorded two `+1`s for a single client and reported it as two live instances. Client creation is a per-client lifecycle event rather than a per-call dispatch, so the chain now notifies each distinct handler at most once per client, deduplicating by handler identity. Distinct handler objects are still notified independently. Three tests pin the repeated-handler, distinct-handler, and no-token cases. The `DIAGNOSTICS-CONTRACT.md` D7 table claimed the counter was "independent of how many handler objects exist", which overstated the guarantee: two distinct metrics handlers built from the same meter each record their own `+1`, exactly as they each record their own duration and request-charge samples. Both the contract and `MetricsOptions::with_active_instance_metric` now state what is actually guaranteed and point users at one metrics handler per meter. The tracing span operation-label precedence entry was filed under the driver's Bugs Fixed, but that behavior lives in `azure_data_cosmos/src/diagnostics/tracing/span_builder.rs`; listing it in the driver changelog misstates what driver-only consumers receive. Moved to the SDK changelog, with the symptom spelled out. Also splits a doc comment on the `active_instance_value` test helper that had two `///` prefixes collapsed onto one line. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: b5586b49-f116-489f-81a4-d51fb2511a71 --- sdk/cosmos/azure_data_cosmos/CHANGELOG.md | 2 + .../src/diagnostics/handler.rs | 105 +++++++++++++++++- .../src/diagnostics/metrics/handler.rs | 3 +- .../src/diagnostics/metrics/options.rs | 8 +- .../azure_data_cosmos_driver/CHANGELOG.md | 1 - .../DIAGNOSTICS-CONTRACT.md | 2 +- 6 files changed, 112 insertions(+), 9 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md index a705812e60f..e0795bffbba 100644 --- a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md @@ -21,6 +21,8 @@ ### Bugs Fixed +- The Cosmos tracing span's operation label now prefers the caller-facing `CosmosOperationContext` identity over the driver-recorded name, matching how the `db.operation.name` metric attribute is resolved. Previously an aggregate whose surfaced sub-operation differed from the caller's operation — such as a PATCH that fails during its internal read — could label the span `read_item` while the metric reported `patch_item`. ([#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874)) + ### Other Changes - Cosmos HTTP error messages now include the service's own explanation from the response body, normalized to a single line and length-bounded, so a `400` no longer renders as a bare `Cosmos DB returned HTTP 400: Unknown`. ([#4904](https://github.com/Azure/azure-sdk-for-rust/pull/4904)) diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/handler.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/handler.rs index 4f2071eb60c..8ed94c7e777 100644 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics/handler.rs +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/handler.rs @@ -156,7 +156,8 @@ pub trait DiagnosticsHandler: Send + Sync { /// Notifies the handler that a new [`CosmosClient`](crate::CosmosClient) was /// constructed with this handler registered. /// - /// Called exactly once per client, at construction. Return a + /// Called exactly once per client, at construction — including when the same + /// handler was registered on that client's chain more than once. Return a /// [`ClientLifetimeToken`] to be notified when that client — and every /// database/container client derived from it — has been dropped; return /// `None` (the default) when the handler does not track client lifetimes. @@ -250,16 +251,29 @@ impl DiagnosticsHandlerChain { /// Notifies every handler that a client was created, collecting the lifetime /// tokens they hand back. /// + /// A handler that appears in the chain more than once — the chain is + /// additive, so the same `Arc` can be registered twice — is notified only + /// once, since this is a per-client lifecycle event rather than a per-call + /// dispatch. Distinct handler objects are always notified independently. + /// /// The returned tokens must be stored for the client's lifetime; dropping /// them is what signals client teardown to the handlers. pub(crate) fn dispatch_client_created( &self, client: &CosmosClientInfo, ) -> Arc<[ClientLifetimeToken]> { - self.handlers - .iter() - .filter_map(|handler| handler.on_client_created(client)) - .collect() + let mut notified: Vec<&Arc> = Vec::new(); + let mut tokens = Vec::new(); + for handler in self.handlers.iter() { + if notified.iter().any(|seen| Arc::ptr_eq(seen, handler)) { + continue; + } + notified.push(handler); + if let Some(token) = handler.on_client_created(client) { + tokens.push(token); + } + } + tokens.into() } } @@ -375,4 +389,85 @@ mod tests { let recorded = log.lock().unwrap().clone(); assert_eq!(recorded, vec![("first", op.clone()), ("second", op)]); } + + /// A handler that counts `on_client_created` calls and hands back a token + /// which records the matching teardown. + struct LifecycleHandler { + created: Arc>, + dropped: Arc>, + } + + impl DiagnosticsHandler for LifecycleHandler { + fn handle(&self, _diagnostics: &DiagnosticsContext, _cx: &Context<'_>) {} + + fn on_client_created(&self, _client: &CosmosClientInfo) -> Option { + *self.created.lock().unwrap() += 1; + let dropped = Arc::clone(&self.dropped); + Some(ClientLifetimeToken::new(move || { + *dropped.lock().unwrap() += 1; + })) + } + } + + fn client_info() -> CosmosClientInfo { + CosmosClientInfo::from_endpoint( + &azure_core::http::Url::parse("https://account.documents.azure.com/") + .expect("valid test endpoint"), + ) + } + + #[test] + fn client_created_notifies_a_repeated_handler_once() { + let created = Arc::new(Mutex::new(0)); + let dropped = Arc::new(Mutex::new(0)); + let handler: Arc = Arc::new(LifecycleHandler { + created: Arc::clone(&created), + dropped: Arc::clone(&dropped), + }); + + // The chain is additive, so the same handler can land on it twice. That + // is one registration for client-lifecycle purposes: a client must not + // be counted twice just because a handler was added twice. + let chain = DiagnosticsHandlerChain::new() + .with_handler(Arc::clone(&handler)) + .with_handler(Arc::clone(&handler)); + assert_eq!(chain.len(), 2); + + let tokens = chain.dispatch_client_created(&client_info()); + assert_eq!(*created.lock().unwrap(), 1); + assert_eq!(tokens.len(), 1); + + drop(tokens); + assert_eq!(*dropped.lock().unwrap(), 1); + } + + #[test] + fn client_created_notifies_distinct_handlers_independently() { + let created = Arc::new(Mutex::new(0)); + let dropped = Arc::new(Mutex::new(0)); + let make = || -> Arc { + Arc::new(LifecycleHandler { + created: Arc::clone(&created), + dropped: Arc::clone(&dropped), + }) + }; + + // Distinct handler objects are separate sinks, so each is notified. + let chain = DiagnosticsHandlerChain::from_handlers(vec![make(), make()]); + let tokens = chain.dispatch_client_created(&client_info()); + assert_eq!(*created.lock().unwrap(), 2); + assert_eq!(tokens.len(), 2); + + drop(tokens); + assert_eq!(*dropped.lock().unwrap(), 2); + } + + #[test] + fn client_created_is_noop_for_handlers_that_do_not_track_lifetime() { + let log = Arc::new(Mutex::new(Vec::new())); + // `RecordingHandler` does not override `on_client_created`, so the + // default returns `None` and no token is retained. + let chain = DiagnosticsHandlerChain::from_handlers(vec![recording("a", &log)]); + assert!(chain.dispatch_client_created(&client_info()).is_empty()); + } } diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/handler.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/handler.rs index 1cc92a2188a..daa30c111c9 100644 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/handler.rs +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/handler.rs @@ -417,7 +417,8 @@ mod tests { CosmosClientInfo::from_endpoint(&url::Url::parse(&url).expect("valid test endpoint")) } - /// Returns the summed value of the `active_instance.count` up-down counter /// from the most recent export, or `None` if the metric was never emitted. + /// Returns the summed value of the `active_instance.count` up-down counter + /// from the most recent export, or `None` if the metric was never emitted. /// /// The in-memory exporter accumulates one snapshot per `collect()` call, so /// we scan every snapshot and keep the value from the last one — the current diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/options.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/options.rs index adbe5ee0b78..4bb37221fa2 100644 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/options.rs +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/options.rs @@ -79,7 +79,13 @@ impl MetricsOptions { /// database/container client derived from it — has been dropped. The counter /// is keyed on the account endpoint (`server.address`, plus `server.port` /// for non-default ports), so sharing one handler across several clients - /// still reports each client. Off by default. + /// still reports each client, and registering the same handler twice on one + /// client still reports that client once. Off by default. + /// + /// Note that two *distinct* metrics handlers built from the same meter and + /// registered on the same client each record their own `+1`, just as they + /// each record their own duration and request-charge samples. Register a + /// single metrics handler per meter. #[must_use] pub fn with_active_instance_metric(mut self, enabled: bool) -> Self { self.active_instance_metric = enabled; diff --git a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md index 3aada72e79b..5860edbcf61 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md @@ -16,7 +16,6 @@ ### Bugs Fixed - `DiagnosticsContext::operation_name()` is now populated in production (previously always `None`): the operation pipeline sets it from `CosmosOperation::db_operation_name`, so tail-sampling classification and the tracing span have an operation name even when no SDK-supplied `CosmosOperationContext` is present. ([#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874)) -- The tracing span's operation label now prefers the caller-facing `CosmosOperationContext` identity, matching how the `db.operation.name` metric attribute is resolved. ([#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874)) - PATCH operations now report `patch_item` rather than the underlying Replace, on both the aggregated success path and every error path (including read, deserialize, patch-evaluation, serialize, and non-412 replace failures). ([#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874)) ### Other Changes diff --git a/sdk/cosmos/azure_data_cosmos_driver/DIAGNOSTICS-CONTRACT.md b/sdk/cosmos/azure_data_cosmos_driver/DIAGNOSTICS-CONTRACT.md index 58a2bbc58fc..c5753333080 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/DIAGNOSTICS-CONTRACT.md +++ b/sdk/cosmos/azure_data_cosmos_driver/DIAGNOSTICS-CONTRACT.md @@ -351,7 +351,7 @@ attributes. This powers client-side Grafana dashboards (R7) with per-combination | `db.client.operation.duration` | stable | histogram | `s` | End-to-end operation duration — **the primary metric**. | | `db.client.response.returned_rows` | development | histogram | `{row}` | Rows returned in the result set. | | `azure.cosmosdb.client.operation.request_charge` | development | histogram | `{request_unit}` | Request units (RU) consumed by the operation. | -| `azure.cosmosdb.client.active_instance.count` | development | up-down counter | `{instance}` | Number of live Cosmos client instances per account endpoint. *(Opt-in via `MetricsOptions::with_active_instance_metric`; `CosmosMetricsHandler` records +1 when a `CosmosClient` is built with it registered and −1 when that client and every client derived from it is dropped, so the value is independent of how many handler objects exist — [#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874).)* | +| `azure.cosmosdb.client.active_instance.count` | development | up-down counter | `{instance}` | Number of live Cosmos client instances per account endpoint. *(Opt-in via `MetricsOptions::with_active_instance_metric`; `CosmosMetricsHandler` records +1 when a `CosmosClient` is built with it registered and −1 when that client and every client derived from it is dropped, so the value follows client lifetime rather than handler lifetime — one handler shared across N clients reports N, and a handler registered on no client reports nothing — [#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874).)* | **Always-on metric attributes (low cardinality, D7):** `db.operation.name`, `db.response.status_code`, `db.collection.name`, `db.namespace`, `error.type`, `server.address`, From bc4fbf659b562419bd7ea1d89bd57242305f658d Mon Sep 17 00:00:00 2001 From: Nalu Tripician Date: Wed, 29 Jul 2026 16:16:43 -0700 Subject: [PATCH 06/10] Box plan_operation futures to satisfy large_futures After merging main, the `plan_operation` future measures 16432 bytes on Linux -- 48 bytes past clippy's 16384-byte `large_futures` threshold, which the workspace denies. Main's perf-diagnostics work and this branch's operation-name/active-instance state each grew the driver's planning future, and together they cross the limit, failing Build Analyze in `azure_data_cosmos` and `azure_data_cosmos_driver_native`. Box the five direct awaits of `plan_operation` so the future is heap allocated, matching the treatment main already applied to the change-feed tests. The size is platform-dependent (~1.2 KB smaller on Windows), and the lint only fires when clippy runs workspace-wide, so neither a Windows build nor a `-p`-scoped run reproduces it. Verified with the exact CI toolchain (rustc 1.95.0, 59807616e) and the same `cargo clippy --all-features --all-targets --keep-going --no-deps` invocation: all five errors are gone, and 424 cosmos tests pass. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7b490ca0-32f6-4909-80c7-e7f77b661c4a --- .../src/clients/container_client.rs | 34 ++++++++----------- .../src/clients/cosmos_client.rs | 17 ++++------ .../src/clients/database_client.rs | 17 ++++------ .../src/submit.rs | 10 ++++-- 4 files changed, 35 insertions(+), 43 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs b/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs index 4428837e003..ba40eabf083 100644 --- a/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs +++ b/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs @@ -889,16 +889,13 @@ impl ContainerClient { if let Some(hint) = options.feed.max_item_count { initial_operation = initial_operation.with_max_item_count(hint); } - let plan = self - .context - .driver - .plan_operation( - initial_operation, - &options.operation, - options.feed.continuation_token.as_ref(), - &options.feed.to_plan_options(), - ) - .await?; + let plan = Box::pin(self.context.driver.plan_operation( + initial_operation, + &options.operation, + options.feed.continuation_token.as_ref(), + &options.feed.to_plan_options(), + )) + .await?; Ok(QueryItemIterator::new( self.context.driver.clone(), Some(self.container_ref.clone()), @@ -1076,16 +1073,13 @@ impl ContainerClient { // precedence. The driver owns the mapping to wire headers. initial_operation = initial_operation.with_change_feed_start(start_from); - let plan = self - .context - .driver - .plan_operation( - initial_operation, - &options.operation, - options.feed.continuation_token.as_ref(), - &options.feed.to_plan_options(), - ) - .await?; + let plan = Box::pin(self.context.driver.plan_operation( + initial_operation, + &options.operation, + options.feed.continuation_token.as_ref(), + &options.feed.to_plan_options(), + )) + .await?; Ok(ChangeFeedPageIterator::new( self.context.driver.clone(), diff --git a/sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client.rs b/sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client.rs index f80626205b9..f35966c1aa6 100644 --- a/sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client.rs +++ b/sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client.rs @@ -186,16 +186,13 @@ impl CosmosClient { CosmosOperation::query_databases(account).with_body(serde_json::to_vec(&query)?); let operation_options = options.operation; - let plan = self - .context - .driver - .plan_operation( - initial_operation, - &operation_options, - None, - &PlanOptions::default(), - ) - .await?; + let plan = Box::pin(self.context.driver.plan_operation( + initial_operation, + &operation_options, + None, + &PlanOptions::default(), + )) + .await?; Ok(QueryItemIterator::new( self.context.driver.clone(), diff --git a/sdk/cosmos/azure_data_cosmos/src/clients/database_client.rs b/sdk/cosmos/azure_data_cosmos/src/clients/database_client.rs index cd5e1880614..f756ce6f22d 100644 --- a/sdk/cosmos/azure_data_cosmos/src/clients/database_client.rs +++ b/sdk/cosmos/azure_data_cosmos/src/clients/database_client.rs @@ -152,16 +152,13 @@ impl DatabaseClient { .with_body(serde_json::to_vec(&query)?); let operation_options = options.operation; - let plan = self - .context - .driver - .plan_operation( - initial_operation, - &operation_options, - None, - &PlanOptions::default(), - ) - .await?; + let plan = Box::pin(self.context.driver.plan_operation( + initial_operation, + &operation_options, + None, + &PlanOptions::default(), + )) + .await?; Ok(QueryItemIterator::new( self.context.driver.clone(), diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/src/submit.rs b/sdk/cosmos/azure_data_cosmos_driver_native/src/submit.rs index 7f3f32a3547..4fdd03e82dc 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/src/submit.rs +++ b/sdk/cosmos/azure_data_cosmos_driver_native/src/submit.rs @@ -349,9 +349,13 @@ pub extern "C" fn cosmos_submit_operation( // the continuation token through the planner and retains the // plan so we can mint the next-page token. let container = operation.container().cloned(); - let mut plan = driver_arc - .plan_operation(operation, &options, continuation.as_ref(), &plan_options) - .await?; + let mut plan = Box::pin(driver_arc.plan_operation( + operation, + &options, + continuation.as_ref(), + &plan_options, + )) + .await?; let page = driver_arc .execute_plan(&mut plan, container, options) .await?; From c05a1982bca2fccd52a67b83c377271d8fd60246 Mon Sep 17 00:00:00 2001 From: Nalu Tripician Date: Thu, 30 Jul 2026 10:56:19 -0700 Subject: [PATCH 07/10] Pin the large future inside plan_operation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses review feedback: rather than pinning at each call site, box the planning future once inside `CosmosDriver::plan_operation`. `plan_operation` holds the whole pipeline-builder state across several await points, making it ~16.4 KB — over the 16384-byte `clippy::large_futures` threshold. The previous fix pinned at the five callers, which only covered the sites that happened to exist and left future callers to rediscover the problem. The body moves to a private `plan_operation_inner`; the public `plan_operation` boxes it, so every caller now awaits a pointer-sized future. This also removes the five call-site `Box::pin` wrappers and a pre-existing redundant one inside `execute_operation`, which would otherwise double-box. Verified against the CI toolchain (rust 1.95.0, 59807616e) with the workspace-wide clippy invocation CI runs: reverting only this change reproduces all five errors at 16432 bytes, and restoring it yields zero. Note that `cargo clippy -p ` does not reproduce the lint, and neither does any Windows host. --- .../src/clients/container_client.rs | 34 +++++++++++-------- .../src/clients/cosmos_client.rs | 17 ++++++---- .../src/clients/database_client.rs | 17 ++++++---- .../src/driver/cosmos_driver.rs | 22 ++++++++++-- .../src/submit.rs | 10 ++---- 5 files changed, 62 insertions(+), 38 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs b/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs index ba40eabf083..4428837e003 100644 --- a/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs +++ b/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs @@ -889,13 +889,16 @@ impl ContainerClient { if let Some(hint) = options.feed.max_item_count { initial_operation = initial_operation.with_max_item_count(hint); } - let plan = Box::pin(self.context.driver.plan_operation( - initial_operation, - &options.operation, - options.feed.continuation_token.as_ref(), - &options.feed.to_plan_options(), - )) - .await?; + let plan = self + .context + .driver + .plan_operation( + initial_operation, + &options.operation, + options.feed.continuation_token.as_ref(), + &options.feed.to_plan_options(), + ) + .await?; Ok(QueryItemIterator::new( self.context.driver.clone(), Some(self.container_ref.clone()), @@ -1073,13 +1076,16 @@ impl ContainerClient { // precedence. The driver owns the mapping to wire headers. initial_operation = initial_operation.with_change_feed_start(start_from); - let plan = Box::pin(self.context.driver.plan_operation( - initial_operation, - &options.operation, - options.feed.continuation_token.as_ref(), - &options.feed.to_plan_options(), - )) - .await?; + let plan = self + .context + .driver + .plan_operation( + initial_operation, + &options.operation, + options.feed.continuation_token.as_ref(), + &options.feed.to_plan_options(), + ) + .await?; Ok(ChangeFeedPageIterator::new( self.context.driver.clone(), diff --git a/sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client.rs b/sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client.rs index f35966c1aa6..f80626205b9 100644 --- a/sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client.rs +++ b/sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client.rs @@ -186,13 +186,16 @@ impl CosmosClient { CosmosOperation::query_databases(account).with_body(serde_json::to_vec(&query)?); let operation_options = options.operation; - let plan = Box::pin(self.context.driver.plan_operation( - initial_operation, - &operation_options, - None, - &PlanOptions::default(), - )) - .await?; + let plan = self + .context + .driver + .plan_operation( + initial_operation, + &operation_options, + None, + &PlanOptions::default(), + ) + .await?; Ok(QueryItemIterator::new( self.context.driver.clone(), diff --git a/sdk/cosmos/azure_data_cosmos/src/clients/database_client.rs b/sdk/cosmos/azure_data_cosmos/src/clients/database_client.rs index f756ce6f22d..cd5e1880614 100644 --- a/sdk/cosmos/azure_data_cosmos/src/clients/database_client.rs +++ b/sdk/cosmos/azure_data_cosmos/src/clients/database_client.rs @@ -152,13 +152,16 @@ impl DatabaseClient { .with_body(serde_json::to_vec(&query)?); let operation_options = options.operation; - let plan = Box::pin(self.context.driver.plan_operation( - initial_operation, - &operation_options, - None, - &PlanOptions::default(), - )) - .await?; + let plan = self + .context + .driver + .plan_operation( + initial_operation, + &operation_options, + None, + &PlanOptions::default(), + ) + .await?; Ok(QueryItemIterator::new( self.context.driver.clone(), diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs index d9a1b206cd4..8a4d9b11fcf 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs @@ -2394,9 +2394,9 @@ impl CosmosDriver { // We need to do some refactoring here to shrink the future size and avoid this heap allocation if possible. let response = Box::pin(async { let container = operation.container().cloned(); - let mut plan = - Box::pin(self.plan_operation(operation, &options, None, &PlanOptions::default())) - .await?; + let mut plan = self + .plan_operation(operation, &options, None, &PlanOptions::default()) + .await?; self.execute_plan(&mut plan, container, options).await }) .await?; @@ -2958,6 +2958,22 @@ impl CosmosDriver { options: &OperationOptions, continuation: Option<&ContinuationToken>, plan_options: &PlanOptions, + ) -> crate::error::Result { + // Planning holds the whole pipeline-builder state across several await + // points, which makes it one of the largest futures in the driver — + // large enough to trip `clippy::large_futures` in callers. Box it once + // here so every caller awaits a pointer-sized future instead of having + // to pin at its own call site and rediscover this each time the state + // grows. + Box::pin(self.plan_operation_inner(operation, options, continuation, plan_options)).await + } + + async fn plan_operation_inner( + &self, + operation: CosmosOperation, + options: &OperationOptions, + continuation: Option<&ContinuationToken>, + plan_options: &PlanOptions, ) -> crate::error::Result { if !self.initialized.load(Ordering::Acquire) { let endpoint = AccountEndpoint::from(self.options.account()); diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/src/submit.rs b/sdk/cosmos/azure_data_cosmos_driver_native/src/submit.rs index 4fdd03e82dc..7f3f32a3547 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/src/submit.rs +++ b/sdk/cosmos/azure_data_cosmos_driver_native/src/submit.rs @@ -349,13 +349,9 @@ pub extern "C" fn cosmos_submit_operation( // the continuation token through the planner and retains the // plan so we can mint the next-page token. let container = operation.container().cloned(); - let mut plan = Box::pin(driver_arc.plan_operation( - operation, - &options, - continuation.as_ref(), - &plan_options, - )) - .await?; + let mut plan = driver_arc + .plan_operation(operation, &options, continuation.as_ref(), &plan_options) + .await?; let page = driver_arc .execute_plan(&mut plan, container, options) .await?; From c377939b03dd144cc811a33ece3033b2ec384435 Mon Sep 17 00:00:00 2001 From: Nalu Tripician Date: Thu, 30 Jul 2026 12:15:54 -0700 Subject: [PATCH 08/10] Name PATCH sub-operations in OTel attempts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A PATCH is one user-facing operation implemented as a read followed by a replace. We previously relabeled both halves `patch_item`, which matches the caller's mental model but erases the fact that two round trips happened. Bare `read_item` / `replace_item` would have been worse: they look like standalone point operations the caller never issued. Encode both facts instead. The operation span and the operation metric keep reporting `patch_item`; the two attempts now report `patch_read_item` and `patch_replace_item`, so the read-modify-write decomposition is visible without being mistakable for a caller request. Per-request operation identity did not exist before, so this needed a change at each layer the name has to survive: - `cosmos_operation.rs` — an `is_patch_sub_operation` marker that `db_operation_name` consults for the Read and Replace arms. - `patch_handler.rs` — the sub-operation builders set the marker; the aggregate keeps `patch_item`. - `diagnostics_context.rs` — `RequestDiagnostics` carries an optional `operation_name`. `aggregate_sub_operations` stamps it as it folds the sub-contexts together, and `preserve_request_operation_names` pushes a displaced name down when a context is relabeled, so the single sub-operation error path keeps its identity too. `None` means "same as the context", so non-PATCH diagnostics JSON is byte-identical. - `compaction.rs` — `CompactionKey` includes the operation name. Otherwise a PATCH's read and replace hitting the same endpoint with the same status could collapse into one run whose RU total and duration percentiles blend a read with a write. - `span_builder.rs` — attempt spans prefer the request's own name and fall back to the operation's. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7b490ca0-32f6-4909-80c7-e7f77b661c4a --- sdk/cosmos/azure_data_cosmos/CHANGELOG.md | 2 +- .../src/diagnostics/tracing/mod.rs | 117 ++++++ .../src/diagnostics/tracing/span_builder.rs | 7 +- .../azure_data_cosmos_driver/CHANGELOG.md | 3 +- .../src/diagnostics/compaction.rs | 10 +- .../src/diagnostics/diagnostics_context.rs | 361 +++++++++++++++++- .../src/driver/pipeline/patch_handler.rs | 37 +- .../src/models/cosmos_operation.rs | 116 ++++++ 8 files changed, 644 insertions(+), 9 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md index 6b7dfd3dd68..455cd44becf 100644 --- a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md @@ -21,7 +21,7 @@ ### Bugs Fixed -- The Cosmos tracing span's operation label now prefers the caller-facing `CosmosOperationContext` identity over the driver-recorded name, matching how the `db.operation.name` metric attribute is resolved. Previously an aggregate whose surfaced sub-operation differed from the caller's operation — such as a PATCH that fails during its internal read — could label the span `read_item` while the metric reported `patch_item`. ([#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874)) +- The Cosmos tracing span's operation label now prefers the caller-facing `CosmosOperationContext` identity over the driver-recorded name, matching how the `db.operation.name` metric attribute is resolved. Previously an aggregate whose surfaced sub-operation differed from the caller's operation — such as a PATCH that fails during its internal read — could label the span `read_item` while the metric reported `patch_item`. Attempt spans now carry the operation that issued them, so a PATCH's attempts report `db.operation.name` of `patch_read_item` / `patch_replace_item` while its operation span and metric stay `patch_item`; attempts of every other operation continue to inherit the operation's own name. ([#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874)) ### Other Changes diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/mod.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/mod.rs index 7d724c0ba50..f15992c5fa7 100644 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/mod.rs +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/mod.rs @@ -282,6 +282,123 @@ mod tests { ); } + #[test] + fn patch_sub_operations_are_visible_on_attempt_spans() { + // A PATCH is one caller-facing operation (`patch_item`) implemented as a + // read + replace. The root span keeps the caller's name, but each + // attempt span must say which half of the read-modify-write it was, so + // an operator can tell a slow/failing Read from a slow/failing Replace + // without the decomposition being flattened away. + let (provider, exporter) = exportable(); + let tracer = provider.tracer("test"); + + let now_instant = Instant::now(); + let now_system = SystemTime::now(); + let started = now_instant - Duration::from_millis(1500); + let requests = vec![ + RequestDiagnostics::for_testing( + "https://acct.documents.azure.com:443/", + Some(Region::new("West US 2")), + CosmosStatus::new(StatusCode::Ok), + RequestCharge::new(1.0), + started, + started + Duration::from_millis(500), + ) + .for_testing_with_operation_name("patch_read_item"), + RequestDiagnostics::for_testing( + "https://acct.documents.azure.com:443/", + Some(Region::new("West US 2")), + CosmosStatus::new(StatusCode::Ok), + RequestCharge::new(4.0), + started + Duration::from_millis(600), + started + Duration::from_millis(1500), + ) + .for_testing_with_operation_name("patch_replace_item"), + ]; + let ctx = DiagnosticsContext::for_testing_with_requests( + ActivityId::new_uuid(), + Duration::from_millis(1500), + Some(CosmosStatus::new(StatusCode::Ok)), + Some("patch_item"), + requests, + ); + let op = CosmosOperationContext::new().with_operation_name("patch_item"); + + emit_backdated_span_tree(&tracer, &ctx, Some(&op), None, now_instant, now_system); + provider.force_flush().unwrap(); + + let spans = exporter.get_finished_spans().unwrap(); + let root = spans + .iter() + .find(|s| s.name == "patch_item") + .expect("root span keeps the caller-facing operation name"); + assert!( + root.attributes.iter().any(|kv| { + kv.key.as_str() == attributes::DB_OPERATION_NAME + && kv.value.as_str() == "patch_item" + }), + "the operation the caller invoked is still `patch_item`" + ); + + let mut child_names: Vec = spans + .iter() + .filter(|s| s.name == "cosmosdb.request") + .filter_map(|s| { + s.attributes + .iter() + .find(|kv| kv.key.as_str() == attributes::DB_OPERATION_NAME) + .map(|kv| kv.value.as_str().to_string()) + }) + .collect(); + child_names.sort(); + assert_eq!( + child_names, + vec!["patch_read_item", "patch_replace_item"], + "each attempt span names the sub-operation that issued it" + ); + } + + #[test] + fn attempt_spans_fall_back_to_the_operation_name() { + // The non-PATCH case: attempts carry no per-request name, so every child + // inherits the operation's identity exactly as before. This is the + // overwhelmingly common path and must not regress. + let (provider, exporter) = exportable(); + let tracer = provider.tracer("test"); + + let now_instant = Instant::now(); + let now_system = SystemTime::now(); + let ctx = context( + Duration::from_millis(1500), + Some(CosmosStatus::new(StatusCode::Ok)), + Some("read_item"), + &[ + (1500, 700, CosmosStatus::new(StatusCode::TooManyRequests)), + (700, 700, CosmosStatus::new(StatusCode::Ok)), + ], + now_instant, + ); + + emit_backdated_span_tree(&tracer, &ctx, None, None, now_instant, now_system); + provider.force_flush().unwrap(); + + let spans = exporter.get_finished_spans().unwrap(); + let children: Vec<_> = spans + .iter() + .filter(|s| s.name == "cosmosdb.request") + .collect(); + assert_eq!(children.len(), 2); + for child in children { + assert!( + child.attributes.iter().any(|kv| { + kv.key.as_str() == attributes::DB_OPERATION_NAME + && kv.value.as_str() == "read_item" + }), + "an unnamed attempt inherits the operation name" + ); + } + } + #[test] fn incomplete_context_is_not_sampled_even_when_slow() { // A finalized context with neither a status nor any attempts does not diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/span_builder.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/span_builder.rs index 85673608480..aeac38709d8 100644 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/span_builder.rs +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/span_builder.rs @@ -223,6 +223,11 @@ pub(crate) fn emit_backdated_span_tree( let parent_cx = Context::current().with_remote_span_context(root.span_context().clone()); // --- Attempt (child) spans --- + // Each child prefers its own request-level operation name, which is set + // only when the operation aggregates requests from several sub-operations + // (a PATCH's `patch_read_item` / `patch_replace_item`). Everywhere else it + // is unset and the child inherits the operation's name, so a retry storm on + // a plain `read_item` still labels every attempt `read_item`. for req in requests.iter() { let child_start = to_system(req.started_at()); let child_end = child_end_of(req); @@ -238,7 +243,7 @@ pub(crate) fn emit_backdated_span_tree( ), KeyValue::new(attributes::REQUEST_CHARGE, req.request_charge().value()), ]; - if let Some(name) = op_name_ref { + if let Some(name) = req.operation_name().or(op_name_ref) { child_attrs.push(KeyValue::new( attributes::DB_OPERATION_NAME, name.to_string(), diff --git a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md index 15f2be9364a..fcff78621bd 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md @@ -8,6 +8,7 @@ - Added a schema-agnostic Cosmos binary JSON codec (`binary_json`) and driver-side binary encoding via `OperationOptions.binary_encoding` (`BinaryEncodingOptions`). When enabled, the driver transcodes item request/response bodies between text and Cosmos binary JSON and negotiates the wire format; it is honored only for point `Document` item operations. Off by default and inert on the wire when unset. ([#4671](https://github.com/Azure/azure-sdk-for-rust/pull/4671)) - Added `PlanOptions` (with `DEFAULT_MAX_FAN_OUT`) to `CosmosDriver::plan_operation`, enforcing a maximum fan-out on fresh cross-partition plans. A fresh plan spanning more leaf request nodes than `PlanOptions::max_fan_out` (default 100) is rejected with the new `CosmosStatus::CLIENT_CROSS_PARTITION_FAN_OUT_EXCEEDED` (HTTP 400). The limit is enforced only at initial plan time: resuming from a continuation token skips the check, and a partition split that raises the fan-out mid-execution does not abort the operation. ([#4855](https://github.com/Azure/azure-sdk-for-rust/pull/4855)) - Added `CosmosOperation::db_operation_name`, returning the canonical OpenTelemetry `db.operation.name` (e.g. `read_item`, `query_items`, `execute_batch`, and `read_all_items_of_logical_partition` for a read feed scoped to one logical partition) for an operation. ([#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874)) +- Added `RequestDiagnostics::operation_name`, naming the operation that issued an individual attempt. It is set only where one `DiagnosticsContext` aggregates attempts from more than one operation — today a PATCH, whose attempts report `patch_read_item` / `patch_replace_item` while the context reports `patch_item` — and is `None` otherwise, meaning the attempt shares the context's operation name. `CosmosOperation::is_patch_sub_operation` reports the same distinction on the operation itself. ([#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874)) ### Breaking Changes @@ -16,7 +17,7 @@ ### Bugs Fixed - `DiagnosticsContext::operation_name()` is now populated in production (previously always `None`): the operation pipeline sets it from `CosmosOperation::db_operation_name`, so tail-sampling classification and the tracing span have an operation name even when no SDK-supplied `CosmosOperationContext` is present. ([#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874)) -- PATCH operations now report `patch_item` rather than the underlying Replace, on both the aggregated success path and every error path (including read, deserialize, patch-evaluation, serialize, and non-412 replace failures). ([#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874)) +- PATCH operations now report `patch_item` rather than the underlying Replace, on both the aggregated success path and every error path (including read, deserialize, patch-evaluation, serialize, and non-412 replace failures). The two internal sub-operations report `patch_read_item` and `patch_replace_item` on their own attempt diagnostics, so the read-modify-write decomposition stays visible underneath the caller-facing operation instead of being flattened to a single name. ([#4874](https://github.com/Azure/azure-sdk-for-rust/pull/4874)) ### Other Changes diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/compaction.rs b/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/compaction.rs index dddab7c698c..3e91d238937 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/compaction.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/compaction.rs @@ -123,13 +123,20 @@ pub struct CompactedRun { } /// The key that defines a "near-identical" run: same region, endpoint, status -/// (incl. sub-status) and execution context. +/// (incl. sub-status), execution context, and issuing operation. +/// +/// The operation is part of the key because a run is meant to be a storm of +/// retries of *the same* attempt. A PATCH's `patch_read_item` and +/// `patch_replace_item` attempts can otherwise land on the same endpoint with +/// the same status and be rolled up together, producing a run whose RU total +/// and duration percentiles silently mix a read with a write. #[derive(Clone, PartialEq, Eq, Hash)] struct CompactionKey { region: Option, endpoint: String, status: CosmosStatus, execution_context: ExecutionContext, + operation_name: Option, } impl CompactionKey { @@ -139,6 +146,7 @@ impl CompactionKey { endpoint: req.endpoint().to_string(), status: *req.status(), execution_context: req.execution_context(), + operation_name: req.operation_name().map(str::to_string), } } } diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/diagnostics_context.rs b/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/diagnostics_context.rs index 39505657c68..423b27bd4c6 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/diagnostics_context.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/diagnostics_context.rs @@ -397,6 +397,20 @@ pub struct RequestDiagnostics { /// Context describing why this request was made. execution_context: ExecutionContext, + /// Canonical `db.operation.name` of the operation that issued this request. + /// + /// Normally redundant with the owning [`DiagnosticsContext`]'s + /// `operation_name`, and therefore left unset. It is populated when an + /// aggregate context spans requests from more than one operation, so the + /// per-request identity is not lost to the aggregate's single name — today + /// that means a PATCH, whose requests are the `patch_read_item` and + /// `patch_replace_item` sub-ops of one caller-facing `patch_item`. + #[serde( + skip_serializing_if = "Option::is_none", + serialize_with = "serialize_optional_shared_str" + )] + operation_name: Option>, + /// The pipeline type used for this request. pipeline_type: PipelineType, @@ -493,6 +507,7 @@ impl RequestDiagnostics { ) -> Self { Self { execution_context, + operation_name: None, pipeline_type, transport_security, transport_kind, @@ -542,6 +557,7 @@ impl RequestDiagnostics { .as_millis() as u64; Self { execution_context: ExecutionContext::Initial, + operation_name: None, pipeline_type: PipelineType::DataPlane, transport_security: TransportSecurity::Secure, transport_kind: TransportKind::Gateway, @@ -568,6 +584,20 @@ impl RequestDiagnostics { } } + /// **Internal test helper — do not call.** + /// + /// Stamps this attempt with the sub-operation that issued it, mirroring what + /// [`DiagnosticsContext::aggregate_sub_operations`] does in production. Lets + /// emission-layer tests exercise per-request naming without reaching into + /// the driver's crate-private aggregation path. + #[cfg(feature = "__internal_test_diagnostics_construction")] + #[doc(hidden)] + #[must_use] + pub fn for_testing_with_operation_name(mut self, operation_name: impl Into>) -> Self { + self.operation_name = Some(operation_name.into()); + self + } + /// Records completion of this request. /// /// Since we received a response, the request was definitely sent. @@ -693,6 +723,19 @@ impl RequestDiagnostics { self.execution_context } + /// Returns the canonical `db.operation.name` of the operation that issued + /// this request, when it differs from the owning context's operation name. + /// + /// This is set only where a single [`DiagnosticsContext`] aggregates + /// requests from more than one operation. Today that is the PATCH handler: + /// the context reports the caller-facing `patch_item` while its requests + /// report the `patch_read_item` / `patch_replace_item` sub-op that produced + /// them. `None` — the common case — means the request shares the owning + /// context's [`operation_name`](DiagnosticsContext::operation_name). + pub fn operation_name(&self) -> Option<&str> { + self.operation_name.as_deref() + } + /// Returns the pipeline type used for this request. pub fn pipeline_type(&self) -> PipelineType { self.pipeline_type @@ -1050,6 +1093,24 @@ fn is_zero_u32(value: &u32) -> bool { *value == 0 } +/// Serializes an `Option>` as a plain optional string. +/// +/// `serde` only implements `Serialize` for `Arc` under its `rc` feature, +/// which this crate does not enable, so the shared string is written through +/// its `str` view instead. +fn serialize_optional_shared_str( + value: &Option>, + serializer: S, +) -> Result +where + S: serde::Serializer, +{ + match value { + Some(value) => serializer.serialize_str(value), + None => serializer.serialize_none(), + } +} + impl RequestEvent { /// Creates a new request event. pub fn new(event_type: RequestEventType) -> Self { @@ -1998,9 +2059,24 @@ impl DiagnosticsContext { /// Returns `None` only when `sources` is empty. pub(crate) fn aggregate_sub_operations(sources: &[Arc]) -> Option { let last = sources.last()?; + // Carry each source's operation name down onto the requests it + // contributed. The aggregate reports a single operation name (the + // caller-facing one — `patch_item`), so without this the sub-op + // identity would be lost the moment the contexts are concatenated and + // every attempt span would inherit the aggregate's name. Sources that + // are themselves aggregates may already carry per-request names; those + // are preserved rather than overwritten. let aggregated_requests: Vec = sources .iter() - .flat_map(|c| c.requests.iter().cloned()) + .flat_map(|c| { + c.requests.iter().map(|req| { + let mut req = req.clone(); + if req.operation_name.is_none() { + req.operation_name = c.operation_name.clone(); + } + req + }) + }) .collect(); let aggregated_duration = sources .iter() @@ -2256,10 +2332,56 @@ impl DiagnosticsContext { /// operation's own name (`patch_item`). Consumes `self` before it is shared /// via `Arc`, preserving the type's immutability contract. pub(crate) fn with_operation_name(mut self, operation_name: Option>) -> Self { + self.requests = Self::preserve_request_operation_names( + &self.requests, + self.operation_name.as_ref(), + operation_name.as_ref(), + ); self.operation_name = operation_name; self } + /// Pushes a context-level operation name down onto the requests that were + /// issued under it, so relabeling the context does not erase where its + /// requests came from. + /// + /// Relabeling happens when a virtual operation is assembled from real + /// sub-operations: a PATCH context is stamped `patch_item`, but its + /// requests were issued by the `patch_read_item` / `patch_replace_item` + /// sub-ops. Without this, the attempt-level view would report the + /// aggregate's name for every request and the read/modify/write + /// decomposition would be invisible. + /// + /// Requests that already carry their own name keep it (they came from a + /// context that was itself an aggregate). When the name is unchanged, or + /// there is no displaced name to record, the existing `Arc` is shared + /// rather than the request list being cloned. + fn preserve_request_operation_names( + requests: &Arc>, + previous: Option<&Arc>, + replacement: Option<&Arc>, + ) -> Arc> { + let Some(previous) = previous else { + return Arc::clone(requests); + }; + if replacement.is_some_and(|new| new == previous) { + return Arc::clone(requests); + } + if requests.iter().all(|req| req.operation_name.is_some()) { + return Arc::clone(requests); + } + Arc::new( + requests + .iter() + .map(|req| { + let mut req = req.clone(); + req.operation_name.get_or_insert_with(|| previous.clone()); + req + }) + .collect(), + ) + } + /// Returns a copy of this context with its canonical `db.operation.name` /// replaced, leaving every other field — including status, hedging /// diagnostics, and compaction metadata — intact. @@ -2275,7 +2397,11 @@ impl DiagnosticsContext { DiagnosticsContext { activity_id: self.activity_id.clone(), duration: self.duration, - requests: Arc::clone(&self.requests), + requests: Self::preserve_request_operation_names( + &self.requests, + self.operation_name.as_ref(), + operation_name.as_ref(), + ), total_request_charge: self.total_request_charge, regions_contacted: self.regions_contacted.clone(), status: self.status, @@ -4237,6 +4363,85 @@ mod tests { } } + #[test] + fn compaction_does_not_collapse_distinct_patch_sub_operations() { + // A run is meant to be a storm of retries of *the same* attempt. A + // PATCH's Read and Replace hit the same endpoint and can return the + // same status, so without the issuing operation in the compaction key + // the aggregate's re-bounding pass collapses them into one run whose RU + // total and duration percentiles silently mix a read with a write. + // + // Sized so neither sub-op compacts on its own (9 < cap) but their + // concatenation does (18 > cap), which is the only path where a single + // compaction pass ever sees requests from more than one operation. + let cap = 16; + let per_sub_op = 9; + let mut read_b = DiagnosticsContextBuilder::new( + ActivityId::from_string("patch-read".to_string()), + options_with_cap(cap), + ); + record_run( + &mut read_b, + ExecutionContext::Retry, + "East US", + "https://east/", + CosmosStatus::new(StatusCode::Ok), + 1.0, + per_sub_op, + ); + read_b.set_operation_name("patch_read_item"); + read_b.set_operation_status(StatusCode::Ok, None); + let read_ctx = Arc::new(read_b.complete()); + assert!( + read_ctx.compaction().is_none(), + "sub-op must be under the cap so the aggregate does the compacting" + ); + + let mut replace_b = DiagnosticsContextBuilder::new( + ActivityId::from_string("patch-replace".to_string()), + options_with_cap(cap), + ); + record_run( + &mut replace_b, + ExecutionContext::Retry, + "East US", + "https://east/", + CosmosStatus::new(StatusCode::Ok), + 10.0, + per_sub_op, + ); + replace_b.set_operation_name("patch_replace_item"); + replace_b.set_operation_status(StatusCode::Ok, None); + let replace_ctx = Arc::new(replace_b.complete()); + assert!(replace_ctx.compaction().is_none()); + + let aggregated = DiagnosticsContext::aggregate_sub_operations(&[read_ctx, replace_ctx]) + .expect("aggregation of two contexts yields Some") + .with_operation_name(Some(Arc::from("patch_item"))); + + let info = aggregated + .compaction() + .expect("18 concatenated attempts past a cap of 16 must compact"); + assert_eq!( + info.runs.len(), + 2, + "the read and the replace must be reported as separate runs" + ); + + // Each run's RU total reflects one sub-op, not a blend of both. + let mut charges: Vec = info + .runs + .iter() + .map(|r| r.total_request_charge.value()) + .collect(); + charges.sort_by(f64::total_cmp); + assert_eq!( + charges, + vec![9.0, 90.0], + "runs must not blend the read's 1 RU attempts with the replace's 10 RU attempts" + ); + } + #[test] fn retry_storm_429_is_bounded_and_lossless() { // A single partition hammered with 429 for the whole retry budget: one @@ -4666,6 +4871,158 @@ mod tests { assert_eq!(stamped.operation_name(), Some("patch_item")); } + #[test] + fn aggregate_sub_operations_preserves_per_request_operation_names() { + // A PATCH reports `patch_item` at the operation level, but its requests + // were issued by the `patch_read_item` / `patch_replace_item` sub-ops. + // Aggregation must push each source's name down onto the requests it + // contributed, or the decomposition is lost the moment the contexts are + // concatenated. + let read_ctx = Arc::new(make_context_with(ActivityId::new_uuid(), |builder| { + builder.start_test_request( + ExecutionContext::Initial, + Some(Region::WEST_US_2), + "https://test.westus2.documents.azure.com", + ); + builder.set_operation_name("patch_read_item"); + builder.set_operation_status(StatusCode::Ok, None); + })); + let replace_ctx = Arc::new(make_context_with(ActivityId::new_uuid(), |builder| { + builder.start_test_request( + ExecutionContext::Initial, + Some(Region::WEST_US_2), + "https://test.westus2.documents.azure.com", + ); + builder.set_operation_name("patch_replace_item"); + builder.set_operation_status(StatusCode::Ok, None); + })); + + let aggregated = DiagnosticsContext::aggregate_sub_operations(&[read_ctx, replace_ctx]) + .expect("aggregation of two contexts yields Some") + .with_operation_name(Some(Arc::from("patch_item"))); + + assert_eq!(aggregated.operation_name(), Some("patch_item")); + let requests = aggregated.requests(); + let names: Vec> = requests + .iter() + .map(RequestDiagnostics::operation_name) + .collect(); + assert_eq!( + names, + vec![Some("patch_read_item"), Some("patch_replace_item")], + "each request must keep the sub-op that issued it" + ); + } + + #[test] + fn single_operation_requests_carry_no_redundant_name() { + // The overwhelmingly common case: one operation, N attempts. The + // per-request name stays unset so it costs nothing and the diagnostics + // JSON is unchanged; consumers fall back to the context's name. + let ctx = make_context_with(ActivityId::new_uuid(), |builder| { + builder.start_test_request( + ExecutionContext::Initial, + Some(Region::WEST_US_2), + "https://test.westus2.documents.azure.com", + ); + builder.start_test_request( + ExecutionContext::Retry, + Some(Region::WEST_US_2), + "https://test.westus2.documents.azure.com", + ); + builder.set_operation_name("read_item"); + builder.set_operation_status(StatusCode::Ok, None); + }); + + assert_eq!(ctx.operation_name(), Some("read_item")); + assert!( + ctx.requests() + .iter() + .all(|req| req.operation_name().is_none()), + "a single-operation context must not duplicate its name onto every request" + ); + + // Re-stamping with the *same* name is a no-op rather than a reason to + // populate every request. + let restamped = ctx.clone_with_operation_name(Some(Arc::from("read_item"))); + assert!(restamped + .requests() + .iter() + .all(|req| req.operation_name().is_none())); + } + + #[test] + fn clone_with_operation_name_preserves_displaced_request_identity() { + // The PATCH error path re-stamps a single sub-op context (it does not + // aggregate when only one context exists). The requests must still + // remember they came from the Read, otherwise a PATCH that fails during + // its internal Read reports `patch_item` on the attempt span and the + // read/replace split disappears exactly when it matters most. + let read_ctx = make_context_with(ActivityId::new_uuid(), |builder| { + builder.start_test_request( + ExecutionContext::Initial, + Some(Region::WEST_US_2), + "https://test.westus2.documents.azure.com", + ); + builder.set_operation_name("patch_read_item"); + builder.set_operation_status(StatusCode::NotFound, None); + }); + + let stamped = read_ctx.clone_with_operation_name(Some(Arc::from("patch_item"))); + + assert_eq!(stamped.operation_name(), Some("patch_item")); + assert_eq!( + stamped.requests()[0].operation_name(), + Some("patch_read_item") + ); + // The source context is untouched. + assert_eq!(read_ctx.requests()[0].operation_name(), None); + } + + #[test] + fn nested_aggregation_keeps_the_innermost_request_identity() { + // Aggregating an aggregate (a PATCH whose sub-ops were themselves + // aggregated) must not overwrite names that are already more specific + // than the enclosing context's. + let read_ctx = Arc::new(make_context_with(ActivityId::new_uuid(), |builder| { + builder.start_test_request( + ExecutionContext::Initial, + Some(Region::WEST_US_2), + "https://test.westus2.documents.azure.com", + ); + builder.set_operation_name("patch_read_item"); + builder.set_operation_status(StatusCode::Ok, None); + })); + let replace_ctx = Arc::new(make_context_with(ActivityId::new_uuid(), |builder| { + builder.start_test_request( + ExecutionContext::Initial, + Some(Region::WEST_US_2), + "https://test.westus2.documents.azure.com", + ); + builder.set_operation_name("patch_replace_item"); + builder.set_operation_status(StatusCode::Ok, None); + })); + + let inner = Arc::new( + DiagnosticsContext::aggregate_sub_operations(&[read_ctx, replace_ctx]) + .expect("aggregation of two contexts yields Some") + .with_operation_name(Some(Arc::from("patch_item"))), + ); + let outer = DiagnosticsContext::aggregate_sub_operations(&[inner]) + .expect("aggregation of one context yields Some") + .with_operation_name(Some(Arc::from("patch_item"))); + + let requests = outer.requests(); + let names: Vec> = requests + .iter() + .map(RequestDiagnostics::operation_name) + .collect(); + assert_eq!( + names, + vec![Some("patch_read_item"), Some("patch_replace_item")] + ); + } + #[test] fn is_failure_reflects_operation_status() { let ok = make_context_with(ActivityId::new_uuid(), |b| { diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/patch_handler.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/patch_handler.rs index 9bf90176828..b2fd5d15122 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/patch_handler.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/patch_handler.rs @@ -208,8 +208,12 @@ pub(crate) async fn execute_with_dispatcher( Vec::with_capacity(2 * attempts as usize); // The aggregated context concatenates the Read + Replace sub-ops and would - // otherwise inherit the *last* sub-op's `db.operation.name` (`replace_item`). - // Stamp the virtual PATCH operation's own canonical name instead. + // otherwise inherit the *last* sub-op's `db.operation.name`. Stamp the + // virtual PATCH operation's own canonical name (`patch_item`) instead, so + // the operation level reports what the caller actually invoked. The + // individual sub-ops keep their own `patch_read_item` / `patch_replace_item` + // identity on their per-request diagnostics, so the read/modify/write + // decomposition stays visible underneath the aggregate. let operation_name: Option> = operation.db_operation_name().map(Arc::from); for _ in 0..attempts { @@ -582,7 +586,7 @@ fn build_read_sub_op( item_ref: crate::models::ItemReference, caller_session_token: Option, ) -> CosmosOperation { - let mut op = CosmosOperation::read_item(item_ref); + let mut op = CosmosOperation::read_item(item_ref).as_patch_sub_operation(); if let Some(token) = caller_session_token { op = op.with_session_token(token); } @@ -600,6 +604,7 @@ fn build_replace_sub_op( read_response_session_token: Option, ) -> CosmosOperation { let mut op = CosmosOperation::replace_item(item_ref) + .as_patch_sub_operation() .with_body(merged_bytes) .with_precondition(Precondition::if_match(etag)); if let Some(token) = read_response_session_token { @@ -891,6 +896,32 @@ mod tests { assert!(op.request_headers().session_token.is_none()); } + #[test] + fn sub_ops_report_patch_scoped_operation_names() { + // The RMW sub-ops are dispatched exactly like standalone point + // operations, so without the marker their telemetry would be + // indistinguishable from a `read_item` / `replace_item` the caller + // issued directly. The `patch_` prefix keeps them attributable to the + // PATCH while still naming which half of the read-modify-write they + // are. + let read = build_read_sub_op(test_item_ref(), None); + assert!(read.is_patch_sub_operation()); + assert_eq!(read.db_operation_name(), Some("patch_read_item")); + + let replace = build_replace_sub_op( + test_item_ref(), + b"{\"id\":\"doc1\"}".to_vec(), + Etag::from("\"abc\""), + None, + ); + assert!(replace.is_patch_sub_operation()); + assert_eq!(replace.db_operation_name(), Some("patch_replace_item")); + + // The caller-facing operation keeps its own name; that is what the + // aggregate context, root span, and operation metric report. + assert_eq!(canonical_patch_op().db_operation_name(), Some("patch_item")); + } + #[test] fn is_precondition_failed_matches_real_412() { // the RMW loop's 412 detection runs on the `Err(_)` produced diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_operation.rs b/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_operation.rs index 96ddc4119aa..281b2b41129 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_operation.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_operation.rs @@ -152,6 +152,13 @@ pub struct CosmosOperation { /// token so never-polled partitions can re-apply it on resume. `None` for /// non-change-feed operations. change_feed_start: Option, + /// `true` when this operation is one of the internal sub-operations the + /// PATCH handler's Read-Modify-Write loop dispatches, rather than an + /// operation the caller requested directly. Set by + /// [`as_patch_sub_operation`](Self::as_patch_sub_operation); it only + /// affects [`db_operation_name`](Self::db_operation_name), so the sub-op + /// is dispatched exactly like the standalone Read/Replace it is. + is_patch_sub_operation: bool, } impl CosmosOperation { @@ -185,11 +192,30 @@ impl CosmosOperation { /// but an offer operation carries only the account and the offer ID, so the /// scope is not recoverable here. The SDK, which knows whether the caller /// addressed a container or a database, supplies those names instead. + /// + /// # PATCH sub-operations + /// + /// PATCH is a single caller-facing operation that this driver implements as + /// a Read followed by an ETag-guarded Replace. The two sub-operations report + /// `patch_read_item` and `patch_replace_item` rather than the bare + /// `read_item` / `replace_item`, so telemetry encodes *both* facts: that the + /// work belongs to a PATCH, and which half of the read-modify-write it is. + /// Naming them `read_item`/`replace_item` would make them indistinguishable + /// from standalone point operations the caller never issued; naming them + /// `patch_item` would hide the decomposition entirely. The operation the + /// caller actually invoked keeps reporting `patch_item` on the root span and + /// the operation metric. pub fn db_operation_name(&self) -> Option<&'static str> { let name = match (self.operation_type, self.resource_type) { // Data-plane item operations. (OperationType::Create, ResourceType::Document) => "create_item", + (OperationType::Read, ResourceType::Document) if self.is_patch_sub_operation => { + "patch_read_item" + } (OperationType::Read, ResourceType::Document) => "read_item", + (OperationType::Replace, ResourceType::Document) if self.is_patch_sub_operation => { + "patch_replace_item" + } (OperationType::Replace, ResourceType::Document) => "replace_item", (OperationType::Delete, ResourceType::Document) => "delete_item", (OperationType::Upsert, ResourceType::Document) => "upsert_item", @@ -438,6 +464,24 @@ impl CosmosOperation { self.patch_max_attempts } + /// Marks this operation as an internal sub-operation of a PATCH's + /// Read-Modify-Write loop. + /// + /// The only effect is on [`db_operation_name`](Self::db_operation_name), + /// which then reports `patch_read_item` / `patch_replace_item` instead of + /// `read_item` / `replace_item`. Routing, retries, and the wire request are + /// unchanged — a PATCH sub-op *is* an ordinary point Read or Replace. + pub(crate) fn as_patch_sub_operation(mut self) -> Self { + self.is_patch_sub_operation = true; + self + } + + /// Returns `true` when this operation is an internal sub-operation of a + /// PATCH's Read-Modify-Write loop. + pub fn is_patch_sub_operation(&self) -> bool { + self.is_patch_sub_operation + } + // ===== Factory Methods ===== /// Creates a new operation with the specified type, resource reference, and target. @@ -464,6 +508,7 @@ impl CosmosOperation { patch_max_attempts: None, is_change_feed: false, change_feed_start: None, + is_patch_sub_operation: false, } } @@ -1275,6 +1320,77 @@ mod tests { ); } + #[test] + fn db_operation_name_distinguishes_patch_sub_operations() { + let item = + || ItemReference::from_name(&test_container(), PartitionKey::from("pk1"), "doc1"); + + // A PATCH is one caller-facing operation implemented as a Read plus an + // ETag-guarded Replace. The sub-ops report names that encode both the + // owning PATCH and which half of the read-modify-write they are, so + // telemetry neither hides the decomposition nor makes the sub-ops look + // like standalone point operations the caller never issued. + assert_eq!( + CosmosOperation::read_item(item()) + .as_patch_sub_operation() + .db_operation_name(), + Some("patch_read_item") + ); + assert_eq!( + CosmosOperation::replace_item(item()) + .as_patch_sub_operation() + .db_operation_name(), + Some("patch_replace_item") + ); + + // The operation the caller actually invoked is unaffected. + assert_eq!( + CosmosOperation::patch_item(item()).db_operation_name(), + Some("patch_item") + ); + assert!(!CosmosOperation::patch_item(item()).is_patch_sub_operation()); + } + + #[test] + fn patch_sub_operation_marker_is_off_by_default() { + let item = + || ItemReference::from_name(&test_container(), PartitionKey::from("pk1"), "doc1"); + + assert!(!CosmosOperation::read_item(item()).is_patch_sub_operation()); + assert!(!CosmosOperation::replace_item(item()).is_patch_sub_operation()); + assert!(CosmosOperation::read_item(item()) + .as_patch_sub_operation() + .is_patch_sub_operation()); + } + + #[test] + fn patch_sub_operation_marker_only_renames_read_and_replace() { + let item = + || ItemReference::from_name(&test_container(), PartitionKey::from("pk1"), "doc1"); + + // The marker is only ever set on the two sub-ops the PATCH handler + // dispatches. Guard the mapping anyway so a stray marker on any other + // operation cannot silently invent a name. + assert_eq!( + CosmosOperation::create_item(item()) + .as_patch_sub_operation() + .db_operation_name(), + Some("create_item") + ); + assert_eq!( + CosmosOperation::upsert_item(item()) + .as_patch_sub_operation() + .db_operation_name(), + Some("upsert_item") + ); + assert_eq!( + CosmosOperation::delete_item(item()) + .as_patch_sub_operation() + .db_operation_name(), + Some("delete_item") + ); + } + #[test] fn db_operation_name_maps_feed_and_query_operations() { assert_eq!( From 4b9a09e736c2569b2ac0d6e712e611dd86450e3c Mon Sep 17 00:00:00 2001 From: Nalu Tripician Date: Fri, 31 Jul 2026 11:58:09 -0700 Subject: [PATCH 09/10] Preserve patch span order Keep the tracing assertion aligned with the actual read-modify-write execution sequence instead of sorting away ordering regressions. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 5440044f-ae66-4eee-8f44-f06bf90a8e8b --- sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/mod.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/mod.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/mod.rs index f15992c5fa7..1e47437f229 100644 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/mod.rs +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/mod.rs @@ -340,7 +340,7 @@ mod tests { "the operation the caller invoked is still `patch_item`" ); - let mut child_names: Vec = spans + let child_names: Vec = spans .iter() .filter(|s| s.name == "cosmosdb.request") .filter_map(|s| { @@ -350,7 +350,6 @@ mod tests { .map(|kv| kv.value.as_str().to_string()) }) .collect(); - child_names.sort(); assert_eq!( child_names, vec!["patch_read_item", "patch_replace_item"], From 216e038d91786fa16c4de75d1bce1bc1f4643dd8 Mon Sep 17 00:00:00 2001 From: Nalu Tripician Date: Fri, 31 Jul 2026 12:05:02 -0700 Subject: [PATCH 10/10] Complete PR review follow-up Record completion of the reviewer nit follow-up after verifying the existing behavioral commit. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 49ca1a2f-4e01-44a7-a1a5-959602b42935