diff --git a/Cargo.lock b/Cargo.lock index 5a906b143e7..eab3854d1c1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -527,6 +527,8 @@ dependencies = [ "base64", "clap", "futures", + "opentelemetry 0.32.0", + "opentelemetry_sdk 0.32.1", "pin-project", "reqwest", "serde", diff --git a/sdk/cosmos/.cspell.json b/sdk/cosmos/.cspell.json index ab818543976..f2625a4d978 100644 --- a/sdk/cosmos/.cspell.json +++ b/sdk/cosmos/.cspell.json @@ -117,6 +117,7 @@ "firewalled", "fixdate", "flamegraph", + "flatbuffer", "fmix", "fract", "francecentral", @@ -175,6 +176,7 @@ "MEMORYSTATUSEX", "Mgmt", "malloc", + "materializer", "mmap", "moka", "mpmc", @@ -220,6 +222,7 @@ "pgcosmos", "PIDS", "Pigw", + "pingpong", "pinning", "Pkrange", "pkranges", @@ -256,12 +259,14 @@ "rhost", "RNTBD", "rgby", + "rollup", "roundtripped", "roundtrips", "RTLD", "RUPM", "rwcache", "sbyte", + "semconv", "serviceunavailable", "serviceversion", "sess", @@ -307,6 +312,7 @@ "unavail", "uncollapsed", "uncontended", + "underspecified", "undrained", "unfaulted", "ungoverned", @@ -345,6 +351,8 @@ "westus", "widenings", "winget", + "workstream", + "workstreams", "writability", "xorshift", "xpart", diff --git a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md index 38a3496ffb5..19bbda6ae61 100644 --- a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md @@ -4,6 +4,10 @@ ### 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 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 ### Bugs Fixed diff --git a/sdk/cosmos/azure_data_cosmos/Cargo.toml b/sdk/cosmos/azure_data_cosmos/Cargo.toml index ae71ef1e03c..d67a20f1061 100644 --- a/sdk/cosmos/azure_data_cosmos/Cargo.toml +++ b/sdk/cosmos/azure_data_cosmos/Cargo.toml @@ -20,6 +20,7 @@ azure_core = { workspace = true, default-features = false } azure_data_cosmos_driver = { version = "0.7.0", path = "../azure_data_cosmos_driver", default-features = false } base64.workspace = true futures.workspace = true +opentelemetry = { workspace = true, optional = true, features = ["metrics"] } pin-project.workspace = true reqwest = { workspace = true, optional = true } serde.workspace = true @@ -51,6 +52,10 @@ azure_data_cosmos_driver = { path = "../azure_data_cosmos_driver", default-featu ] } azure_identity.workspace = true clap.workspace = true +# Enabled here (dev-only) so the metrics and tracing handler tests can drive +# in-memory OpenTelemetry exporters. `testing` provides `InMemoryMetricExporter` +# and `InMemorySpanExporter`. +opentelemetry_sdk = { workspace = true, features = ["testing"] } reqwest = { workspace = true, features = ["http2"] } serde = { workspace = true, features = ["derive"] } serde_bytes.workspace = true @@ -79,11 +84,17 @@ rustls = ["reqwest", "azure_data_cosmos_driver/rustls", "__tls"] native_tls = ["reqwest", "azure_data_cosmos_driver/native_tls", "__tls"] key_auth = [ ] # Enables support for key-based authentication (Primary Keys and Resource Tokens) +distributed_tracing = [ + "dep:opentelemetry", +] # Enables the OpenTelemetry distributed-tracing diagnostics handler (off by default) hmac_rust = ["azure_core/hmac_rust"] hmac_openssl = ["azure_core/hmac_openssl"] fault_injection = [ "azure_data_cosmos_driver/fault_injection", ] # Enables support for fault injection testing +metrics = [ + "dep:opentelemetry", +] # Enables the OpenTelemetry metrics handler (`CosmosMetricsHandler`). Off by default: with the feature disabled the metrics module is not compiled, so there is zero cost. preview_dtx = [ "azure_data_cosmos_driver/preview_dtx", ] # Enables preview Distributed Transaction APIs. Disabled by default and not production-ready. @@ -95,10 +106,12 @@ __tls = [] [package.metadata.docs.rs] features = [ + "distributed_tracing", "fault_injection", "hmac_openssl", "hmac_rust", "key_auth", + "metrics", "native_tls", "rustls", ] 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 a69327a7543..ea9e56b496e 100644 --- a/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs +++ b/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs @@ -3,6 +3,7 @@ use crate::{ clients::{offers_client, ClientContext}, + diagnostics::CosmosOperationContext, feed::{ChangeFeedPageIterator, FeedRange, FeedScope, QueryItemIterator}, models::TransactionalBatch, models::{BatchResponse, ChangeFeedItem, ItemResponse, ResourceResponse}, @@ -62,6 +63,16 @@ impl ContainerClient { }) } + /// Builds the SDK-side [`CosmosOperationContext`] for this container's + /// operations, carrying the operation name plus the database and container + /// identity the driver context does not know. + fn operation_context(&self, operation_name: &'static str) -> CosmosOperationContext { + CosmosOperationContext::new() + .with_operation_name(operation_name) + .with_database_name(self.container_ref.database_name().to_string()) + .with_container_name(self.container_ref.name().to_string()) + } + /// Reads the properties of the container. /// /// # Arguments @@ -85,14 +96,15 @@ impl ContainerClient { let options = options.unwrap_or_default(); let operation = CosmosOperation::read_container(self.container_ref.clone()); - let driver_response = self + let driver_result = self .context .driver .execute_singleton_operation(operation, options.operation) - .await?; + .await; Ok(ResourceResponse::new( - crate::driver_bridge::driver_response_to_cosmos_response(driver_response), + self.context + .complete_result(driver_result, || self.operation_context("read_container"))?, )) } @@ -138,14 +150,16 @@ impl ContainerClient { operation_options.content_response_on_write = Some(azure_data_cosmos_driver::options::ContentResponseOnWrite::Enabled); - let driver_response = self + let driver_result = self .context .driver .execute_singleton_operation(operation, operation_options) - .await?; + .await; Ok(ResourceResponse::new( - crate::driver_bridge::driver_response_to_cosmos_response(driver_response), + self.context.complete_result(driver_result, || { + self.operation_context("replace_container") + })?, )) } @@ -161,10 +175,11 @@ impl ContainerClient { ) -> crate::Result> { let options = options.unwrap_or_default(); offers_client::find_offer( - &self.context.driver, + &self.context, self.container_ref.account(), self.container_ref.rid(), options.operation, + self.operation_context("read_throughput"), ) .await } @@ -202,11 +217,12 @@ impl ContainerClient { let options = options.unwrap_or_default(); offers_client::begin_replace( - self.context.driver.clone(), + self.context.clone(), self.container_ref.account().clone(), self.container_ref.rid(), throughput, options.operation, + self.operation_context("replace_throughput"), ) .await } @@ -224,14 +240,15 @@ impl ContainerClient { let options = options.unwrap_or_default(); let operation = CosmosOperation::delete_container(self.container_ref.clone()); - let driver_response = self + let driver_result = self .context .driver .execute_singleton_operation(operation, options.operation) - .await?; + .await; Ok(ResourceResponse::new( - crate::driver_bridge::driver_response_to_cosmos_response(driver_response), + self.context + .complete_result(driver_result, || self.operation_context("delete_container"))?, )) } @@ -322,15 +339,16 @@ impl ContainerClient { let operation = apply_item_options(operation, options.session_token, options.precondition); // Execute through the driver. - let driver_response = self + let driver_result = self .context .driver .execute_singleton_operation(operation, options.operation) - .await?; + .await; // Bridge the driver response to the SDK response type. Ok(ItemResponse::new( - crate::driver_bridge::driver_response_to_cosmos_response(driver_response), + self.context + .complete_result(driver_result, || self.operation_context("create_item"))?, )) } @@ -420,15 +438,16 @@ impl ContainerClient { let operation = apply_item_options(operation, options.session_token, options.precondition); // Execute through the driver. - let driver_response = self + let driver_result = self .context .driver .execute_singleton_operation(operation, options.operation) - .await?; + .await; // Bridge the driver response to the SDK response type. Ok(ItemResponse::new( - crate::driver_bridge::driver_response_to_cosmos_response(driver_response), + self.context + .complete_result(driver_result, || self.operation_context("replace_item"))?, )) } @@ -527,14 +546,15 @@ impl ContainerClient { // session token. let operation = apply_item_options(operation, options.session_token, None); - let driver_response = self + let driver_result = self .context .driver .execute_singleton_operation(operation, options.operation) - .await?; + .await; Ok(ItemResponse::new( - crate::driver_bridge::driver_response_to_cosmos_response(driver_response), + self.context + .complete_result(driver_result, || self.operation_context("patch_item"))?, )) } @@ -628,15 +648,16 @@ impl ContainerClient { let operation = apply_item_options(operation, options.session_token, options.precondition); // Execute through the driver. - let driver_response = self + let driver_result = self .context .driver .execute_singleton_operation(operation, options.operation) - .await?; + .await; // Bridge the driver response to the SDK response type. Ok(ItemResponse::new( - crate::driver_bridge::driver_response_to_cosmos_response(driver_response), + self.context + .complete_result(driver_result, || self.operation_context("upsert_item"))?, )) } @@ -688,15 +709,16 @@ impl ContainerClient { let operation = apply_item_options(operation, options.session_token, options.precondition); // Execute through the driver. - let driver_response = self + let driver_result = self .context .driver .execute_singleton_operation(operation, options.operation) - .await?; + .await; // Bridge the driver response to the SDK response type. Ok(ItemResponse::new( - crate::driver_bridge::driver_response_to_cosmos_response(driver_response), + self.context + .complete_result(driver_result, || self.operation_context("read_item"))?, )) } @@ -740,15 +762,16 @@ impl ContainerClient { let operation = apply_item_options(operation, options.session_token, options.precondition); // Execute through the driver. - let driver_response = self + let driver_result = self .context .driver .execute_singleton_operation(operation, options.operation) - .await?; + .await; // Bridge the driver response to the SDK response type. Ok(ItemResponse::new( - crate::driver_bridge::driver_response_to_cosmos_response(driver_response), + self.context + .complete_result(driver_result, || self.operation_context("delete_item"))?, )) } @@ -858,6 +881,8 @@ impl ContainerClient { Some(self.container_ref.clone()), plan, options.operation, + self.context.diagnostics_handlers.clone(), + self.operation_context("query_items"), )) } @@ -951,6 +976,8 @@ impl ContainerClient { Some(self.container_ref.clone()), plan, options.operation, + self.context.diagnostics_handlers.clone(), + self.operation_context("query_change_feed"), )) } @@ -1008,14 +1035,15 @@ impl ContainerClient { CosmosOperation::batch(self.container_ref.clone(), driver_pk).with_body(body); let operation = apply_batch_options(operation, &options); - let driver_response = self + let driver_result = self .context .driver .execute_singleton_operation(operation, options.operation) - .await?; + .await; Ok(BatchResponse::new( - crate::driver_bridge::driver_response_to_cosmos_response(driver_response), + self.context + .complete_result(driver_result, || self.operation_context("execute_batch"))?, )) } 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 753176d3421..a300fa4aa96 100644 --- a/sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client.rs +++ b/sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client.rs @@ -3,6 +3,7 @@ use crate::{ clients::{ClientContext, DatabaseClient}, + diagnostics::CosmosOperationContext, feed::QueryItemIterator, models::DatabaseProperties, models::ResourceResponse, @@ -188,6 +189,8 @@ impl CosmosClient { None, plan, operation_options, + self.context.diagnostics_handlers.clone(), + CosmosOperationContext::new().with_operation_name("query_databases"), )) } @@ -221,15 +224,20 @@ impl CosmosClient { operation_options.content_response_on_write = Some(azure_data_cosmos_driver::options::ContentResponseOnWrite::Enabled); - let driver_response = self + let driver_result = self .context .driver .execute_singleton_operation(operation, operation_options) - .await?; + .await; - Ok(ResourceResponse::new( - crate::driver_bridge::driver_response_to_cosmos_response(driver_response), - )) + Ok(ResourceResponse::new(self.context.complete_result( + driver_result, + || { + CosmosOperationContext::new() + .with_operation_name("create_database") + .with_database_name(id.to_string()) + }, + )?)) } } 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 e68479e7fad..21f7a86da37 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 @@ -3,11 +3,11 @@ //! Builder for creating [`CosmosClient`] instances. -#[cfg(feature = "fault_injection")] use std::sync::Arc; use crate::{ clients::ClientContext, + diagnostics::DiagnosticsHandler, options::{ CosmosClientOptions, OperationOptions, PartitionFailoverOptions, ThroughputControlGroupOptions, UserAgentSuffix, @@ -173,6 +173,22 @@ impl CosmosClientBuilder { self } + /// Registers a [`DiagnosticsHandler`](crate::diagnostics::DiagnosticsHandler) + /// that is invoked once per operation at completion with the operation's + /// completed [`DiagnosticsContext`](crate::diagnostics::DiagnosticsContext). + /// + /// Handlers run in registration order; call this multiple times to build an + /// ordered chain. With no handler registered the completion path does + /// nothing beyond checking whether a handler is present. + /// + /// # Arguments + /// + /// * `handler` - The handler to append to this client's diagnostics chain. + pub fn with_diagnostics_handler(mut self, handler: Arc) -> Self { + self.options.diagnostics_handlers = self.options.diagnostics_handlers.with_handler(handler); + self + } + /// Configures fault injection for testing. /// /// Accepts a vector of [`FaultInjectionRule`](crate::fault_injection::FaultInjectionRule) @@ -289,7 +305,10 @@ impl CosmosClientBuilder { let driver = runtime.into_inner().create_driver(driver_options).await?; Ok(CosmosClient { - context: ClientContext { driver }, + context: ClientContext { + driver, + diagnostics_handlers: self.options.diagnostics_handlers, + }, }) } } 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 42dea72dd2b..4f9d851a82d 100644 --- a/sdk/cosmos/azure_data_cosmos/src/clients/database_client.rs +++ b/sdk/cosmos/azure_data_cosmos/src/clients/database_client.rs @@ -3,6 +3,7 @@ use crate::{ clients::{offers_client, ClientContext, ContainerClient}, + diagnostics::CosmosOperationContext, feed::QueryItemIterator, models::ResourceResponse, models::{ContainerProperties, DatabaseProperties, ThroughputProperties}, @@ -38,6 +39,15 @@ impl DatabaseClient { } } + /// Builds the SDK-side [`CosmosOperationContext`] for this database's + /// operations, carrying the operation name plus the database identity the + /// driver context does not know. + fn operation_context(&self, operation_name: &'static str) -> CosmosOperationContext { + CosmosOperationContext::new() + .with_operation_name(operation_name) + .with_database_name(self.database_id.clone()) + } + /// Gets a [`ContainerClient`] that can be used to access the collection with the specified name. /// /// This method eagerly resolves immutable container metadata (resource ID and partition key @@ -83,14 +93,15 @@ impl DatabaseClient { let options = options.unwrap_or_default(); let operation = CosmosOperation::read_database(self.database_ref.clone()); - let driver_response = self + let driver_result = self .context .driver .execute_singleton_operation(operation, options.operation) - .await?; + .await; Ok(ResourceResponse::new( - crate::driver_bridge::driver_response_to_cosmos_response(driver_response), + self.context + .complete_result(driver_result, || self.operation_context("read_database"))?, )) } @@ -139,6 +150,8 @@ impl DatabaseClient { None, plan, operation_options, + self.context.diagnostics_handlers.clone(), + self.operation_context("query_containers"), )) } @@ -173,15 +186,19 @@ impl DatabaseClient { operation_options.content_response_on_write = Some(azure_data_cosmos_driver::options::ContentResponseOnWrite::Enabled); - let driver_response = self + let driver_result = self .context .driver .execute_singleton_operation(operation, operation_options) - .await?; + .await; - Ok(ResourceResponse::new( - crate::driver_bridge::driver_response_to_cosmos_response(driver_response), - )) + Ok(ResourceResponse::new(self.context.complete_result( + driver_result, + || { + self.operation_context("create_container") + .with_container_name(properties.id.clone()) + }, + )?)) } /// Deletes this database. @@ -197,14 +214,15 @@ impl DatabaseClient { let options = options.unwrap_or_default(); let operation = CosmosOperation::delete_database(self.database_ref.clone()); - let driver_response = self + let driver_result = self .context .driver .execute_singleton_operation(operation, options.operation) - .await?; + .await; Ok(ResourceResponse::new( - crate::driver_bridge::driver_response_to_cosmos_response(driver_response), + self.context + .complete_result(driver_result, || self.operation_context("delete_database"))?, )) } @@ -224,10 +242,11 @@ impl DatabaseClient { let resource_id = resource_id_or_error(db.system_properties.resource_id, "database")?; offers_client::find_offer( - &self.context.driver, + &self.context, self.context.driver.account(), &resource_id, options.operation, + self.operation_context("read_throughput"), ) .await } @@ -268,11 +287,12 @@ impl DatabaseClient { let resource_id = resource_id_or_error(db.system_properties.resource_id, "database")?; offers_client::begin_replace( - self.context.driver.clone(), + self.context.clone(), self.context.driver.account().clone(), &resource_id, throughput, options.operation, + self.operation_context("replace_throughput"), ) .await } diff --git a/sdk/cosmos/azure_data_cosmos/src/clients/distributed_transaction.rs b/sdk/cosmos/azure_data_cosmos/src/clients/distributed_transaction.rs index 4eb742febae..379b2fae8af 100644 --- a/sdk/cosmos/azure_data_cosmos/src/clients/distributed_transaction.rs +++ b/sdk/cosmos/azure_data_cosmos/src/clients/distributed_transaction.rs @@ -16,7 +16,7 @@ use azure_data_cosmos_driver::models as driver_models; use serde::{de::DeserializeOwned, Serialize}; use crate::clients::{ClientContext, ContainerClient}; -use crate::diagnostics::DiagnosticsContext; +use crate::diagnostics::{CosmosOperationContext, DiagnosticsContext}; use crate::models::{PartitionKey, PatchInstructions, ResponseHeaders}; use crate::options::{Precondition, SessionToken}; @@ -323,11 +323,22 @@ pub(crate) async fn commit_distributed_write( driver_models::DistributedTransactionType::Write, operations, ); - let response = context + match context .driver .execute_distributed_transaction(request, Default::default()) - .await?; - Ok(DistributedTransactionResponse::from_driver(response)) + .await + { + Ok(response) => { + let response = DistributedTransactionResponse::from_driver(response); + dispatch_transaction_diagnostics(context, &response, "commit_distributed_write"); + Ok(response) + } + Err(err) => { + let err = crate::CosmosError::from(err); + context.dispatch_error(&err, || transaction_op_context("commit_distributed_write")); + Err(err) + } + } } pub(crate) async fn execute_distributed_read( @@ -340,11 +351,43 @@ pub(crate) async fn execute_distributed_read( driver_models::DistributedTransactionType::Read, operations, ); - let response = context + match context .driver .execute_distributed_transaction(request, Default::default()) - .await?; - Ok(DistributedTransactionResponse::from_driver(response)) + .await + { + Ok(response) => { + let response = DistributedTransactionResponse::from_driver(response); + dispatch_transaction_diagnostics(context, &response, "execute_distributed_read"); + Ok(response) + } + Err(err) => { + let err = crate::CosmosError::from(err); + context.dispatch_error(&err, || transaction_op_context("execute_distributed_read")); + Err(err) + } + } +} + +/// Operation identity for a distributed transaction. A transaction can span +/// multiple containers/accounts, so only the operation name is carried. +fn transaction_op_context(operation_name: &'static str) -> CosmosOperationContext { + CosmosOperationContext::new().with_operation_name(operation_name) +} + +/// Dispatches a successful distributed transaction's diagnostics to the handler +/// chain, guarded so the empty-chain default takes no diagnostics handle. +fn dispatch_transaction_diagnostics( + context: &ClientContext, + response: &DistributedTransactionResponse, + operation_name: &'static str, +) { + if context.diagnostics_handlers.is_empty() { + return; + } + if let Some(diagnostics) = response.diagnostics() { + context.dispatch_diagnostics(&diagnostics, || transaction_op_context(operation_name)); + } } fn validate_transaction_account( diff --git a/sdk/cosmos/azure_data_cosmos/src/clients/mod.rs b/sdk/cosmos/azure_data_cosmos/src/clients/mod.rs index 0ef8a892932..ffbe25d5e4a 100644 --- a/sdk/cosmos/azure_data_cosmos/src/clients/mod.rs +++ b/sdk/cosmos/azure_data_cosmos/src/clients/mod.rs @@ -40,6 +40,9 @@ use std::sync::Arc; use azure_data_cosmos_driver::CosmosDriver; +use crate::diagnostics::{CosmosOperationContext, DiagnosticsContext, DiagnosticsHandlerChain}; +use crate::models::CosmosResponse; + /// Shared infrastructure threaded from [`CosmosClient`](super::CosmosClient) /// through [`DatabaseClient`](super::DatabaseClient) to /// [`ContainerClient`](super::ContainerClient). @@ -49,4 +52,113 @@ use azure_data_cosmos_driver::CosmosDriver; #[derive(Clone, Debug)] pub(crate) struct ClientContext { pub(crate) driver: Arc, + /// Diagnostics emission handlers invoked once per operation at completion. + /// + /// Empty by default, in which case the completion path does nothing beyond + /// checking whether a handler is present. + pub(crate) diagnostics_handlers: DiagnosticsHandlerChain, +} + +impl ClientContext { + /// Converts a completed driver response into the SDK + /// [`CosmosResponse`](crate::models::CosmosResponse) and invokes the + /// diagnostics handler chain for the operation. + /// + /// This is the per-operation completion seam for the singleton + /// (non-paginated) data- and control-plane operations: the handler chain + /// observes the operation's finalized [`DiagnosticsContext`] exactly once. + /// + /// `make_op_context` supplies the SDK-side operation identity + /// ([`CosmosOperationContext`]) — operation name, database, container — that + /// the driver context does not carry. It is a closure so the identity is + /// only materialized when at least one handler is registered, so the + /// completion path does no extra work when the chain is empty. + pub(crate) fn complete_operation( + &self, + driver_response: azure_data_cosmos_driver::models::CosmosResponse, + make_op_context: impl FnOnce() -> CosmosOperationContext, + ) -> CosmosResponse { + let response = crate::driver_bridge::driver_response_to_cosmos_response(driver_response); + // Guard the diagnostics `Arc` clone behind the emptiness check so the + // default (no-handler) path clones nothing. + if !self.diagnostics_handlers.is_empty() { + let diagnostics = response.diagnostics(); + self.dispatch_diagnostics(&diagnostics, make_op_context); + } + response + } + + /// Result-aware completion seam for singleton operations. + /// + /// Dispatches the handler chain exactly once for **both** outcomes: on + /// success from the bridged response's finalized [`DiagnosticsContext`], and + /// on failure from the context the driver attaches to the returned error + /// ([`crate::Error::diagnostics`]). The singleton call sites propagate the + /// error with `?` *after* this seam, so without it the failure-triggered + /// tracing and sampled logging would never run. + /// + /// Zero-overhead no-op when no handler is registered: neither the error's + /// diagnostics nor the operation context is materialized. + pub(crate) fn complete_result( + &self, + driver_result: Result, + make_op_context: impl FnOnce() -> CosmosOperationContext, + ) -> crate::Result + where + crate::CosmosError: From, + { + match driver_result { + Ok(driver_response) => Ok(self.complete_operation(driver_response, make_op_context)), + Err(err) => { + let err = crate::CosmosError::from(err); + if !self.diagnostics_handlers.is_empty() { + if let Some(diagnostics) = err.diagnostics() { + self.dispatch_diagnostics(&diagnostics, make_op_context); + } + } + Err(err) + } + } + } + + /// Invokes the registered diagnostics handlers with a completed context and + /// the SDK-supplied operation identity. + /// + /// Zero-overhead no-op when no handlers are registered: neither the trace + /// [`Context`](azure_core::http::Context) nor the + /// [`CosmosOperationContext`] is constructed unless at least one handler + /// will observe them. + pub(crate) fn dispatch_diagnostics( + &self, + diagnostics: &DiagnosticsContext, + make_op_context: impl FnOnce() -> CosmosOperationContext, + ) { + if self.diagnostics_handlers.is_empty() { + return; + } + let cx = azure_core::http::Context::new().with_value(make_op_context()); + self.diagnostics_handlers.dispatch(diagnostics, &cx); + } + + /// Result-aware failure completion seam for call sites that do not funnel + /// through [`complete_result`](Self::complete_result) — e.g. the offers + /// helpers, whose success path bridges the response itself. + /// + /// Dispatches the handler chain from the diagnostics the driver attached to + /// the returned error ([`crate::Error::diagnostics`]), so failure-triggered + /// tracing and sampled logging still run. Zero-overhead no-op when no handler + /// is registered: neither the error's diagnostics nor the operation context + /// is materialized. + pub(crate) fn dispatch_error( + &self, + err: &crate::CosmosError, + make_op_context: impl FnOnce() -> CosmosOperationContext, + ) { + if self.diagnostics_handlers.is_empty() { + return; + } + if let Some(diagnostics) = err.diagnostics() { + self.dispatch_diagnostics(&diagnostics, make_op_context); + } + } } diff --git a/sdk/cosmos/azure_data_cosmos/src/clients/offers_client.rs b/sdk/cosmos/azure_data_cosmos/src/clients/offers_client.rs index 8d3f0e84458..3fc92de8e8c 100644 --- a/sdk/cosmos/azure_data_cosmos/src/clients/offers_client.rs +++ b/sdk/cosmos/azure_data_cosmos/src/clients/offers_client.rs @@ -6,20 +6,24 @@ //! These functions are used by container and database clients to read and //! replace throughput offers. All operations go through the Cosmos driver. +use crate::clients::ClientContext; +use crate::diagnostics::CosmosOperationContext; use crate::{feed::FeedBody, models::CosmosResponse, models::ThroughputProperties, Query}; use azure_data_cosmos_driver::models::{AccountReference, CosmosOperation}; use azure_data_cosmos_driver::options::OperationOptions; -use azure_data_cosmos_driver::CosmosDriver; -use std::sync::Arc; /// Queries the offer for a given resource ID (RID) via the driver. /// -/// Returns `None` if no offer is configured for the resource. +/// Returns `None` if no offer is configured for the resource. The offer-query +/// operation is dispatched to the client's diagnostics handler chain — on both +/// success and failure — under the supplied `op_context` identity, so throughput +/// reads honor the same once-per-operation contract as singleton operations. pub(crate) async fn find_offer( - driver: &CosmosDriver, + context: &ClientContext, account: &AccountReference, resource_id: &str, operation_options: OperationOptions, + op_context: CosmosOperationContext, ) -> crate::Result> { let query = Query::from("SELECT * FROM c WHERE c.offerResourceId = @rid") .with_parameter("@rid", resource_id)?; @@ -27,62 +31,82 @@ pub(crate) async fn find_offer( let operation = CosmosOperation::query_offers(account.clone()).with_body(body); - let driver_response = driver + match context + .driver .execute_operation(operation, operation_options) - .await?; - let Some(driver_response) = driver_response else { - // No offer found for this resource - return Ok(None); - }; - tracing::debug!( - activity_id = ?driver_response.headers().activity_id, - request_charge = ?driver_response.headers().request_charge, - "offer query completed" - ); - let feed: FeedBody = driver_response.into_body().into_single()?; - Ok(feed.items.into_iter().next()) + .await + { + Ok(Some(driver_response)) => { + tracing::debug!( + activity_id = ?driver_response.headers().activity_id, + request_charge = ?driver_response.headers().request_charge, + "offer query completed" + ); + let response = context.complete_operation(driver_response, || op_context); + let feed: FeedBody = response.into_model()?; + Ok(feed.items.into_iter().next()) + } + // No offer found for this resource: no operation completed to dispatch. + Ok(None) => Ok(None), + Err(err) => { + let err = crate::CosmosError::from(err); + context.dispatch_error(&err, || op_context); + Err(err) + } + } } /// Reads a specific offer by its RID via the driver, returning the full response. +/// +/// The read is routed through the result-aware completion seam so the offer-read +/// operation reaches the diagnostics handler chain on both success and failure. pub(crate) async fn read_offer_by_id( - driver: &CosmosDriver, + context: &ClientContext, account: &AccountReference, offer_id: &str, + op_context: CosmosOperationContext, ) -> crate::Result { let operation = CosmosOperation::read_offer(account.clone(), offer_id.to_owned()); - let driver_response = driver + let driver_result = context + .driver .execute_singleton_operation(operation, OperationOptions::default()) - .await?; - Ok(crate::driver_bridge::driver_response_to_cosmos_response( - driver_response, - )) + .await; + context.complete_result(driver_result, || op_context) } /// Replaces the throughput for a resource and returns a [`ThroughputPoller`] to track the operation. /// /// Reads the current offer, validates the offer RID, applies the new throughput, and /// executes the replace via the driver. Returns a poller for async completion tracking. +/// +/// The offer-query, offer-replace, and every subsequent poll are dispatched to the +/// client's diagnostics handler chain under `op_context`, so an asynchronous +/// throughput replace surfaces each of its wire operations to registered handlers. pub(crate) async fn begin_replace( - driver: Arc, + context: ClientContext, account: AccountReference, resource_id: &str, throughput: ThroughputProperties, operation_options: OperationOptions, + op_context: CosmosOperationContext, ) -> crate::Result { - let mut current_throughput = - find_offer(&driver, &account, resource_id, operation_options.clone()) - .await? - .ok_or_else(|| { - // No offer exists for the resource — typically the caller - // pointed at a resource that doesn't support throughput - // (e.g. a serverless or shared-throughput container). - crate::DriverCosmosError::builder() - .with_status( - crate::error::CosmosStatus::CLIENT_NO_THROUGHPUT_OFFER_FOR_RESOURCE, - ) - .with_message("no throughput offer found for this resource") - .build() - })?; + let mut current_throughput = find_offer( + &context, + &account, + resource_id, + operation_options.clone(), + op_context.clone(), + ) + .await? + .ok_or_else(|| { + // No offer exists for the resource — typically the caller + // pointed at a resource that doesn't support throughput + // (e.g. a serverless or shared-throughput container). + crate::DriverCosmosError::builder() + .with_status(crate::error::CosmosStatus::CLIENT_NO_THROUGHPUT_OFFER_FOR_RESOURCE) + .with_message("no throughput offer found for this resource") + .build() + })?; if current_throughput.offer_id.is_empty() { // Service contract violation: an offer was returned but it has @@ -111,13 +135,13 @@ pub(crate) async fn begin_replace( opts }; - let driver_response = driver + let driver_result = context + .driver .execute_singleton_operation(operation, replace_options) - .await?; - - let response = crate::driver_bridge::driver_response_to_cosmos_response(driver_response); + .await; + let response = context.complete_result(driver_result, || op_context.clone())?; Ok(crate::clients::ThroughputPoller::new( - response, driver, account, offer_id, + response, context, account, offer_id, op_context, )) } diff --git a/sdk/cosmos/azure_data_cosmos/src/clients/throughput_poller.rs b/sdk/cosmos/azure_data_cosmos/src/clients/throughput_poller.rs index 80a46247b58..99e906de835 100644 --- a/sdk/cosmos/azure_data_cosmos/src/clients/throughput_poller.rs +++ b/sdk/cosmos/azure_data_cosmos/src/clients/throughput_poller.rs @@ -12,18 +12,18 @@ use crate::{ clients::offers_client, + clients::ClientContext, + diagnostics::CosmosOperationContext, models::ThroughputProperties, models::{CosmosResponse, ResourceResponse}, }; use azure_core::http::StatusCode; use azure_core::time::Duration; use azure_data_cosmos_driver::models::AccountReference; -use azure_data_cosmos_driver::CosmosDriver; use futures::{stream::BoxStream, Stream, StreamExt}; use std::{ future::{Future, IntoFuture}, pin::Pin, - sync::Arc, task, }; @@ -73,16 +73,21 @@ impl ThroughputPoller { /// /// The `offer_id` is provided by the caller (extracted before the replace request) /// to enable efficient single-GET polling without re-querying. + /// + /// `context` and `op_context` are threaded so each poll's offer-read reaches + /// the diagnostics handler chain under the operation's identity; the initial + /// replace response was already dispatched by the caller. pub(crate) fn new( initial_response: CosmosResponse, - driver: Arc, + context: ClientContext, account: AccountReference, offer_id: String, + op_context: CosmosOperationContext, ) -> Self { let is_pending = is_offer_replace_pending(&initial_response); if is_pending { - Self::pending(initial_response, driver, account, offer_id) + Self::pending(initial_response, context, account, offer_id, op_context) } else { Self::completed(initial_response) } @@ -99,18 +104,20 @@ impl ThroughputPoller { /// Creates a poller for an operation that is still pending. fn pending( initial_response: CosmosResponse, - driver: Arc, + context: ClientContext, account: AccountReference, offer_id: String, + op_context: CosmosOperationContext, ) -> Self { let polling_interval = DEFAULT_POLLING_INTERVAL; let stream = futures::stream::unfold( Some(PollState::Initial(Box::new(initial_response))), move |state| { - let driver = driver.clone(); + let context = context.clone(); let account = account.clone(); let offer_id = offer_id.clone(); + let op_context = op_context.clone(); async move { let state = state?; match state { @@ -119,8 +126,10 @@ impl ThroughputPoller { } PollState::Polling => { azure_core::sleep::sleep(polling_interval).await; - let result = - offers_client::read_offer_by_id(&driver, &account, &offer_id).await; + let result = offers_client::read_offer_by_id( + &context, &account, &offer_id, op_context, + ) + .await; match result { Ok(response) => { if is_offer_replace_pending(&response) { diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics.rs deleted file mode 100644 index 5c65e119d87..00000000000 --- a/sdk/cosmos/azure_data_cosmos/src/diagnostics.rs +++ /dev/null @@ -1,18 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -//! Per-operation diagnostics surfaced by the Cosmos DB SDK. -//! -//! Every fallible Cosmos operation produces a [`DiagnosticsContext`] capturing -//! request tracking, retries, regions contacted, and other observability -//! signals from the request pipeline. The context is reachable from -//! [`CosmosError`](crate::CosmosError) on failure, and from the -//! [`FeedPage`](crate::feed::FeedPage), [`ItemResponse`](crate::models::ItemResponse), and -//! similar response wrappers on success. - -// ========================================================================= -// Public API -// ========================================================================= - -#[doc(inline)] -pub use azure_data_cosmos_driver::diagnostics::{DiagnosticsContext, TransportKind}; diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/attributes.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/attributes.rs new file mode 100644 index 00000000000..b2f8c789716 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/attributes.rs @@ -0,0 +1,69 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Shared OpenTelemetry semantic-convention attribute-name literals. +//! +//! Both the metrics and the distributed-tracing handlers emit the same +//! database-client semantic-convention attributes. Keeping the string literals in +//! one place gives the SDK a single source of truth — a rename is a one-line +//! change, and the two handlers can never drift. +//! +//! Some names are used by only one handler, so under a single feature a few of +//! these are unused; the module-level `allow(dead_code)` keeps that from being a +//! warning rather than forcing per-const gating. + +#![allow(dead_code)] + +/// `db.system.name` — identifies the database system. Always +/// [`DB_SYSTEM_NAME_VALUE`] for Cosmos DB. +pub(crate) const DB_SYSTEM_NAME: &str = "db.system.name"; + +/// The stable `db.system.name` value for Azure Cosmos DB. +pub(crate) const DB_SYSTEM_NAME_VALUE: &str = "azure.cosmosdb"; + +/// `db.operation.name` — the canonical operation name (e.g. `read_item`). +pub(crate) const DB_OPERATION_NAME: &str = "db.operation.name"; + +/// `db.namespace` — the database name. +pub(crate) const DB_NAMESPACE: &str = "db.namespace"; + +/// `db.collection.name` — the container name. +pub(crate) const DB_COLLECTION_NAME: &str = "db.collection.name"; + +/// `db.response.status_code` — the HTTP status code of the response, as a string. +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"; + +/// `error.type` — a low-cardinality identifier of the error (the status code). +pub(crate) const ERROR_TYPE: &str = "error.type"; + +/// Fallback value for [`ERROR_TYPE`] when a failure carries no status anywhere +/// (per semantic conventions). +pub(crate) const ERROR_TYPE_OTHER: &str = "_OTHER"; + +/// `azure.cosmosdb.consistency.level` — effective consistency level. +pub(crate) const CONSISTENCY_LEVEL: &str = "azure.cosmosdb.consistency.level"; + +/// `azure.cosmosdb.connection.mode` — gateway vs. direct connection mode. +pub(crate) const CONNECTION_MODE: &str = "azure.cosmosdb.connection.mode"; + +/// `azure.cosmosdb.operation.contacted_regions` — regions contacted (ordered `string[]`). +pub(crate) const CONTACTED_REGIONS: &str = "azure.cosmosdb.operation.contacted_regions"; + +/// `azure.cosmosdb.response.sub_status_code` — the Cosmos sub-status code. +pub(crate) const SUB_STATUS_CODE: &str = "azure.cosmosdb.response.sub_status_code"; + +/// `azure.cosmosdb.operation.request_charge` — Request Units consumed (span attribute). +pub(crate) const REQUEST_CHARGE: &str = "azure.cosmosdb.operation.request_charge"; + +/// `azure.cosmosdb.request.activity_id` — the per-attempt activity id. +pub(crate) const ACTIVITY_ID: &str = "azure.cosmosdb.request.activity_id"; + +/// `azure.cosmosdb.machine_id` — the client/machine instance identifier. +pub(crate) const MACHINE_ID: &str = "azure.cosmosdb.machine_id"; + +/// `azure.cosmosdb.sampling.reason` — why the operation was tail-sampled for +/// emission (a failure, or which threshold it crossed). +pub(crate) const SAMPLING_REASON: &str = "azure.cosmosdb.sampling.reason"; diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/handler.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/handler.rs new file mode 100644 index 00000000000..78564ee2654 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/handler.rs @@ -0,0 +1,260 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! The [`DiagnosticsHandler`] emission extension point and its +//! [`DiagnosticsHandlerChain`]. +//! +//! The driver produces a canonical [`DiagnosticsContext`] for every operation +//! and always materializes it. This module adds the SDK-side seam that decides +//! *what to emit* from that context: applications register one or more +//! [`DiagnosticsHandler`]s, and the SDK invokes them — in registration order — +//! exactly once per operation at completion. +//! +//! The surface is deliberately small, additive, and swappable (Cosmos-local for +//! now; a candidate for promotion into `azure_core` later). Built-in handlers +//! (metrics, tracing, sampled logging) are layered on top in separate modules. + +use std::fmt; +use std::sync::Arc; + +use azure_core::http::Context; + +use crate::diagnostics::DiagnosticsContext; + +/// A sink that consumes a completed [`DiagnosticsContext`] for a single Cosmos +/// operation. +/// +/// Handlers are the SDK's emission extension point: the driver produces the +/// context and the handler decides what telemetry (metrics, spans, logs, …) to +/// emit from it. [`handle`](DiagnosticsHandler::handle) is called once per +/// operation, after the operation has completed, with the finalized context. +/// +/// Implementations must be cheap and non-blocking — they run on the operation's +/// completion path. Handlers should never panic; a panicking handler will +/// propagate out of the invoking operation. +/// +/// # Examples +/// +/// ``` +/// use std::sync::Arc; +/// use azure_data_cosmos::diagnostics::{ +/// DiagnosticsContext, DiagnosticsHandler, DiagnosticsHandlerChain, +/// }; +/// use azure_core::http::Context; +/// +/// struct LoggingHandler; +/// +/// impl DiagnosticsHandler for LoggingHandler { +/// fn handle(&self, diagnostics: &DiagnosticsContext, _cx: &Context) { +/// tracing::debug!( +/// duration_ms = diagnostics.duration().as_millis() as u64, +/// "cosmos operation completed" +/// ); +/// } +/// } +/// +/// let chain = DiagnosticsHandlerChain::new().with_handler(Arc::new(LoggingHandler)); +/// assert_eq!(chain.len(), 1); +/// ``` +pub trait DiagnosticsHandler: Send + Sync { + /// Consumes the completed diagnostics for one operation. + /// + /// * `diagnostics` - The finalized context for the just-completed operation. + /// * `cx` - A [`Context`] carrying the SDK-supplied + /// [`CosmosOperationContext`](crate::diagnostics::CosmosOperationContext) + /// for the operation (operation name, database, container) when one is + /// available. This is a diagnostics-scoped context built at completion, not + /// 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<'_>); +} + +/// An ordered, cheaply cloneable chain of [`DiagnosticsHandler`]s. +/// +/// The chain is the unit the SDK invokes at operation completion. Handlers run +/// in registration order, so ordering is deterministic. An empty chain — the +/// default when no handler is registered — does nothing. +/// +/// The chain is backed by a shared, reference-counted slice, so cloning it (for +/// example when a [`CosmosClient`](crate::CosmosClient) hands state down to a +/// [`DatabaseClient`](crate::clients::DatabaseClient) or +/// [`ContainerClient`](crate::clients::ContainerClient)) is a single atomic +/// increment rather than a deep copy. +/// +/// Register handlers via +/// [`CosmosClientBuilder::with_diagnostics_handler`](crate::CosmosClientBuilder::with_diagnostics_handler). +#[derive(Clone)] +pub struct DiagnosticsHandlerChain { + handlers: Arc<[Arc]>, +} + +impl DiagnosticsHandlerChain { + /// Creates an empty chain. + /// + /// An empty chain performs no work when invoked. + pub fn new() -> Self { + Self { + handlers: Vec::new().into(), + } + } + + /// Creates a chain from an ordered list of handlers. + /// + /// Handlers are invoked in the order supplied. + pub fn from_handlers(handlers: Vec>) -> Self { + Self { + handlers: handlers.into(), + } + } + + /// Returns a new chain with `handler` appended to the end. + /// + /// This is additive: existing handlers keep their relative order and the new + /// handler runs last. + #[must_use] + pub fn with_handler(&self, handler: Arc) -> Self { + let mut handlers: Vec> = self.handlers.to_vec(); + handlers.push(handler); + Self { + handlers: handlers.into(), + } + } + + /// Returns the registered handlers in invocation order. + pub fn handlers(&self) -> &[Arc] { + &self.handlers + } + + /// Returns `true` when no handlers are registered. + pub fn is_empty(&self) -> bool { + self.handlers.is_empty() + } + + /// Returns the number of registered handlers. + pub fn len(&self) -> usize { + self.handlers.len() + } + + /// Invokes every handler, in order, with the completed diagnostics. + /// + /// A no-op when the chain is empty. + pub(crate) fn dispatch(&self, diagnostics: &DiagnosticsContext, cx: &Context<'_>) { + for handler in self.handlers.iter() { + handler.handle(diagnostics, cx); + } + } +} + +impl Default for DiagnosticsHandlerChain { + fn default() -> Self { + Self::new() + } +} + +impl fmt::Debug for DiagnosticsHandlerChain { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + // Handlers are `dyn` trait objects without a `Debug` bound, so surface + // only the count to keep the trait surface minimal. + f.debug_struct("DiagnosticsHandlerChain") + .field("handlers", &self.handlers.len()) + .finish() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use azure_data_cosmos_driver::models::ActivityId; + use std::sync::Mutex; + + /// A test handler that appends `(label, activity_id)` to a shared log every + /// time it is invoked, so tests can assert both per-operation receipt and + /// cross-handler ordering. + struct RecordingHandler { + label: &'static str, + log: Arc>>, + } + + impl DiagnosticsHandler for RecordingHandler { + fn handle(&self, diagnostics: &DiagnosticsContext, _cx: &Context<'_>) { + self.log + .lock() + .unwrap() + .push((self.label, diagnostics.activity_id().clone())); + } + } + + fn recording( + label: &'static str, + log: &Arc>>, + ) -> Arc { + Arc::new(RecordingHandler { + label, + log: Arc::clone(log), + }) + } + + #[test] + fn empty_chain_is_noop() { + let chain = DiagnosticsHandlerChain::new(); + assert!(chain.is_empty()); + assert_eq!(chain.len(), 0); + + // Dispatching against an empty chain must not panic and must do nothing. + let ctx = DiagnosticsContext::for_testing(ActivityId::new_uuid()); + chain.dispatch(&ctx, &Context::new()); + } + + #[test] + fn handlers_receive_completed_context_in_registration_order() { + let log = Arc::new(Mutex::new(Vec::new())); + let chain = DiagnosticsHandlerChain::from_handlers(vec![ + recording("a", &log), + recording("b", &log), + recording("c", &log), + ]); + assert_eq!(chain.len(), 3); + + // Two distinct "operations", each with its own completed context. + let op1 = ActivityId::new_uuid(); + let op2 = ActivityId::new_uuid(); + let cx = Context::new(); + chain.dispatch(&DiagnosticsContext::for_testing(op1.clone()), &cx); + chain.dispatch(&DiagnosticsContext::for_testing(op2.clone()), &cx); + + let recorded = log.lock().unwrap().clone(); + // Every handler saw each operation's own completed context, and the + // handlers ran in deterministic registration order (a, b, c) per op. + assert_eq!( + recorded, + vec![ + ("a", op1.clone()), + ("b", op1.clone()), + ("c", op1), + ("a", op2.clone()), + ("b", op2.clone()), + ("c", op2), + ] + ); + } + + #[test] + fn with_handler_appends_and_is_additive() { + let log = Arc::new(Mutex::new(Vec::new())); + let base = DiagnosticsHandlerChain::new().with_handler(recording("first", &log)); + let extended = base.with_handler(recording("second", &log)); + + // `with_handler` returns a new chain and leaves the original untouched. + assert_eq!(base.len(), 1); + assert_eq!(extended.len(), 2); + + let op = ActivityId::new_uuid(); + extended.dispatch( + &DiagnosticsContext::for_testing(op.clone()), + &Context::new(), + ); + + let recorded = log.lock().unwrap().clone(); + assert_eq!(recorded, vec![("first", op.clone()), ("second", op)]); + } +} diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/logging/handler.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/logging/handler.rs new file mode 100644 index 00000000000..101c48e26fe --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/logging/handler.rs @@ -0,0 +1,253 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Composable, tail-sampled diagnostics logging. +//! +//! [`SamplingLogHandler`] is a **wrapper**: it applies tail-based sampling and a +//! per-window rate limit, then delegates the actual emission to an inner +//! [`DiagnosticsHandler`]. [`TracingLogHandler`] is the default **leaf**: it +//! writes a compact diagnostics line through [`tracing`]. Composing them +//! (`SamplingLogHandler` around `TracingLogHandler`) reproduces the built-in +//! "sampled log" behavior, while letting callers swap in any inner handler. +//! +//! [`tracing`]: https://docs.rs/tracing + +use std::sync::Arc; +use std::time::Instant; + +use azure_core::http::Context; +use azure_data_cosmos_driver::{ + diagnostics::DiagnosticsContext, DiagnosticsThresholds, DiagnosticsVerbosity, +}; + +use crate::diagnostics::rate_limiter::{RateLimiter, RateLimiterConfig}; +use crate::diagnostics::reason::EmitReason; +use crate::diagnostics::{CosmosOperationContext, DiagnosticsHandler}; + +/// `tracing` target for emitted sampled-diagnostics lines. +const SAMPLED_TARGET: &str = "azure_data_cosmos::diagnostics::sampled"; + +/// `tracing` target for the "suppressed N until reset" notice. +const SUPPRESSED_TARGET: &str = "azure_data_cosmos::diagnostics::suppressed"; + +/// A leaf [`DiagnosticsHandler`] that writes a compact diagnostics line for the +/// context it is handed, through the [`tracing`](https://docs.rs/tracing) +/// ecosystem. +/// +/// It performs **no** sampling or rate limiting of its own — it emits for every +/// context passed to [`handle`](DiagnosticsHandler::handle). Wrap it in a +/// [`SamplingLogHandler`] (the default composition) to gate emission on failures +/// / threshold breaches and cap the rate under a storm, or register it directly +/// to log every completed operation. +/// +/// Failed operations are logged at `warn`, everything else at `info`; both use +/// the `azure_data_cosmos::diagnostics::sampled` target. +/// +/// # Examples +/// +/// ``` +/// use std::sync::Arc; +/// use azure_data_cosmos::diagnostics::TracingLogHandler; +/// +/// let handler = Arc::new(TracingLogHandler::new()); +/// # let _ = handler; +/// ``` +#[derive(Clone, Copy, Debug, Default)] +pub struct TracingLogHandler; + +impl TracingLogHandler { + /// Creates a `TracingLogHandler`. + pub fn new() -> Self { + Self + } +} + +impl DiagnosticsHandler for TracingLogHandler { + fn handle(&self, diagnostics: &DiagnosticsContext, cx: &Context<'_>) { + // The sampling wrapper stamps *why* the operation was sampled onto the + // context; when this handler is used directly (unsampled), fall back to a + // stable marker so the field is always present. + let reason = cx + .value::() + .copied() + .map(EmitReason::as_str) + .unwrap_or("unsampled"); + + // Compute the JSON line inside each macro call so `to_json_string` is + // only evaluated when a subscriber is actually listening for the event + // (the `tracing` macros run an "is enabled" check before evaluating + // field expressions). + if diagnostics.is_failure() { + tracing::warn!(target: SAMPLED_TARGET, reason, diagnostics = %diagnostics.to_json_string(Some(DiagnosticsVerbosity::Summary)), "cosmos operation diagnostics"); + } else { + tracing::info!(target: SAMPLED_TARGET, reason, diagnostics = %diagnostics.to_json_string(Some(DiagnosticsVerbosity::Summary)), "cosmos operation diagnostics"); + } + } +} + +/// A wrapping [`DiagnosticsHandler`] that applies tail-based sampling and a +/// per-window rate limit, delegating emission to an inner handler. +/// +/// For each completed operation it first applies the same tail-based sampling +/// gate the tracing handler uses — emit if, and only if, the operation *is +/// completed* and *failed* or breached a [`DiagnosticsThresholds`]. It then +/// limits how many emissions are dispatched during a time window (≈100/min by +/// default); when emissions are suppressed a single "suppressed N until reset" +/// warning is emitted per window. Only when both gates allow does it call the +/// inner handler. +/// +/// The inner handler defaults to a [`TracingLogHandler`], so a bare +/// `SamplingLogHandler::new()` reproduces the built-in sampled-log behavior. Pass +/// your own inner handler with [`with_handler`](Self::with_handler) (or the +/// `*_and_handler` constructors) to feed sampled, rate-limited contexts to any +/// other sink. +/// +/// Register it with +/// [`CosmosClientBuilder::with_diagnostics_handler`](crate::CosmosClientBuilder::with_diagnostics_handler). +/// +/// # Examples +/// +/// ``` +/// use std::sync::Arc; +/// use azure_data_cosmos::diagnostics::SamplingLogHandler; +/// +/// // Default: samples + rate-limits, then logs via `tracing`. +/// let handler = Arc::new(SamplingLogHandler::new()); +/// # let _ = handler; +/// ``` +pub struct SamplingLogHandler { + thresholds: DiagnosticsThresholds, + limiter: RateLimiter, + inner: Arc, +} + +impl SamplingLogHandler { + /// Creates a handler with default thresholds and rate limiting (~100/min) + /// wrapping a default [`TracingLogHandler`]. + pub fn new() -> Self { + Self::with_thresholds(DiagnosticsThresholds::default()) + } + + /// Creates a handler with the supplied sampling thresholds, default rate + /// limiting, and a default [`TracingLogHandler`] inner. + pub fn with_thresholds(thresholds: DiagnosticsThresholds) -> Self { + Self::with_thresholds_and_rate_limit(thresholds, RateLimiterConfig::default()) + } + + /// Creates a handler with the supplied sampling thresholds and rate-limiter + /// configuration, wrapping a default [`TracingLogHandler`] inner. + pub fn with_thresholds_and_rate_limit( + thresholds: DiagnosticsThresholds, + rate_limit: RateLimiterConfig, + ) -> Self { + Self::with_thresholds_rate_limit_and_handler( + thresholds, + rate_limit, + Arc::new(TracingLogHandler::new()), + ) + } + + /// Creates a handler that samples and rate-limits with the defaults, then + /// delegates emission to `inner`. + pub fn with_handler(inner: Arc) -> Self { + Self::with_thresholds_rate_limit_and_handler( + DiagnosticsThresholds::default(), + RateLimiterConfig::default(), + inner, + ) + } + + /// Creates a handler with the supplied thresholds and default rate limiting, + /// delegating emission to `inner`. + pub fn with_thresholds_and_handler( + thresholds: DiagnosticsThresholds, + inner: Arc, + ) -> Self { + Self::with_thresholds_rate_limit_and_handler( + thresholds, + RateLimiterConfig::default(), + inner, + ) + } + + /// Creates a handler with the supplied thresholds and rate-limiter + /// configuration, delegating emission to `inner`. + pub fn with_thresholds_rate_limit_and_handler( + thresholds: DiagnosticsThresholds, + rate_limit: RateLimiterConfig, + inner: Arc, + ) -> Self { + Self { + thresholds, + limiter: RateLimiter::new(rate_limit), + inner, + } + } + + /// Returns the sampling thresholds this handler applies. + pub fn thresholds(&self) -> &DiagnosticsThresholds { + &self.thresholds + } + + /// Returns whether the given completed context would pass the tail-based + /// sampling gate (before rate limiting). + pub fn should_log(&self, diagnostics: &DiagnosticsContext) -> bool { + should_log(diagnostics, &self.thresholds, None) + } +} + +impl Default for SamplingLogHandler { + fn default() -> Self { + Self::new() + } +} + +impl DiagnosticsHandler for SamplingLogHandler { + fn handle(&self, diagnostics: &DiagnosticsContext, cx: &Context<'_>) { + let op = cx.value::(); + if !should_log(diagnostics, &self.thresholds, op) { + return; + } + + let decision = self.limiter.check(diagnostics.is_failure(), Instant::now()); + + if let Some(suppressed) = decision.suppression_notice { + tracing::warn!( + target: SUPPRESSED_TARGET, + suppressed, + "cosmos diagnostics: suppressed {suppressed} sampled log line(s) until window reset" + ); + } + + if decision.emit { + // Stamp *why* the operation was sampled onto the context so the inner + // handler can surface it. `should_log` passed, so a reason exists. + match EmitReason::of(diagnostics, &self.thresholds, op) { + Some(reason) => { + let cx = cx.clone().with_value(reason); + self.inner.handle(diagnostics, &cx); + } + None => self.inner.handle(diagnostics, cx), + } + } + } +} + +/// The tail-based sampling gate: log iff the operation is completed and either +/// failed or crossed a sampling threshold. +/// +/// `op` supplies the SDK-side operation identity so the threshold classifier +/// can distinguish point from non-point operations; production driver contexts +/// do not carry the operation name. +pub(crate) fn should_log( + diagnostics: &DiagnosticsContext, + thresholds: &DiagnosticsThresholds, + op: Option<&CosmosOperationContext>, +) -> bool { + diagnostics.is_completed() + && (diagnostics.is_failure() + || diagnostics.is_threshold_violated_for( + thresholds, + op.and_then(CosmosOperationContext::operation_name), + )) +} diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/logging/mod.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/logging/mod.rs new file mode 100644 index 00000000000..98908f7867c --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/logging/mod.rs @@ -0,0 +1,278 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Rate-limited, tail-sampled diagnostics logging for Cosmos DB operations. +//! +//! Provides two composable [`DiagnosticsHandler`](crate::diagnostics::DiagnosticsHandler)s: +//! +//! - [`TracingLogHandler`] — a leaf that writes a compact diagnostics line +//! through the [`tracing`] ecosystem for whatever context it is handed. +//! - [`SamplingLogHandler`] — a wrapper that applies tail-based sampling (only +//! failures and threshold breaches) plus a per-window rate limit, delegating +//! the actual emission to an inner handler (a [`TracingLogHandler`] by +//! default). +//! +//! [`tracing`]: https://docs.rs/tracing + +mod handler; + +pub use handler::{SamplingLogHandler, TracingLogHandler}; + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + use std::time::{Duration, Instant}; + + use azure_core::http::{Context, StatusCode}; + use azure_data_cosmos_driver::diagnostics::{DiagnosticsContext, RequestDiagnostics}; + use azure_data_cosmos_driver::models::{ActivityId, RequestCharge}; + use azure_data_cosmos_driver::options::Region; + use azure_data_cosmos_driver::{CosmosStatus, DiagnosticsThresholds}; + use tracing_subscriber::layer::{Context as LayerContext, SubscriberExt}; + use tracing_subscriber::Layer; + + use super::handler::{should_log, SamplingLogHandler}; + use crate::diagnostics::rate_limiter::RateLimiterConfig; + use crate::diagnostics::{CosmosOperationContext, DiagnosticsHandler}; + + /// Builds a completed single-attempt context with the given duration, status, + /// and request charge. + fn context(duration: Duration, status: CosmosStatus, charge: f64) -> DiagnosticsContext { + let now = Instant::now(); + let request = RequestDiagnostics::for_testing( + "https://acct.documents.azure.com:443/", + Some(Region::new("West US 2")), + status, + RequestCharge::new(charge), + now - duration, + now, + ); + DiagnosticsContext::for_testing_with_requests( + ActivityId::new_uuid(), + duration, + Some(status), + Some("read_item"), + vec![request], + ) + } + + #[test] + fn gate_skips_fast_success_and_admits_failures_and_breaches() { + let thresholds = DiagnosticsThresholds::default(); + + // Fast, cheap success: nothing to log. + let ok = context( + Duration::from_millis(5), + CosmosStatus::new(StatusCode::Ok), + 2.0, + ); + assert!(!should_log(&ok, &thresholds, None)); + + // Failure: log. + let failed = context( + Duration::from_millis(5), + CosmosStatus::new(StatusCode::TooManyRequests), + 2.0, + ); + assert!(should_log(&failed, &thresholds, None)); + + // Threshold breach (RU over a low threshold): log. + let strict = DiagnosticsThresholds::default().with_request_charge(1.0); + let expensive = context( + Duration::from_millis(5), + CosmosStatus::new(StatusCode::Ok), + 500.0, + ); + assert!(should_log(&expensive, &strict, None)); + } + + #[test] + fn non_point_threshold_uses_operation_name_from_context() { + // A production driver context carries no operation name, so a 2s + // operation classifies as a point op (1s) and is wrongly logged. The + // SDK operation identity switches it to the non-point (3s) threshold. + let thresholds = DiagnosticsThresholds::default(); + let now = Instant::now(); + let request = RequestDiagnostics::for_testing( + "https://acct.documents.azure.com:443/", + Some(Region::new("West US 2")), + CosmosStatus::new(StatusCode::Ok), + RequestCharge::new(2.0), + now - Duration::from_millis(2000), + now, + ); + let slow_non_point = DiagnosticsContext::for_testing_with_requests( + ActivityId::new_uuid(), + Duration::from_millis(2000), + Some(CosmosStatus::new(StatusCode::Ok)), + None, + vec![request], + ); + + // No operation identity → point (1s) fallback → logged. + assert!(should_log(&slow_non_point, &thresholds, None)); + // Query op identity → non-point (3s) threshold → not logged at 2s. + let op = CosmosOperationContext::new().with_operation_name("query_items"); + assert!(!should_log(&slow_non_point, &thresholds, Some(&op))); + } + + /// Counts sampled lines and suppression notices by `tracing` target. + struct CountingLayer { + sampled: Arc, + suppressed: Arc, + } + + impl Layer for CountingLayer { + fn on_event(&self, event: &tracing::Event<'_>, _cx: LayerContext<'_, S>) { + match event.metadata().target() { + "azure_data_cosmos::diagnostics::sampled" => { + self.sampled.fetch_add(1, Ordering::SeqCst); + } + "azure_data_cosmos::diagnostics::suppressed" => { + self.suppressed.fetch_add(1, Ordering::SeqCst); + } + _ => {} + } + } + } + + #[test] + fn storm_is_capped_with_one_suppression_notice() { + let sampled = Arc::new(AtomicUsize::new(0)); + let suppressed = Arc::new(AtomicUsize::new(0)); + let layer = CountingLayer { + sampled: Arc::clone(&sampled), + suppressed: Arc::clone(&suppressed), + }; + let subscriber = tracing_subscriber::registry().with(layer); + + // Cap 5 per 200ms window, no failure reserve. Drive threshold-breaching + // successes (RU 2 > 1) so they log without tapping the failure reserve. + let handler = SamplingLogHandler::with_thresholds_and_rate_limit( + DiagnosticsThresholds::default().with_request_charge(1.0), + RateLimiterConfig { + max_per_window: 5, + window: Duration::from_millis(200), + failure_reserve: 0, + }, + ); + let ctx = context( + Duration::from_millis(5), + CosmosStatus::new(StatusCode::Ok), + 2.0, + ); + + tracing::subscriber::with_default(subscriber, || { + // 30 events in one window: only 5 admitted, 25 suppressed. + for _ in 0..30 { + handler.handle(&ctx, &Context::new()); + } + // Roll into the next window and log once more, flushing the notice. + std::thread::sleep(Duration::from_millis(250)); + handler.handle(&ctx, &Context::new()); + }); + + // 5 (window 1) + 1 (window 2) sampled lines, and exactly one notice. + assert_eq!(sampled.load(Ordering::SeqCst), 6); + assert_eq!(suppressed.load(Ordering::SeqCst), 1); + } + + /// Counts how many times it is invoked, standing in for an arbitrary inner + /// sink wrapped by `SamplingLogHandler`. + #[derive(Default)] + struct CountingInner(AtomicUsize); + + impl DiagnosticsHandler for CountingInner { + fn handle(&self, _diagnostics: &DiagnosticsContext, _cx: &Context<'_>) { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + + #[test] + fn wrapper_delegates_to_inner_only_when_sampled() { + let inner = Arc::new(CountingInner::default()); + let handler = SamplingLogHandler::with_handler(inner.clone()); + + // A fast, cheap success is not "interesting": the wrapper must NOT call + // the inner handler. + let ok = context( + Duration::from_millis(5), + CosmosStatus::new(StatusCode::Ok), + 2.0, + ); + handler.handle(&ok, &Context::new()); + assert_eq!(inner.0.load(Ordering::SeqCst), 0); + + // A failure passes the sampling gate: the inner handler is invoked once. + let failed = context( + Duration::from_millis(5), + CosmosStatus::new(StatusCode::TooManyRequests), + 2.0, + ); + handler.handle(&failed, &Context::new()); + assert_eq!(inner.0.load(Ordering::SeqCst), 1); + } + + /// Captures the `reason` field value of the most recent sampled log event. + struct ReasonCapture(Arc>>); + + impl Layer for ReasonCapture { + fn on_event(&self, event: &tracing::Event<'_>, _cx: LayerContext<'_, S>) { + if event.metadata().target() != "azure_data_cosmos::diagnostics::sampled" { + return; + } + struct Visitor<'a>(&'a mut Option); + impl tracing::field::Visit for Visitor<'_> { + fn record_str(&mut self, field: &tracing::field::Field, value: &str) { + if field.name() == "reason" { + *self.0 = Some(value.to_string()); + } + } + fn record_debug( + &mut self, + field: &tracing::field::Field, + value: &dyn std::fmt::Debug, + ) { + if field.name() == "reason" { + *self.0 = Some(format!("{value:?}")); + } + } + } + let mut slot = self.0.lock().unwrap(); + event.record(&mut Visitor(&mut slot)); + } + } + + #[test] + fn sampled_line_carries_the_emit_reason() { + let captured = Arc::new(std::sync::Mutex::new(None)); + let layer = ReasonCapture(Arc::clone(&captured)); + let subscriber = tracing_subscriber::registry().with(layer); + + // Threshold breaches log at info, so the info level must be enabled. + let strict = DiagnosticsThresholds::default().with_request_charge(1.0); + let handler = SamplingLogHandler::with_thresholds(strict); + + tracing::subscriber::with_default(subscriber, || { + // A failure is reported with reason "failure". + let failed = context( + Duration::from_millis(5), + CosmosStatus::new(StatusCode::TooManyRequests), + 2.0, + ); + handler.handle(&failed, &Context::new()); + assert_eq!(captured.lock().unwrap().as_deref(), Some("failure")); + + // A successful-but-expensive operation is reported with the specific + // threshold it crossed. + let expensive = context( + Duration::from_millis(5), + CosmosStatus::new(StatusCode::Ok), + 500.0, + ); + handler.handle(&expensive, &Context::new()); + assert_eq!(captured.lock().unwrap().as_deref(), Some("request_charge")); + }); + } +} diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/attributes.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/attributes.rs new file mode 100644 index 00000000000..6286e098e92 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/attributes.rs @@ -0,0 +1,101 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Centralized metric- and attribute-name string literals for Cosmos DB OTel +//! metrics. +//! +//! The OpenTelemetry semantic conventions for database clients define the exact +//! metric and attribute names emitted here. Metric names and instrument units are +//! metrics-specific and defined here; the shared **attribute** names are re-exported +//! from the crate-internal `diagnostics::attributes` module so metrics and tracing +//! can't drift. +//! +//! Two tiers, matching the semantic conventions: +//! - **Stable** names — emitted unconditionally (operation-scope, low cardinality). +//! - **Optional** names — emitted only when the matching +//! [`MetricsOptions`](super::MetricsOptions) toggle is opted into. + +use crate::diagnostics::attributes; + +// ========================================================================= +// Metric names +// ========================================================================= + +/// Stable histogram (seconds): total client-observed duration of an operation. +/// +/// This is the primary Cosmos metric — the one to graph for latency SLOs. +pub const METRIC_OPERATION_DURATION: &str = "db.client.operation.duration"; + +/// Optional histogram (request units): request charge (RU) for an operation. +pub const METRIC_OPERATION_REQUEST_CHARGE: &str = "azure.cosmosdb.client.operation.request_charge"; + +/// Optional histogram (rows): number of rows/items returned by an operation. +pub const METRIC_RESPONSE_RETURNED_ROWS: &str = "db.client.response.returned_rows"; + +// ========================================================================= +// Instrument units +// ========================================================================= +// +// Unit strings follow the Unified Code for Units of Measure (UCUM): +// . OpenTelemetry adopts UCUM for instrument units; see +// . + +/// Unit for [`METRIC_OPERATION_DURATION`] — seconds. +pub const UNIT_SECONDS: &str = "s"; + +/// Unit for [`METRIC_OPERATION_REQUEST_CHARGE`] — Cosmos request units. +pub const UNIT_REQUEST_UNIT: &str = "{request_unit}"; + +/// Unit for [`METRIC_RESPONSE_RETURNED_ROWS`] — rows. +pub const UNIT_ROW: &str = "{row}"; + +// ========================================================================= +// Stable attributes (always emitted; operation scope, low cardinality) +// +// These alias the shared semconv literals in `crate::diagnostics::attributes` +// so the string lives in exactly one place. +// ========================================================================= + +/// `db.system.name` — identifies the database system. +pub const ATTR_DB_SYSTEM_NAME: &str = attributes::DB_SYSTEM_NAME; + +/// Value for [`ATTR_DB_SYSTEM_NAME`] on every Cosmos metric. +pub const DB_SYSTEM_NAME_VALUE: &str = attributes::DB_SYSTEM_NAME_VALUE; + +/// `db.operation.name` — canonical operation name (e.g. `read_item`). +pub const ATTR_DB_OPERATION_NAME: &str = attributes::DB_OPERATION_NAME; + +/// `db.collection.name` — the container name. +pub const ATTR_DB_COLLECTION_NAME: &str = attributes::DB_COLLECTION_NAME; + +/// `db.namespace` — the database name. +pub const ATTR_DB_NAMESPACE: &str = attributes::DB_NAMESPACE; + +/// `db.response.status_code` — the HTTP status code of the response. +pub const ATTR_DB_RESPONSE_STATUS_CODE: &str = attributes::DB_RESPONSE_STATUS_CODE; + +/// `error.type` — present only when the operation failed. +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; + +/// 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; + +// ========================================================================= +// Extended attributes (opt-in; may be higher cardinality) +// ========================================================================= + +/// `azure.cosmosdb.consistency.level` — effective consistency level. +pub const ATTR_CONSISTENCY_LEVEL: &str = attributes::CONSISTENCY_LEVEL; + +/// `azure.cosmosdb.operation.contacted_regions` — regions contacted. +pub const ATTR_CONTACTED_REGIONS: &str = attributes::CONTACTED_REGIONS; + +/// `azure.cosmosdb.response.sub_status_code` — Cosmos sub-status code. +pub const ATTR_SUB_STATUS_CODE: &str = attributes::SUB_STATUS_CODE; + +/// `azure.cosmosdb.connection.mode` — gateway vs. direct connection mode. +pub const ATTR_CONNECTION_MODE: &str = attributes::CONNECTION_MODE; diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/handler.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/handler.rs new file mode 100644 index 00000000000..1198782d289 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/handler.rs @@ -0,0 +1,541 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! [`CosmosMetricsHandler`] — emits OpenTelemetry metrics from a completed +//! [`DiagnosticsContext`]. + +use azure_core::http::Context; +use opentelemetry::metrics::Meter; +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}; + +/// Instrumentation scope name used for the Cosmos [`Meter`]. +const METER_NAME: &str = "azure_data_cosmos"; + +/// A [`DiagnosticsHandler`] that records OpenTelemetry metrics for every +/// completed Cosmos DB operation. +/// +/// Register it via +/// [`CosmosClientBuilder::with_diagnostics_handler`](crate::CosmosClientBuilder::with_diagnostics_handler): +/// +/// ``` +/// use std::sync::Arc; +/// use azure_data_cosmos::diagnostics::CosmosMetricsHandler; +/// +/// // Uses the globally-registered OpenTelemetry meter provider. +/// let handler = Arc::new(CosmosMetricsHandler::new()); +/// # let _handler: Arc = handler; +/// ``` +/// +/// The handler always records the stable `db.client.operation.duration` +/// histogram. The optional per-signal metrics and the extended attribute set are +/// opt-in via [`MetricsOptions`] (see [`with_options`](CosmosMetricsHandler::with_options)). +/// +/// 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 +/// no-op even after a real provider is installed later, so metrics would be +/// silently dropped. If you need to build the handler before the global provider +/// is ready, bind it to an explicit meter with +/// [`with_meter`](CosmosMetricsHandler::with_meter) instead. +pub struct CosmosMetricsHandler { + instruments: Instruments, + options: MetricsOptions, +} + +impl CosmosMetricsHandler { + /// Creates a handler backed by the globally-registered meter provider, with + /// default [`MetricsOptions`] (stable metric only). + pub fn new() -> Self { + Self::with_options(MetricsOptions::default()) + } + + /// Creates a handler backed by the global meter provider with the given + /// [`MetricsOptions`]. + pub fn with_options(options: MetricsOptions) -> Self { + Self::from_meter(&global::meter(METER_NAME), options) + } + + /// Creates a handler that records into a specific [`Meter`], with default + /// [`MetricsOptions`]. + /// + /// Useful for tests (against an in-memory meter) or to bind metrics to a + /// meter provider other than the global one. + pub fn with_meter(meter: Meter) -> Self { + Self::from_meter(&meter, MetricsOptions::default()) + } + + /// Creates a handler that records into a specific [`Meter`] with the given + /// [`MetricsOptions`]. + pub fn with_meter_and_options(meter: Meter, options: MetricsOptions) -> Self { + Self::from_meter(&meter, options) + } + + fn from_meter(meter: &Meter, options: MetricsOptions) -> Self { + Self { + instruments: Instruments::new(meter), + options, + } + } + + /// Resolves `server.address`: the operation-context override if present, + /// otherwise the host of the last contacted endpoint. + fn server_address( + &self, + diagnostics: &DiagnosticsContext, + op: Option<&CosmosOperationContext>, + ) -> Option { + if let Some(address) = op.and_then(CosmosOperationContext::server_address) { + return Some(address.to_string()); + } + let requests = diagnostics.requests(); + requests + .last() + .and_then(|request| host_of(request.endpoint())) + } + + /// Builds the operation-scope attribute set shared by the duration and + /// per-operation optional histograms. + fn build_attributes( + &self, + diagnostics: &DiagnosticsContext, + op: Option<&CosmosOperationContext>, + server_address: Option<&str>, + ) -> Vec { + let mut attrs = Vec::with_capacity(8); + + // db.system.name is constant for every Cosmos metric. + attrs.push(KeyValue::new( + attributes::ATTR_DB_SYSTEM_NAME, + attributes::DB_SYSTEM_NAME_VALUE, + )); + + // Operation-scope identity supplied by the SDK. + if let Some(op) = op { + if let Some(name) = op.operation_name() { + attrs.push(KeyValue::new( + attributes::ATTR_DB_OPERATION_NAME, + name.to_string(), + )); + } + if let Some(container) = op.container_name() { + attrs.push(KeyValue::new( + attributes::ATTR_DB_COLLECTION_NAME, + container.to_string(), + )); + } + if let Some(database) = op.database_name() { + attrs.push(KeyValue::new( + attributes::ATTR_DB_NAMESPACE, + database.to_string(), + )); + } + } + + // Response status and, on failure, error.type (per semantic conventions). + // Use the effective status (operation status, else the terminal attempt's) + // so status-less error-finalization paths still report an accurate status + // and error.type instead of the _OTHER catch-all. + match diagnostics.effective_status() { + Some(status) => { + let code = u16::from(status.status_code()); + attrs.push(KeyValue::new( + attributes::ATTR_DB_RESPONSE_STATUS_CODE, + code.to_string(), + )); + if !status.is_success() { + attrs.push(KeyValue::new(attributes::ATTR_ERROR_TYPE, code.to_string())); + } + } + None => { + // No status anywhere — a client/transport failure with no HTTP + // response and no attempt. Classify it as the semconv catch-all. + attrs.push(KeyValue::new( + attributes::ATTR_ERROR_TYPE, + attributes::ERROR_TYPE_OTHER, + )); + } + } + + if let Some(address) = server_address { + attrs.push(KeyValue::new( + attributes::ATTR_SERVER_ADDRESS, + address.to_string(), + )); + } + + // Extended attributes are opt-in (higher cardinality; D7). + if self.options.extended_attributes_enabled() { + if let Some(op) = op { + if let Some(level) = op.consistency_level() { + attrs.push(KeyValue::new( + attributes::ATTR_CONSISTENCY_LEVEL, + level.to_string(), + )); + } + if let Some(mode) = op.connection_mode() { + attrs.push(KeyValue::new( + attributes::ATTR_CONNECTION_MODE, + mode.to_string(), + )); + } + } + if let Some(sub_status) = diagnostics.effective_status().and_then(|s| s.sub_status()) { + attrs.push(KeyValue::new( + attributes::ATTR_SUB_STATUS_CODE, + i64::from(sub_status.value()), + )); + } + let regions = diagnostics.regions_contacted(); + if !regions.is_empty() { + // Semantic conventions define contacted_regions as an ordered + // string[]; emit an OpenTelemetry array, not a joined scalar. + let values: Vec = regions + .iter() + .map(|region| StringValue::from(region.as_str().to_string())) + .collect(); + attrs.push(KeyValue::new( + attributes::ATTR_CONTACTED_REGIONS, + Value::Array(Array::String(values)), + )); + } + } + + attrs + } +} + +impl Default for CosmosMetricsHandler { + fn default() -> Self { + Self::new() + } +} + +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 + // only the configuration. + f.debug_struct("CosmosMetricsHandler") + .field("options", &self.options) + .finish_non_exhaustive() + } +} + +impl DiagnosticsHandler for CosmosMetricsHandler { + fn handle(&self, diagnostics: &DiagnosticsContext, cx: &Context<'_>) { + let op = cx.value::(); + let server_address = self.server_address(diagnostics, op); + let attributes = self.build_attributes(diagnostics, op, server_address.as_deref()); + + // Stable metric: always recorded. + self.instruments + .operation_duration + .record(diagnostics.duration().as_secs_f64(), &attributes); + + // Optional per-signal metrics (each opt-in). + if self.options.request_charge_metric_enabled() { + self.instruments + .request_charge + .record(diagnostics.total_request_charge().value(), &attributes); + } + + if self.options.returned_rows_metric_enabled() { + if let Some(rows) = op.and_then(CosmosOperationContext::returned_item_count) { + self.instruments.returned_rows.record(rows, &attributes); + } + } + } +} + +/// Extracts the host portion of an endpoint URI for `server.address`. +fn host_of(endpoint: &str) -> Option { + url::Url::parse(endpoint) + .ok() + .and_then(|url| url.host_str().map(str::to_owned)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::CosmosStatus; + use azure_core::http::StatusCode; + use azure_data_cosmos_driver::models::ActivityId; + use opentelemetry::metrics::MeterProvider as _; + use opentelemetry_sdk::metrics::data::{AggregatedMetrics, MetricData, ResourceMetrics}; + use opentelemetry_sdk::metrics::{InMemoryMetricExporter, PeriodicReader, SdkMeterProvider}; + use std::collections::HashMap; + use std::time::Duration; + + /// A test meter provider wired to an in-memory exporter, plus the meter to + /// build a handler from. + struct TestMeter { + provider: SdkMeterProvider, + exporter: InMemoryMetricExporter, + meter: Meter, + } + + fn test_meter() -> TestMeter { + let exporter = InMemoryMetricExporter::default(); + let reader = PeriodicReader::builder(exporter.clone()).build(); + let provider = SdkMeterProvider::builder().with_reader(reader).build(); + let meter = provider.meter("test"); + TestMeter { + provider, + exporter, + meter, + } + } + + impl TestMeter { + /// Flushes and returns the collected resource metrics. + fn collect(&self) -> Vec { + self.provider.force_flush().unwrap(); + self.exporter.get_finished_metrics().unwrap() + } + } + + /// Completed context for a successful `read_item` returning HTTP 200. + fn completed(status_code: u16) -> DiagnosticsContext { + DiagnosticsContext::for_testing_completed( + ActivityId::new_uuid(), + Duration::from_millis(42), + Some(CosmosStatus::new(StatusCode::from(status_code))), + ) + } + + fn operation_context() -> CosmosOperationContext { + CosmosOperationContext::new() + .with_operation_name("read_item") + .with_database_name("my_db") + .with_container_name("my_container") + .with_server_address("my-account.documents.azure.com") + } + + fn metric_names(metrics: &[ResourceMetrics]) -> Vec { + let mut names = Vec::new(); + for rm in metrics { + for sm in rm.scope_metrics() { + for m in sm.metrics() { + names.push(m.name().to_string()); + } + } + } + names + } + + /// Returns the attributes (as a string map) and count of the first data + /// point of the `db.client.operation.duration` histogram, if present. + fn duration_point(metrics: &[ResourceMetrics]) -> Option<(HashMap, u64)> { + for rm in metrics { + for sm in rm.scope_metrics() { + for m in sm.metrics() { + if m.name() != attributes::METRIC_OPERATION_DURATION { + continue; + } + if let AggregatedMetrics::F64(MetricData::Histogram(histogram)) = m.data() { + if let Some(point) = histogram.data_points().next() { + let attrs = point + .attributes() + .map(|kv| { + (kv.key.as_str().to_string(), kv.value.as_str().into_owned()) + }) + .collect(); + return Some((attrs, point.count())); + } + } + } + } + } + None + } + + #[test] + fn stable_duration_metric_carries_expected_attributes() { + let harness = test_meter(); + let handler = CosmosMetricsHandler::with_meter(harness.meter.clone()); + + let cx = Context::new().with_value(operation_context()); + handler.handle(&completed(200), &cx); + + let metrics = harness.collect(); + let (attrs, count) = + duration_point(&metrics).expect("db.client.operation.duration should be emitted"); + + // Exactly one operation recorded. + assert_eq!(count, 1); + + // Stable attributes match semconv exactly. + assert_eq!( + attrs + .get(attributes::ATTR_DB_SYSTEM_NAME) + .map(String::as_str), + Some(attributes::DB_SYSTEM_NAME_VALUE) + ); + assert_eq!( + attrs + .get(attributes::ATTR_DB_OPERATION_NAME) + .map(String::as_str), + Some("read_item") + ); + assert_eq!( + attrs + .get(attributes::ATTR_DB_COLLECTION_NAME) + .map(String::as_str), + Some("my_container") + ); + assert_eq!( + attrs.get(attributes::ATTR_DB_NAMESPACE).map(String::as_str), + Some("my_db") + ); + assert_eq!( + attrs + .get(attributes::ATTR_DB_RESPONSE_STATUS_CODE) + .map(String::as_str), + Some("200") + ); + assert_eq!( + attrs + .get(attributes::ATTR_SERVER_ADDRESS) + .map(String::as_str), + Some("my-account.documents.azure.com") + ); + + // Success => no error.type, and no extended attributes by default. + assert!(!attrs.contains_key(attributes::ATTR_ERROR_TYPE)); + assert!(!attrs.contains_key(attributes::ATTR_CONSISTENCY_LEVEL)); + assert!(!attrs.contains_key(attributes::ATTR_SUB_STATUS_CODE)); + } + + #[test] + fn failure_sets_error_type_to_status_code() { + let harness = test_meter(); + let handler = CosmosMetricsHandler::with_meter(harness.meter.clone()); + + let cx = Context::new().with_value(operation_context()); + handler.handle(&completed(404), &cx); + + let metrics = harness.collect(); + let (attrs, _) = duration_point(&metrics).expect("duration metric should be emitted"); + + assert_eq!( + attrs + .get(attributes::ATTR_DB_RESPONSE_STATUS_CODE) + .map(String::as_str), + Some("404") + ); + assert_eq!( + attrs.get(attributes::ATTR_ERROR_TYPE).map(String::as_str), + Some("404") + ); + } + + #[test] + fn development_metrics_off_by_default() { + let harness = test_meter(); + let handler = CosmosMetricsHandler::with_meter(harness.meter.clone()); + + let cx = Context::new().with_value(operation_context().with_returned_item_count(7)); + handler.handle(&completed(200), &cx); + + let names = metric_names(&harness.collect()); + assert!(names + .iter() + .any(|n| n == attributes::METRIC_OPERATION_DURATION)); + // Only the stable metric is emitted with default options. + assert!(!names + .iter() + .any(|n| n == attributes::METRIC_OPERATION_REQUEST_CHARGE)); + assert!(!names + .iter() + .any(|n| n == attributes::METRIC_RESPONSE_RETURNED_ROWS)); + } + + #[test] + fn development_metrics_emitted_when_enabled() { + let harness = test_meter(); + let options = MetricsOptions::default() + .with_request_charge_metric(true) + .with_returned_rows_metric(true) + .with_extended_attributes(true); + let handler = CosmosMetricsHandler::with_meter_and_options(harness.meter.clone(), options); + + let cx = Context::new().with_value(operation_context().with_returned_item_count(7)); + handler.handle(&completed(200), &cx); + + let names = metric_names(&harness.collect()); + assert!(names + .iter() + .any(|n| n == attributes::METRIC_OPERATION_DURATION)); + assert!(names + .iter() + .any(|n| n == attributes::METRIC_OPERATION_REQUEST_CHARGE)); + assert!(names + .iter() + .any(|n| n == attributes::METRIC_RESPONSE_RETURNED_ROWS)); + } + + #[test] + fn optional_metrics_toggle_independently() { + // Enabling only the request-charge signal emits it — and NOT returned_rows — + // proving the per-signal toggles are independent. + let harness = test_meter(); + let options = MetricsOptions::default().with_request_charge_metric(true); + let handler = CosmosMetricsHandler::with_meter_and_options(harness.meter.clone(), options); + + let cx = Context::new().with_value(operation_context().with_returned_item_count(7)); + handler.handle(&completed(200), &cx); + + let names = metric_names(&harness.collect()); + assert!(names + .iter() + .any(|n| n == attributes::METRIC_OPERATION_DURATION)); + assert!(names + .iter() + .any(|n| n == attributes::METRIC_OPERATION_REQUEST_CHARGE)); + assert!(!names + .iter() + .any(|n| n == attributes::METRIC_RESPONSE_RETURNED_ROWS)); + } + + #[test] + fn missing_operation_context_still_emits_duration() { + // With no CosmosOperationContext on the pipeline context, the handler + // must still emit the duration metric (identity attributes just absent). + let harness = test_meter(); + let handler = CosmosMetricsHandler::with_meter(harness.meter.clone()); + + handler.handle(&completed(200), &Context::new()); + + let metrics = harness.collect(); + let (attrs, count) = duration_point(&metrics).expect("duration metric should be emitted"); + assert_eq!(count, 1); + assert_eq!( + attrs + .get(attributes::ATTR_DB_SYSTEM_NAME) + .map(String::as_str), + Some(attributes::DB_SYSTEM_NAME_VALUE) + ); + assert!(!attrs.contains_key(attributes::ATTR_DB_OPERATION_NAME)); + } + + #[test] + fn no_meter_provider_is_a_noop() { + // Building from the global meter with no provider registered yields a + // no-op meter; recording must not panic (the exporter-absent path). + let handler = CosmosMetricsHandler::new(); + handler.handle(&completed(200), &Context::new()); + } + + #[test] + fn host_of_extracts_host() { + assert_eq!( + host_of("https://my-account.documents.azure.com:443/dbs/x"), + Some("my-account.documents.azure.com".to_string()) + ); + assert_eq!(host_of("not a url"), None); + } +} diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/instruments.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/instruments.rs new file mode 100644 index 00000000000..78e60c598e0 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/instruments.rs @@ -0,0 +1,63 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! The OpenTelemetry instruments the metrics handler records into. +//! +//! Instruments are built once, from a single [`Meter`], when the handler is +//! constructed, and then cheaply shared (each instrument is internally +//! 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 crate::diagnostics::metrics::attributes; + +/// The full set of Cosmos metric instruments, created from one [`Meter`]. +/// +/// 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. +/// 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. +#[derive(Clone)] +pub(crate) struct Instruments { + /// Stable: `db.client.operation.duration` (seconds). + pub(crate) operation_duration: Histogram, + + /// Development: `azure.cosmosdb.client.operation.request_charge` (RU). + pub(crate) request_charge: Histogram, + + /// Development: `db.client.response.returned_rows` (rows). + pub(crate) returned_rows: Histogram, +} + +impl Instruments { + /// Builds every Cosmos instrument from `meter`. + pub(crate) fn new(meter: &Meter) -> Self { + let operation_duration = meter + .f64_histogram(attributes::METRIC_OPERATION_DURATION) + .with_unit(attributes::UNIT_SECONDS) + .with_description("Total client-observed duration of a Cosmos DB operation.") + .build(); + + let request_charge = meter + .f64_histogram(attributes::METRIC_OPERATION_REQUEST_CHARGE) + .with_unit(attributes::UNIT_REQUEST_UNIT) + .with_description("Request charge (RU) consumed by a Cosmos DB operation.") + .build(); + + let returned_rows = meter + .u64_histogram(attributes::METRIC_RESPONSE_RETURNED_ROWS) + .with_unit(attributes::UNIT_ROW) + .with_description("Number of rows/items returned by a Cosmos DB operation.") + .build(); + + Self { + operation_duration, + request_charge, + returned_rows, + } + } +} diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/mod.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/mod.rs new file mode 100644 index 00000000000..ea79f106736 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/mod.rs @@ -0,0 +1,41 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! OpenTelemetry metrics for Cosmos DB operations (feature `metrics`). +//! +//! This module provides [`CosmosMetricsHandler`], a +//! [`DiagnosticsHandler`](crate::diagnostics::DiagnosticsHandler) that maps each +//! completed operation's [`DiagnosticsContext`](crate::diagnostics::DiagnosticsContext) +//! to OpenTelemetry metrics following the database-client semantic conventions. +//! +//! The whole module is compiled only when the `metrics` Cargo feature is +//! enabled, so with the feature off there is no metrics code at all. With the +//! feature on but no meter provider registered, OpenTelemetry's global meter is a +//! no-op. +//! +//! Metrics are emitted through the [`opentelemetry`] metrics API. +//! +//! # Emitted metrics +//! +//! - **Stable (always on):** `db.client.operation.duration` (histogram, seconds). +//! - **Optional (each opt-in via [`MetricsOptions`]):** +//! `azure.cosmosdb.client.operation.request_charge` and +//! `db.client.response.returned_rows`. +//! +//! # Operation-scope identity +//! +//! `db.operation.name`, `db.collection.name`, and `db.namespace` are not carried +//! on the driver's [`DiagnosticsContext`](crate::diagnostics::DiagnosticsContext); +//! they are supplied by the SDK through a +//! [`CosmosOperationContext`](crate::diagnostics::CosmosOperationContext) stored on the pipeline +//! [`Context`](azure_core::http::Context). The handler reads whichever fields are +//! present and omits the rest, so it degrades gracefully. + +pub mod attributes; + +mod handler; +mod instruments; +mod options; + +pub use handler::CosmosMetricsHandler; +pub use options::MetricsOptions; diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/options.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/options.rs new file mode 100644 index 00000000000..2fddee6fc6a --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/metrics/options.rs @@ -0,0 +1,95 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Configuration for the Cosmos DB metrics handler. + +/// Controls which optional metrics and attributes the +/// [`CosmosMetricsHandler`](super::CosmosMetricsHandler) emits. +/// +/// The options are per-signal: each optional metric and the extended attribute +/// set is toggled on its own, focused on *what* is emitted rather than on a +/// "preview/development" tier. Everything optional is **off by default** — only +/// the stable, low-cardinality operation-duration metric is emitted — so enabling +/// metrics never silently multiplies a backend's time-series count (design +/// decision **D7**); each additional signal is an explicit opt-in. +/// +/// # Examples +/// +/// ``` +/// use azure_data_cosmos::diagnostics::MetricsOptions; +/// +/// // Stable metric only (default). +/// let stable = MetricsOptions::default(); +/// assert!(!stable.request_charge_metric_enabled()); +/// assert!(!stable.returned_rows_metric_enabled()); +/// assert!(!stable.extended_attributes_enabled()); +/// +/// // Opt into just the request-charge metric. +/// let charge = MetricsOptions::default().with_request_charge_metric(true); +/// assert!(charge.request_charge_metric_enabled()); +/// assert!(!charge.returned_rows_metric_enabled()); +/// +/// // Opt into everything. +/// let full = MetricsOptions::default() +/// .with_request_charge_metric(true) +/// .with_returned_rows_metric(true) +/// .with_extended_attributes(true); +/// assert!(full.request_charge_metric_enabled()); +/// assert!(full.returned_rows_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, + extended_attributes: bool, +} + +impl MetricsOptions { + /// Returns options with every optional signal disabled — the stable + /// operation-duration metric only. + pub fn new() -> Self { + Self::default() + } + + /// Enables (or disables) the `azure.cosmosdb.client.operation.request_charge` + /// histogram (request units consumed per operation). Off by default. + #[must_use] + pub fn with_request_charge_metric(mut self, enabled: bool) -> Self { + self.request_charge_metric = enabled; + self + } + + /// Enables (or disables) the `db.client.response.returned_rows` histogram + /// (rows/items returned per operation). Off by default. + #[must_use] + pub fn with_returned_rows_metric(mut self, enabled: bool) -> Self { + self.returned_rows_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 + /// default. + #[must_use] + pub fn with_extended_attributes(mut self, enabled: bool) -> Self { + self.extended_attributes = enabled; + self + } + + /// Whether the request-charge histogram is emitted. + pub fn request_charge_metric_enabled(&self) -> bool { + self.request_charge_metric + } + + /// Whether the returned-rows histogram is emitted. + pub fn returned_rows_metric_enabled(&self) -> bool { + self.returned_rows_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/src/diagnostics/mod.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/mod.rs new file mode 100644 index 00000000000..e785d31cfdc --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/mod.rs @@ -0,0 +1,81 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Per-operation diagnostics surfaced by the Cosmos DB SDK. +//! +//! Every fallible Cosmos operation produces a [`DiagnosticsContext`] capturing +//! request tracking, retries, regions contacted, and other observability +//! signals from the request pipeline. The context is reachable from +//! [`CosmosError`](crate::CosmosError) on failure, and from the +//! [`FeedPage`](crate::feed::FeedPage), [`ItemResponse`](crate::models::ItemResponse), and +//! similar response wrappers on success. +//! +//! The SDK also exposes an emission extension point on top of that context: a +//! [`DiagnosticsHandler`] receives each operation's completed +//! [`DiagnosticsContext`], and an ordered [`DiagnosticsHandlerChain`] invokes +//! registered handlers once per operation at completion. Register handlers via +//! [`CosmosClientBuilder::with_diagnostics_handler`](crate::CosmosClientBuilder::with_diagnostics_handler). +//! With no handlers registered the chain does nothing beyond checking whether a +//! handler is present. +//! +//! A built-in OpenTelemetry metrics handler, +//! [`CosmosMetricsHandler`], is available behind +//! the off-by-default `metrics` feature. It emits the stable +//! `db.client.operation.duration` histogram (and, opt-in, development-tier +//! metrics) from each completed context. +//! +//! Two further built-in handlers layer telemetry on top of the chain, both driven by +//! tail-based sampling against [`DiagnosticsThresholds`] (emit only for failed or +//! threshold-breaching operations): +//! +//! - [`SamplingLogHandler`] — a composable wrapper that applies the sampling +//! gate plus a per-window rate limit and delegates emission to an inner +//! handler, defaulting to a [`TracingLogHandler`] that logs a compact +//! diagnostics line through the [`tracing`](https://docs.rs/tracing) ecosystem. +//! - [`CosmosTracingHandler`] — emits a backdated OpenTelemetry span tree +//! (behind the off-by-default `distributed_tracing` feature), rate-limited so +//! an error storm can't overwhelm exporters. + +// ========================================================================= +// Public API +// ========================================================================= + +#[doc(inline)] +pub use azure_data_cosmos_driver::diagnostics::{ + DiagnosticsContext, ThresholdBreach, TransportKind, +}; +#[doc(inline)] +pub use azure_data_cosmos_driver::DiagnosticsThresholds; +pub use handler::{DiagnosticsHandler, DiagnosticsHandlerChain}; +pub use logging::{SamplingLogHandler, TracingLogHandler}; +pub use operation_context::CosmosOperationContext; +pub use rate_limiter::RateLimiterConfig; +#[cfg(feature = "distributed_tracing")] +pub use tracing::CosmosTracingHandler; + +#[cfg(feature = "metrics")] +pub use metrics::{CosmosMetricsHandler, MetricsOptions}; + +// ========================================================================= +// Internal modules +// ========================================================================= + +// Shared semantic-convention attribute-name literals, used by the metrics and +// distributed-tracing handlers (single source of truth). Only needed when one of +// those feature-gated handlers is compiled. +#[cfg(any(feature = "metrics", feature = "distributed_tracing"))] +pub(crate) mod attributes; + +mod handler; +mod logging; +mod operation_context; +mod reason; +// Count-per-interval rate limiter shared by the sampling handlers (logging and, +// when enabled, tracing) so they can bound emission under an error storm. +pub(crate) mod rate_limiter; + +#[cfg(feature = "metrics")] +pub mod metrics; + +#[cfg(feature = "distributed_tracing")] +mod tracing; diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/operation_context.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/operation_context.rs new file mode 100644 index 00000000000..c1fd03072b6 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/operation_context.rs @@ -0,0 +1,150 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! SDK-supplied, operation-scope identity carried to diagnostics handlers. +//! +//! The driver's [`DiagnosticsContext`](crate::diagnostics::DiagnosticsContext) +//! records what happened on the wire, but the caller-facing identity of an +//! operation — its name, database, and container — is known only to the SDK. +//! [`CosmosOperationContext`] carries that identity to the handler chain through +//! the pipeline [`Context`](azure_core::http::Context), so both the metrics and +//! tracing handlers can emit correct operation-scope attributes. +//! +//! This type is always compiled (independent of the `metrics` / +//! `distributed_tracing` features) because the SDK populates it on every completed +//! operation; the feature-gated handlers simply read whichever fields are set. + +use std::borrow::Cow; + +use azure_core::fmt::SafeDebug; + +/// SDK-supplied, operation-scope identity for a single Cosmos operation. +/// +/// The driver's [`DiagnosticsContext`](crate::diagnostics::DiagnosticsContext) +/// captures what happened on the wire (status, duration, regions, request +/// charge) but not the caller-facing identity of the operation — the operation +/// name, database, and container are known only to the SDK. This type carries +/// that identity to the diagnostics handlers via the pipeline +/// [`Context`](azure_core::http::Context): +/// +/// ``` +/// use azure_core::http::Context; +/// use azure_data_cosmos::diagnostics::CosmosOperationContext; +/// +/// let op = CosmosOperationContext::new() +/// .with_operation_name("read_item") +/// .with_database_name("my_db") +/// .with_container_name("my_container"); +/// let cx = Context::new().with_value(op); +/// assert_eq!( +/// cx.value::().and_then(|o| o.operation_name()), +/// Some("read_item"), +/// ); +/// ``` +/// +/// Every field is optional; a handler emits an attribute only for the fields +/// that are set. All setters accept anything convertible into a +/// `Cow<'static, str>`, so canonical static operation names (e.g. `"read_item"`) +/// are stored without allocating. +#[derive(Clone, Default, SafeDebug, PartialEq, Eq)] +pub struct CosmosOperationContext { + operation_name: Option>, + database_name: Option>, + container_name: Option>, + server_address: Option>, + consistency_level: Option>, + connection_mode: Option>, + returned_item_count: Option, +} + +impl CosmosOperationContext { + /// Creates an empty operation context. + pub fn new() -> Self { + Self::default() + } + + /// Sets the canonical operation name (`db.operation.name`), e.g. `read_item`. + #[must_use] + pub fn with_operation_name(mut self, name: impl Into>) -> Self { + self.operation_name = Some(name.into()); + self + } + + /// Sets the database name (`db.namespace`). + #[must_use] + pub fn with_database_name(mut self, name: impl Into>) -> Self { + self.database_name = Some(name.into()); + self + } + + /// Sets the container name (`db.collection.name`). + #[must_use] + pub fn with_container_name(mut self, name: impl Into>) -> Self { + self.container_name = Some(name.into()); + self + } + + /// Sets the server address (`server.address`), overriding the value the + /// handler would otherwise derive from the contacted endpoint. + #[must_use] + pub fn with_server_address(mut self, address: impl Into>) -> Self { + self.server_address = Some(address.into()); + self + } + + /// Sets the effective consistency level (development attribute). + #[must_use] + pub fn with_consistency_level(mut self, level: impl Into>) -> Self { + self.consistency_level = Some(level.into()); + self + } + + /// Sets the connection mode, e.g. `gateway` or `direct` (development attribute). + #[must_use] + pub fn with_connection_mode(mut self, mode: impl Into>) -> Self { + self.connection_mode = Some(mode.into()); + self + } + + /// Sets the number of items returned (feeds `db.client.response.returned_rows`). + #[must_use] + pub fn with_returned_item_count(mut self, count: u64) -> Self { + self.returned_item_count = Some(count); + self + } + + /// The canonical operation name, if set. + pub fn operation_name(&self) -> Option<&str> { + self.operation_name.as_deref() + } + + /// The database name, if set. + pub fn database_name(&self) -> Option<&str> { + self.database_name.as_deref() + } + + /// The container name, if set. + pub fn container_name(&self) -> Option<&str> { + self.container_name.as_deref() + } + + /// The server-address override, if set. + pub fn server_address(&self) -> Option<&str> { + self.server_address.as_deref() + } + + /// The effective consistency level, if set. + pub fn consistency_level(&self) -> Option<&str> { + self.consistency_level.as_deref() + } + + /// The connection mode, if set. + pub fn connection_mode(&self) -> Option<&str> { + self.connection_mode.as_deref() + } + + /// The number of items returned, if set. + pub fn returned_item_count(&self) -> Option { + self.returned_item_count + } +} diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/rate_limiter.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/rate_limiter.rs new file mode 100644 index 00000000000..f171abbced8 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/rate_limiter.rs @@ -0,0 +1,346 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! A count-per-interval rate limiter with a failure reserve. +//! +//! [`RateLimiter`] caps how many diagnostics emissions are allowed per fixed +//! time window during a storm, while always letting a bounded number of +//! *failures* through even after the normal budget is exhausted. When emissions +//! are suppressed within a window, the limiter surfaces a single +//! "suppressed N until reset" notice — exactly once per window — the first time +//! it is consulted after that window ends. +//! +//! The limiter is deterministic and clock-injectable ([`RateLimiter::check`] +//! takes the current [`Instant`]), so its behavior can be unit-tested without +//! sleeping. +//! +//! It is shared by the built-in emission handlers (sampled logging and, when +//! enabled, distributed tracing) so they can all bound their output under an +//! error storm. +//! +//! ## Storm fast path +//! +//! Once a window is fully saturated — the normal budget *and* the failure +//! reserve are both exhausted — no further emission can be admitted until the +//! window rolls over. In that state [`RateLimiter::check`] takes a lock-free fast +//! path (relaxed atomics) that returns "suppress" without contending the mutex, +//! so a 10k-errors/sec storm doesn't serialize every operation task on one lock. +//! The fast path's suppressed count is folded back into the exact total on the +//! next locked call, so the once-per-window notice stays accurate. + +use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering}; +use std::sync::Mutex; +use std::time::{Duration, Instant}; + +/// Default number of sampled diagnostics lines allowed per window (~100/min). +pub(crate) const DEFAULT_MAX_PER_WINDOW: u32 = 100; + +/// Default rate-limiting window length. +pub(crate) const DEFAULT_WINDOW: Duration = Duration::from_secs(60); + +/// Default number of failures always allowed per window, even past the cap. +pub(crate) const DEFAULT_FAILURE_RESERVE: u32 = 10; + +/// Configuration for the rate limiting applied by the built-in sampling handlers +/// ([`SamplingLogHandler`](crate::diagnostics::SamplingLogHandler) and the +/// distributed-tracing handler). +#[derive(Clone, Copy, Debug)] +#[non_exhaustive] +pub struct RateLimiterConfig { + /// Maximum number of emissions allowed per [`window`](Self::window). + pub max_per_window: u32, + /// Length of a rate-limiting window. + pub window: Duration, + /// Number of failures permitted per window *in addition to* the normal cap, + /// so a bounded number of failures is always emitted during a storm. + pub failure_reserve: u32, +} + +impl Default for RateLimiterConfig { + fn default() -> Self { + Self { + max_per_window: DEFAULT_MAX_PER_WINDOW, + window: DEFAULT_WINDOW, + failure_reserve: DEFAULT_FAILURE_RESERVE, + } + } +} + +/// The outcome of a [`RateLimiter::check`]. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct LimitDecision { + /// Whether the caller should emit for this event. + pub emit: bool, + /// When `Some(n)`, the caller should emit exactly one "suppressed `n` until + /// reset" notice for the window that just ended. + pub suppression_notice: Option, +} + +#[derive(Debug)] +struct State { + window_start: Instant, + /// Total emissions this window (including reserve failures). + emitted: u32, + /// Failures emitted this window (counts toward the failure reserve). + failures_emitted: u32, + /// Emissions suppressed this window (exact; includes folded fast-path counts). + suppressed: u32, +} + +/// A count-per-interval limiter with a failure reserve. Cheap to share behind an +/// `Arc`; internally synchronized. +#[derive(Debug)] +pub(crate) struct RateLimiter { + config: RateLimiterConfig, + state: Mutex, + /// Monotonic anchor for converting `Instant`s to the `u64` nanos used by the + /// lock-free fast path. + base: Instant, + /// `true` when the current window is fully saturated (normal budget and + /// failure reserve both exhausted), so nothing more can be admitted until it + /// rolls over. Read on the fast path with relaxed ordering. + over_budget: AtomicBool, + /// Nanoseconds-since-[`base`](Self::base) at which the current window ends. + /// The fast path only trusts [`over_budget`](Self::over_budget) while `now` + /// is still before this instant. + window_end_nanos: AtomicU64, + /// Suppressions taken on the lock-free fast path, folded into + /// [`State::suppressed`] on the next locked call so the notice stays exact. + suppressed_fast: AtomicU32, +} + +impl RateLimiter { + /// Creates a limiter with the given configuration, anchored at "now". + /// + /// Because [`RateLimiterConfig`] is publicly constructible with mutable + /// fields, a caller can supply `window = Duration::ZERO`, which would roll the + /// window over on every call and bypass the cap. Normalize a zero window to + /// the default so the limiter always makes forward progress. + pub(crate) fn new(mut config: RateLimiterConfig) -> Self { + if config.window.is_zero() { + config.window = DEFAULT_WINDOW; + } + let base = Instant::now(); + Self { + config, + state: Mutex::new(State { + window_start: base, + emitted: 0, + failures_emitted: 0, + suppressed: 0, + }), + base, + over_budget: AtomicBool::new(false), + window_end_nanos: AtomicU64::new(config.window.as_nanos() as u64), + suppressed_fast: AtomicU32::new(0), + } + } + + /// Records an emission attempt at `now` and returns whether it is allowed. + /// + /// `is_failure` marks the event as a failure, which may be admitted from the + /// failure reserve once the normal per-window budget is exhausted. + /// + /// When the call is the first after a window rolled over and the previous + /// window suppressed anything, the returned [`LimitDecision::suppression_notice`] + /// carries that window's suppressed count so the caller can emit a single + /// notice. + pub(crate) fn check(&self, is_failure: bool, now: Instant) -> LimitDecision { + // Lock-free fast path: if the window is known-saturated and we're still + // inside it, suppress without taking the mutex. Relaxed ordering is fine + // — a storm only needs a definitive "skip", and the exact accounting is + // reconciled on the next locked call. + let now_nanos = now.saturating_duration_since(self.base).as_nanos() as u64; + if self.over_budget.load(Ordering::Relaxed) + && now_nanos < self.window_end_nanos.load(Ordering::Relaxed) + { + self.suppressed_fast.fetch_add(1, Ordering::Relaxed); + return LimitDecision { + emit: false, + suppression_notice: None, + }; + } + + let mut state = self.state.lock().unwrap(); + + // Reconcile any fast-path suppressions into the exact per-window count + // before we might roll the window over (so they are reflected in the + // notice for the window they belonged to). + let fast = self.suppressed_fast.swap(0, Ordering::Relaxed); + state.suppressed = state.suppressed.saturating_add(fast); + + let mut suppression_notice = None; + if now.saturating_duration_since(state.window_start) >= self.config.window { + if state.suppressed > 0 { + suppression_notice = Some(state.suppressed); + } + state.window_start = now; + state.emitted = 0; + state.failures_emitted = 0; + state.suppressed = 0; + } + + let within_normal_budget = state.emitted < self.config.max_per_window; + // Past the normal cap, a bounded number of failures may still pass from + // the reserve. Only reserve admissions count against `failures_emitted`, + // so the reserve is genuinely *in addition to* the cap (a cap filled + // with failures must not consume it). + let from_reserve = !within_normal_budget + && is_failure + && state.failures_emitted < self.config.failure_reserve; + let allowed = within_normal_budget || from_reserve; + + let decision = if allowed { + state.emitted += 1; + if from_reserve { + state.failures_emitted += 1; + } + LimitDecision { + emit: true, + suppression_notice, + } + } else { + state.suppressed += 1; + LimitDecision { + emit: false, + suppression_notice, + } + }; + + // Refresh the fast-path atomics from the authoritative state. The window + // is saturated only when neither the normal budget nor the failure + // reserve can admit anything more. + let saturated = state.emitted >= self.config.max_per_window + && state.failures_emitted >= self.config.failure_reserve; + self.over_budget.store(saturated, Ordering::Relaxed); + let window_end = state + .window_start + .saturating_duration_since(self.base) + .saturating_add(self.config.window) + .as_nanos() as u64; + self.window_end_nanos.store(window_end, Ordering::Relaxed); + + decision + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn caps_emissions_within_a_window() { + let limiter = RateLimiter::new(RateLimiterConfig { + max_per_window: 5, + window: Duration::from_millis(100), + failure_reserve: 0, + }); + let t0 = Instant::now(); + + let mut emitted = 0; + let mut notices = 0; + for _ in 0..50 { + let d = limiter.check(false, t0); + if d.emit { + emitted += 1; + } + if d.suppression_notice.is_some() { + notices += 1; + } + } + + // Only the cap is admitted; the rest are suppressed and no notice has + // fired yet (the window has not rolled over). + assert_eq!(emitted, 5); + assert_eq!(notices, 0); + } + + #[test] + fn emits_exactly_one_suppression_notice_per_window() { + let limiter = RateLimiter::new(RateLimiterConfig { + max_per_window: 5, + window: Duration::from_millis(100), + failure_reserve: 0, + }); + let t0 = Instant::now(); + for _ in 0..50 { + limiter.check(false, t0); + } + + // First check after the window ends carries the suppressed count once. + let t1 = t0 + Duration::from_millis(150); + let first = limiter.check(false, t1); + assert!(first.emit, "new window admits again"); + assert_eq!(first.suppression_notice, Some(45), "50 - 5 suppressed"); + + // Subsequent checks in the new window do not repeat the notice. + let second = limiter.check(false, t1); + assert_eq!(second.suppression_notice, None); + } + + #[test] + fn failures_admitted_from_reserve_past_cap() { + let limiter = RateLimiter::new(RateLimiterConfig { + max_per_window: 2, + window: Duration::from_secs(60), + failure_reserve: 3, + }); + let t = Instant::now(); + + // Fill the normal cap with successes. + assert!(limiter.check(false, t).emit); + assert!(limiter.check(false, t).emit); + // Further successes are suppressed. + assert!(!limiter.check(false, t).emit); + // Failures still pass, up to the reserve of 3. + assert!(limiter.check(true, t).emit); + assert!(limiter.check(true, t).emit); + assert!(limiter.check(true, t).emit); + // Reserve exhausted: further failures are suppressed. + assert!(!limiter.check(true, t).emit); + } + + #[test] + fn reserve_is_additional_to_a_cap_filled_with_failures() { + let limiter = RateLimiter::new(RateLimiterConfig { + max_per_window: 2, + window: Duration::from_secs(60), + failure_reserve: 3, + }); + let t = Instant::now(); + + // Fill the normal cap with failures. These are admitted by the normal + // budget, so they must NOT consume the failure reserve. + assert!(limiter.check(true, t).emit); + assert!(limiter.check(true, t).emit); + // The full reserve of 3 is still available for post-cap failures. + assert!(limiter.check(true, t).emit); + assert!(limiter.check(true, t).emit); + assert!(limiter.check(true, t).emit); + // Reserve now exhausted. + assert!(!limiter.check(true, t).emit); + } + + #[test] + fn fast_path_suppressions_are_counted_in_the_window_notice() { + // Once saturated, the lock-free fast path suppresses without the mutex; + // those suppressions must still be reflected in the next window's notice. + let limiter = RateLimiter::new(RateLimiterConfig { + max_per_window: 2, + window: Duration::from_millis(100), + failure_reserve: 0, + }); + let t0 = Instant::now(); + + // Two admitted, then the window is saturated (reserve 0) and the fast + // path suppresses the remaining 18. + for _ in 0..20 { + limiter.check(false, t0); + } + + // Roll over: the notice must count all 18 suppressed (fast path + locked). + let t1 = t0 + Duration::from_millis(150); + let rolled = limiter.check(false, t1); + assert_eq!(rolled.suppression_notice, Some(18), "20 - 2 admitted"); + } +} diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/reason.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/reason.rs new file mode 100644 index 00000000000..9dac89ac20b --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/reason.rs @@ -0,0 +1,85 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Why a completed operation was sampled for emission. +//! +//! The sampling handlers emit a diagnostic when an operation *fails* or *crosses +//! a threshold*. [`EmitReason`] captures which of those it was — and, for a +//! threshold breach, which specific bound — so the emitted diagnostic can record +//! *why* it was surfaced, not just that it was. + +use azure_data_cosmos_driver::diagnostics::{DiagnosticsContext, ThresholdBreach}; +use azure_data_cosmos_driver::DiagnosticsThresholds; + +use crate::diagnostics::CosmosOperationContext; + +/// The reason a completed operation passed the tail-based sampling gate. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum EmitReason { + /// The operation failed. + Failure, + /// The operation succeeded but crossed a sampling threshold. + Threshold(ThresholdBreach), +} + +impl EmitReason { + /// A stable, low-cardinality identifier suitable for a log field or span + /// attribute value. + pub(crate) fn as_str(self) -> &'static str { + match self { + EmitReason::Failure => "failure", + EmitReason::Threshold(breach) => match breach { + ThresholdBreach::PointLatency => "point_latency", + ThresholdBreach::NonPointLatency => "non_point_latency", + ThresholdBreach::RequestCharge => "request_charge", + // `ThresholdBreach` is `#[non_exhaustive]`; keep a stable + // fallback if a new breach kind is added upstream. + _ => "threshold", + }, + } + } + + /// Computes why `diagnostics` would be sampled against `thresholds`: a failure + /// takes precedence, otherwise the specific threshold crossed. Returns `None` + /// when the operation is neither a failure nor a threshold breach. + /// + /// `op` supplies the SDK operation identity used to pick the point vs + /// non-point latency threshold, mirroring the sampling gate. + pub(crate) fn of( + diagnostics: &DiagnosticsContext, + thresholds: &DiagnosticsThresholds, + op: Option<&CosmosOperationContext>, + ) -> Option { + if diagnostics.is_failure() { + return Some(EmitReason::Failure); + } + diagnostics + .threshold_breach_for( + thresholds, + op.and_then(CosmosOperationContext::operation_name), + ) + .map(EmitReason::Threshold) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn as_str_is_stable() { + assert_eq!(EmitReason::Failure.as_str(), "failure"); + assert_eq!( + EmitReason::Threshold(ThresholdBreach::PointLatency).as_str(), + "point_latency" + ); + assert_eq!( + EmitReason::Threshold(ThresholdBreach::NonPointLatency).as_str(), + "non_point_latency" + ); + assert_eq!( + EmitReason::Threshold(ThresholdBreach::RequestCharge).as_str(), + "request_charge" + ); + } +} diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/handler.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/handler.rs new file mode 100644 index 00000000000..3f6a9ef7df3 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/handler.rs @@ -0,0 +1,173 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! The [`CosmosTracingHandler`]: a tail-sampled OpenTelemetry tracing handler. + +use std::time::{Instant, SystemTime}; + +use azure_core::http::Context; +use azure_data_cosmos_driver::{diagnostics::DiagnosticsContext, DiagnosticsThresholds}; +use opentelemetry::global; + +use super::span_builder::emit_backdated_span_tree; +use crate::diagnostics::rate_limiter::{RateLimiter, RateLimiterConfig}; +use crate::diagnostics::reason::EmitReason; +use crate::diagnostics::{CosmosOperationContext, DiagnosticsHandler}; + +/// The instrumentation scope name used for the Cosmos tracer. +const TRACER_NAME: &str = "azure_data_cosmos"; + +/// `tracing` target for the "suppressed N span tree(s)" notice. +const SUPPRESSED_TARGET: &str = "azure_data_cosmos::diagnostics::tracing_suppressed"; + +/// A [`DiagnosticsHandler`] that emits a backdated OpenTelemetry span tree for +/// operations that fail or cross a sampling threshold. +/// +/// This handler implements **tail-based sampling**: it inspects the *completed* +/// [`DiagnosticsContext`] and only emits when +/// [`should_emit`](Self::should_emit) is satisfied — i.e. the operation failed +/// or breached one of the configured [`DiagnosticsThresholds`]. A fast, successful +/// point read therefore produces **no** span at all, keeping the common path +/// free of tracing overhead. +/// +/// When it does emit, it reconstructs the operation as a root span with one child +/// span per retained attempt, each backdated to the time the work actually +/// happened (see the module docs). +/// +/// Register it with +/// [`CosmosClientBuilder::with_diagnostics_handler`](crate::CosmosClientBuilder::with_diagnostics_handler). +/// Spans are emitted through the globally-installed OpenTelemetry tracer provider, +/// resolved lazily on each sampled emission — so a provider installed *after* the +/// handler (or client) is constructed is still picked up. With no provider +/// installed, emission is a no-op. +/// +/// Span emission is **rate-limited** per window (≈100/min by default, with a +/// bounded failure reserve) so an error storm can't reconstruct millions of span +/// trees per second and overwhelm exporters. When trees are suppressed, a single +/// "suppressed N" warning is emitted per window on the +/// `azure_data_cosmos::diagnostics::tracing_suppressed` target. Tune it with +/// [`with_thresholds_and_rate_limit`](Self::with_thresholds_and_rate_limit). +/// +/// # Examples +/// +/// ``` +/// use std::sync::Arc; +/// use azure_data_cosmos::diagnostics::CosmosTracingHandler; +/// +/// let handler = Arc::new(CosmosTracingHandler::new()); +/// // let client = CosmosClient::builder(endpoint, credential) +/// // .with_diagnostics_handler(handler) +/// // .build()?; +/// # let _ = handler; +/// ``` +pub struct CosmosTracingHandler { + thresholds: DiagnosticsThresholds, + limiter: RateLimiter, +} + +impl CosmosTracingHandler { + /// Creates a handler using the default sampling thresholds and default span + /// emission rate limiting (~100/min). + pub fn new() -> Self { + Self::with_thresholds(DiagnosticsThresholds::default()) + } + + /// Creates a handler using the supplied sampling thresholds and default span + /// emission rate limiting. + pub fn with_thresholds(thresholds: DiagnosticsThresholds) -> Self { + Self::with_thresholds_and_rate_limit(thresholds, RateLimiterConfig::default()) + } + + /// Creates a handler using the supplied sampling thresholds and span emission + /// rate-limiter configuration. + /// + /// The rate limit bounds how many span trees are reconstructed per window + /// across all operations, so a failure storm can't overwhelm CPU/exporters. + pub fn with_thresholds_and_rate_limit( + thresholds: DiagnosticsThresholds, + rate_limit: RateLimiterConfig, + ) -> Self { + Self { + thresholds, + limiter: RateLimiter::new(rate_limit), + } + } + + /// Returns the sampling thresholds this handler applies. + pub fn thresholds(&self) -> &DiagnosticsThresholds { + &self.thresholds + } + + /// Returns whether the given completed context should emit a span, per the + /// tail-based sampling policy: emit iff the operation failed or crossed a + /// threshold. + pub fn should_emit(&self, diagnostics: &DiagnosticsContext) -> bool { + should_emit_span(diagnostics, &self.thresholds, None) + } +} + +impl Default for CosmosTracingHandler { + fn default() -> Self { + Self::new() + } +} + +impl DiagnosticsHandler for CosmosTracingHandler { + fn handle(&self, diagnostics: &DiagnosticsContext, cx: &Context<'_>) { + let op = cx.value::(); + if !should_emit_span(diagnostics, &self.thresholds, op) { + return; + } + + // Bound span reconstruction across operations so an error storm can't + // emit millions of trees per second. This is a synchronous, per-failure + // cost, so the limit applies before we build anything. + let decision = self.limiter.check(diagnostics.is_failure(), Instant::now()); + if let Some(suppressed) = decision.suppression_notice { + tracing::warn!( + target: SUPPRESSED_TARGET, + suppressed, + "cosmos diagnostics: suppressed {suppressed} tracing span tree(s) until window reset" + ); + } + if !decision.emit { + return; + } + + // Resolve the global tracer lazily, on the (rare) sampled emission path, + // rather than caching it at construction. `global::tracer` binds to + // whatever provider is installed *now*; caching it in the handler would + // permanently capture the no-op default whenever the handler is built + // before `global::set_tracer_provider`, silently dropping every sampled + // span even after a provider is installed later. + let tracer = global::tracer(TRACER_NAME); + let reason = EmitReason::of(diagnostics, &self.thresholds, op).map(EmitReason::as_str); + emit_backdated_span_tree( + &tracer, + diagnostics, + op, + reason, + Instant::now(), + SystemTime::now(), + ); + } +} + +/// The tail-based sampling decision: emit a span iff the operation completed and +/// either failed or crossed one of the sampling thresholds. +/// +/// `op` supplies the SDK-side operation identity so the threshold classifier +/// can distinguish point from non-point operations; production driver contexts +/// do not carry the operation name. +pub(crate) fn should_emit_span( + diagnostics: &DiagnosticsContext, + thresholds: &DiagnosticsThresholds, + op: Option<&CosmosOperationContext>, +) -> bool { + diagnostics.is_completed() + && (diagnostics.is_failure() + || diagnostics.is_threshold_violated_for( + thresholds, + op.and_then(CosmosOperationContext::operation_name), + )) +} diff --git a/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/mod.rs b/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/mod.rs new file mode 100644 index 00000000000..f678868c2d9 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/mod.rs @@ -0,0 +1,341 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Tail-sampled OpenTelemetry tracing for Cosmos DB operations. +//! +//! This module is gated behind the off-by-default `distributed_tracing` feature. It +//! provides [`CosmosTracingHandler`], a [`DiagnosticsHandler`](crate::diagnostics::DiagnosticsHandler) +//! that reconstructs a **backdated** span tree — one operation span plus one child +//! per retained attempt — for operations selected by tail-based sampling (failures +//! and threshold breaches). Fast, successful operations emit nothing. + +mod handler; +mod span_builder; + +pub use handler::CosmosTracingHandler; + +#[cfg(test)] +mod tests { + use std::time::{Duration, Instant, SystemTime}; + + use azure_core::http::StatusCode; + use azure_data_cosmos_driver::diagnostics::{DiagnosticsContext, RequestDiagnostics}; + use azure_data_cosmos_driver::models::{ActivityId, RequestCharge}; + use azure_data_cosmos_driver::options::Region; + use azure_data_cosmos_driver::{CosmosStatus, DiagnosticsThresholds}; + use opentelemetry::trace::{SpanId, TracerProvider}; + use opentelemetry_sdk::trace::{in_memory_exporter::InMemorySpanExporter, SdkTracerProvider}; + + use super::handler::should_emit_span; + use super::span_builder::emit_backdated_span_tree; + use crate::diagnostics::attributes; + use crate::diagnostics::CosmosOperationContext; + + /// Builds a completed context: `duration` long, final `status`, and one + /// synthetic attempt per `(offset_ms, dur_ms, status)` triple. `offset_ms` is + /// how far before `anchor` the attempt started. + fn context( + duration: Duration, + status: Option, + operation_name: Option<&str>, + attempts: &[(u64, u64, CosmosStatus)], + anchor: Instant, + ) -> DiagnosticsContext { + let requests = attempts + .iter() + .map(|(offset_ms, dur_ms, req_status)| { + let started = anchor - Duration::from_millis(*offset_ms); + let completed = started + Duration::from_millis(*dur_ms); + RequestDiagnostics::for_testing( + "https://acct.documents.azure.com:443/", + Some(Region::new("West US 2")), + *req_status, + RequestCharge::new(2.0), + started, + completed, + ) + }) + .collect(); + DiagnosticsContext::for_testing_with_requests( + ActivityId::new_uuid(), + duration, + status, + operation_name, + requests, + ) + } + + fn exportable() -> (SdkTracerProvider, InMemorySpanExporter) { + let exporter = InMemorySpanExporter::default(); + let provider = SdkTracerProvider::builder() + .with_simple_exporter(exporter.clone()) + .build(); + (provider, exporter) + } + + #[test] + fn fast_success_emits_no_span() { + // A 5ms successful point read is below every threshold and is not a + // failure, so tail-based sampling must skip it entirely. + let thresholds = DiagnosticsThresholds::default(); + let now = Instant::now(); + let ctx = context( + Duration::from_millis(5), + Some(CosmosStatus::new(StatusCode::Ok)), + Some("read_item"), + &[(5, 5, CosmosStatus::new(StatusCode::Ok))], + now, + ); + assert!(!should_emit_span(&ctx, &thresholds, None)); + } + + #[test] + fn failure_and_slow_are_sampled() { + let thresholds = DiagnosticsThresholds::default(); + let now = Instant::now(); + + // Failure → emit. + let failed = context( + Duration::from_millis(5), + Some(CosmosStatus::new(StatusCode::TooManyRequests)), + Some("read_item"), + &[(5, 5, CosmosStatus::new(StatusCode::TooManyRequests))], + now, + ); + assert!(should_emit_span(&failed, &thresholds, None)); + + // Slow point op (> 1s) → emit. + let slow = context( + Duration::from_millis(1500), + Some(CosmosStatus::new(StatusCode::Ok)), + Some("read_item"), + &[(1500, 1500, CosmosStatus::new(StatusCode::Ok))], + now, + ); + assert!(should_emit_span(&slow, &thresholds, None)); + } + + #[test] + fn non_point_threshold_uses_operation_name_from_context() { + // Production driver contexts carry no operation name, so a 2s operation + // would be classified as a point op (1s threshold) and wrongly sampled. + // Passing the SDK operation identity switches it to the non-point (3s) + // threshold. + let thresholds = DiagnosticsThresholds::default(); + let now = Instant::now(); + let slow_non_point = context( + Duration::from_millis(2000), + Some(CosmosStatus::new(StatusCode::Ok)), + None, // driver context has no operation name (the production case) + &[(2000, 2000, CosmosStatus::new(StatusCode::Ok))], + now, + ); + + // Without an operation identity: falls back to the point (1s) threshold + // and emits. + assert!(should_emit_span(&slow_non_point, &thresholds, None)); + + // With the SDK operation identity: a query is a non-point op (3s + // threshold), so a 2s operation is NOT sampled. + let op = CosmosOperationContext::new().with_operation_name("query_items"); + assert!(!should_emit_span(&slow_non_point, &thresholds, Some(&op))); + } + + #[test] + fn emits_backdated_parent_child_tree() { + let (provider, exporter) = exportable(); + let tracer = provider.tracer("test"); + + let now_instant = Instant::now(); + let now_system = SystemTime::now(); + // A ~250ms failed operation with two attempts. + let ctx = context( + Duration::from_millis(250), + Some(CosmosStatus::new(StatusCode::TooManyRequests)), + Some("read_item"), + &[ + (250, 100, CosmosStatus::new(StatusCode::TooManyRequests)), + (120, 110, CosmosStatus::new(StatusCode::TooManyRequests)), + ], + 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(); + assert_eq!(spans.len(), 3, "root + two attempt children"); + + // The root span carries the operation name and has no parent. + let root = spans + .iter() + .find(|s| s.name == "read_item") + .expect("root span present"); + assert_eq!(root.parent_span_id, SpanId::INVALID); + + // Root is backdated into the past and lasts a non-zero interval. + assert!(root.end_time <= now_system); + assert!(root.start_time < root.end_time); + let root_span_id = root.span_context.span_id(); + + let children: Vec<_> = spans + .iter() + .filter(|s| s.name == "cosmosdb.request") + .collect(); + assert_eq!(children.len(), 2); + for child in &children { + // Correct parentage and backdated (past) timestamps. + assert_eq!(child.parent_span_id, root_span_id); + assert!(child.end_time <= now_system); + assert!(child.start_time <= child.end_time); + } + + // Semantic-convention attributes are present on the root span. + assert!(root.attributes.iter().any(|kv| { + kv.key.as_str() == attributes::DB_SYSTEM_NAME + && kv.value.as_str() == attributes::DB_SYSTEM_NAME_VALUE + })); + assert!(root + .attributes + .iter() + .any(|kv| kv.key.as_str() == attributes::DB_OPERATION_NAME)); + } + + #[test] + fn op_context_supplies_identity_when_driver_context_lacks_it() { + let (provider, exporter) = exportable(); + let tracer = provider.tracer("test"); + + 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. + let ctx = context( + Duration::from_millis(1500), + Some(CosmosStatus::new(StatusCode::Ok)), + None, + &[(1500, 1500, CosmosStatus::new(StatusCode::Ok))], + now_instant, + ); + let op = CosmosOperationContext::new() + .with_operation_name("read_item") + .with_database_name("my_db") + .with_container_name("my_container"); + + emit_backdated_span_tree(&tracer, &ctx, Some(&op), None, now_instant, now_system); + provider.force_flush().unwrap(); + + let spans = exporter.get_finished_spans().unwrap(); + // The op context names the root span and supplies db.operation.name, + // db.namespace, and db.collection.name. + let root = spans + .iter() + .find(|s| s.name == "read_item") + .expect("root span named from op context"); + assert!(root.attributes.iter().any(|kv| { + kv.key.as_str() == attributes::DB_OPERATION_NAME && kv.value.as_str() == "read_item" + })); + assert!(root.attributes.iter().any(|kv| { + kv.key.as_str() == attributes::DB_NAMESPACE && kv.value.as_str() == "my_db" + })); + assert!(root.attributes.iter().any(|kv| { + kv.key.as_str() == attributes::DB_COLLECTION_NAME && kv.value.as_str() == "my_container" + })); + } + + #[test] + fn incomplete_context_is_not_sampled_even_when_slow() { + // A finalized context with neither a status nor any attempts does not + // represent a completed operation. The tail-sampling gate must not emit + // for it even when its elapsed duration alone crosses a threshold. + let thresholds = DiagnosticsThresholds::default(); + let ctx = DiagnosticsContext::for_testing_with_requests( + ActivityId::new_uuid(), + Duration::from_millis(5000), + None, + None, + Vec::new(), + ); + assert!(!ctx.is_completed()); + assert!(!should_emit_span(&ctx, &thresholds, None)); + } + + #[test] + fn op_context_server_address_overrides_endpoint_host() { + // The tracing root span must honor a caller-supplied server.address + // override (as the metrics handler does), not just the endpoint host. + 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, 1500, CosmosStatus::new(StatusCode::Ok))], + now_instant, + ); + let op = CosmosOperationContext::new() + .with_operation_name("read_item") + .with_server_address("override.example.com"); + + 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 == "read_item") + .expect("root span present"); + assert!( + root.attributes.iter().any(|kv| { + kv.key.as_str() == attributes::SERVER_ADDRESS + && kv.value.as_str() == "override.example.com" + }), + "root server.address must honor the op-context override" + ); + } + + #[test] + fn root_span_contains_children_when_duration_underestimates_window() { + // An aggregate operation's duration() is the SUM of its sub-op durations + // and omits the gaps between them, so `op_end - duration()` can fall + // AFTER the earliest attempt. The reconstructed root must still start no + // later than its earliest child so the span tree stays well-formed. + let (provider, exporter) = exportable(); + let tracer = provider.tracer("test"); + let now_instant = Instant::now(); + let now_system = SystemTime::now(); + // duration (50ms) is far smaller than the 250ms window the attempts span. + let ctx = context( + Duration::from_millis(50), + Some(CosmosStatus::new(StatusCode::TooManyRequests)), + Some("read_item"), + &[ + (250, 40, CosmosStatus::new(StatusCode::TooManyRequests)), + (60, 40, 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 root = spans + .iter() + .find(|s| s.name == "read_item") + .expect("root span present"); + let earliest_child = spans + .iter() + .filter(|s| s.name == "cosmosdb.request") + .map(|s| s.start_time) + .min() + .expect("attempt children present"); + assert!( + root.start_time <= earliest_child, + "root must start no later than its earliest child" + ); + } +} 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 new file mode 100644 index 00000000000..f8152ec7a30 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos/src/diagnostics/tracing/span_builder.rs @@ -0,0 +1,290 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Reconstructs a backdated span tree from a completed [`DiagnosticsContext`]. +//! +//! Because the SDK only decides *whether* to emit a span after the operation +//! has finished (tail-based sampling), the spans it produces must be **backdated** +//! to the times the work actually happened. We use the raw `opentelemetry` +//! [`SpanBuilder`](opentelemetry::trace::SpanBuilder) API — which lets us set an +//! explicit start time and end timestamp — rather than the `azure_core` tracing +//! traits, whose backdating support is not yet available on a shipped release. +//! +//! The tree is: one operation ("root") span, plus one child span per retained +//! attempt in [`DiagnosticsContext::requests`]. Monotonic [`Instant`]s recorded +//! during execution are mapped to wall-clock [`SystemTime`]s using an anchor pair +//! `(now_instant, now_system)` captured at emission time. + +use std::time::{Duration, Instant, SystemTime}; + +use opentelemetry::{ + trace::{Span, SpanKind, Status, TraceContextExt, Tracer}, + Array, Context, KeyValue, StringValue, Value, +}; + +use azure_data_cosmos_driver::diagnostics::{DiagnosticsContext, RequestDiagnostics}; + +use crate::diagnostics::attributes; +use crate::diagnostics::CosmosOperationContext; + +/// Span name for the operation ("root") span when the operation name is unknown. +const DEFAULT_OPERATION_SPAN_NAME: &str = "cosmosdb.operation"; + +/// Span name for each per-attempt ("child") span. +const REQUEST_SPAN_NAME: &str = "cosmosdb.request"; + +/// Emits a backdated operation span with one child span per retained attempt. +/// +/// * `tracer` — the OpenTelemetry tracer to build spans on. Generic so tests can +/// pass a `SdkTracer` backed by an in-memory exporter with no global state. +/// * `diagnostics` — the completed context to reconstruct. +/// * `op` — the SDK-supplied operation identity (name, database, container), if +/// present on the pipeline context. Supplies the operation name when the driver +/// context does not carry one, and the `db.namespace` / `db.collection.name` +/// attributes the driver context never knows. +/// * `now_instant` / `now_system` — an anchor pair captured together at emission +/// time; monotonic instants in the context are converted to wall-clock times +/// relative to this anchor. +pub(crate) fn emit_backdated_span_tree( + tracer: &T, + diagnostics: &DiagnosticsContext, + op: Option<&CosmosOperationContext>, + reason: Option<&str>, + now_instant: Instant, + now_system: SystemTime, +) where + T: Tracer, +{ + // Map a monotonic `Instant` recorded during execution to wall-clock time. + let to_system = |instant: Instant| -> SystemTime { + now_system - now_instant.saturating_duration_since(instant) + }; + + let requests = diagnostics.requests(); + + // Wall-clock end of one attempt: its recorded completion, or its start plus + // reported duration when no completion instant was captured. + let child_end_of = |req: &RequestDiagnostics| -> SystemTime { + let start = to_system(req.started_at()); + let end = match req.completed_at() { + Some(completed) => to_system(completed), + None => start + Duration::from_millis(req.duration_ms()), + }; + end.max(start) + }; + + // Anchor the root's end to the last attempt's completion, NOT to + // `now_system` (the handler's invocation time). Handlers run in registration + // order, so using the invocation time would fold any earlier handler's delay + // into this span while the child timestamps stay anchored to the real request + // instants — inflating and shifting the reconstructed operation span. Fall + // back to `now_system` only when there are no retained attempts. + let op_end = requests + .iter() + .map(&child_end_of) + .max() + .unwrap_or(now_system); + // The root span must fully contain its backdated children. A normal + // operation's `duration()` already reaches back to before its first attempt, + // but an aggregate operation's `duration()` is the SUM of its sub-op + // durations (see `DiagnosticsContext::aggregate_sub_operations`) and omits + // the wall-clock gaps between them, so `op_end - duration()` can land AFTER + // the earliest child. Extend the start back to the earliest retained attempt + // so the reconstructed root spans first-attempt start to last-attempt end. + let duration_start = op_end - diagnostics.duration(); + let op_start = requests + .iter() + .map(|req| to_system(req.started_at())) + .min() + .map(|earliest| earliest.min(duration_start)) + .unwrap_or(duration_start); + 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)); + let op_name = op_name_ref + .unwrap_or(DEFAULT_OPERATION_SPAN_NAME) + .to_string(); + + let mut root_attrs = vec![KeyValue::new( + attributes::DB_SYSTEM_NAME, + attributes::DB_SYSTEM_NAME_VALUE, + )]; + if let Some(name) = op_name_ref { + root_attrs.push(KeyValue::new( + attributes::DB_OPERATION_NAME, + name.to_string(), + )); + } + if let Some(namespace) = op.and_then(CosmosOperationContext::database_name) { + root_attrs.push(KeyValue::new( + attributes::DB_NAMESPACE, + namespace.to_string(), + )); + } + if let Some(collection) = op.and_then(CosmosOperationContext::container_name) { + root_attrs.push(KeyValue::new( + attributes::DB_COLLECTION_NAME, + collection.to_string(), + )); + } + if let Some(status) = diagnostics.effective_status() { + root_attrs.push(KeyValue::new( + attributes::DB_RESPONSE_STATUS_CODE, + u16::from(status.status_code()).to_string(), + )); + if let Some(sub) = status.sub_status() { + root_attrs.push(KeyValue::new( + attributes::SUB_STATUS_CODE, + i64::from(u16::from(sub)), + )); + } + } + root_attrs.push(KeyValue::new( + attributes::REQUEST_CHARGE, + diagnostics.total_request_charge().value(), + )); + let regions = diagnostics.regions_contacted(); + if !regions.is_empty() { + // Semantic conventions define contacted_regions as an ordered string[]; + // emit an OpenTelemetry array, not a joined scalar. + let values: Vec = regions + .iter() + .map(|r| StringValue::from(r.as_str().to_string())) + .collect(); + root_attrs.push(KeyValue::new( + attributes::CONTACTED_REGIONS, + Value::Array(Array::String(values)), + )); + } + // Prefer the caller-supplied server-address override (mirroring the metrics + // handler) before falling back to the host of the first contacted endpoint, + // so an override changes both the metric and the root span consistently. + let server_addr = op + .and_then(CosmosOperationContext::server_address) + .map(str::to_string) + .or_else(|| requests.first().and_then(server_address)); + if let Some(addr) = server_addr { + root_attrs.push(KeyValue::new(attributes::SERVER_ADDRESS, addr)); + } + if let Some(machine_id) = diagnostics.machine_id() { + root_attrs.push(KeyValue::new( + attributes::MACHINE_ID, + machine_id.to_string(), + )); + } + // Record why the operation was sampled (failure, or which threshold), so the + // root span carries the same reason the sampled log line does. + if let Some(reason) = reason { + root_attrs.push(KeyValue::new( + attributes::SAMPLING_REASON, + reason.to_string(), + )); + } + if op_failed { + if let Some(status) = diagnostics.effective_status() { + root_attrs.push(KeyValue::new( + attributes::ERROR_TYPE, + u16::from(status.status_code()).to_string(), + )); + } else { + // A failure with no status anywhere — report the semconv catch-all so + // the Error-marked root span still carries an error.type. + root_attrs.push(KeyValue::new( + attributes::ERROR_TYPE, + attributes::ERROR_TYPE_OTHER, + )); + } + } + + let root_builder = tracer + .span_builder(op_name) + .with_kind(SpanKind::Client) + .with_start_time(op_start) + .with_attributes(root_attrs); + let mut root = tracer.build_with_context(root_builder, &Context::current()); + if op_failed { + root.set_status(Status::error("operation failed")); + } + + // Children hang off the root via its span context so the exporter records the + // correct parent/child linkage. + let parent_cx = Context::current().with_remote_span_context(root.span_context().clone()); + + // --- Attempt (child) spans --- + for req in requests.iter() { + let child_start = to_system(req.started_at()); + let child_end = child_end_of(req); + + let req_status = req.status(); + let req_failed = !req_status.is_success(); + + let mut child_attrs = vec![ + KeyValue::new(attributes::DB_SYSTEM_NAME, attributes::DB_SYSTEM_NAME_VALUE), + KeyValue::new( + attributes::DB_RESPONSE_STATUS_CODE, + u16::from(req_status.status_code()).to_string(), + ), + KeyValue::new(attributes::REQUEST_CHARGE, req.request_charge().value()), + ]; + if let Some(name) = op_name_ref { + child_attrs.push(KeyValue::new( + attributes::DB_OPERATION_NAME, + name.to_string(), + )); + } + if let Some(sub) = req_status.sub_status() { + child_attrs.push(KeyValue::new( + attributes::SUB_STATUS_CODE, + i64::from(u16::from(sub)), + )); + } + if let Some(region) = req.region() { + child_attrs.push(KeyValue::new( + attributes::CONTACTED_REGIONS, + Value::Array(Array::String(vec![StringValue::from( + region.as_str().to_string(), + )])), + )); + } + if let Some(addr) = server_address(req) { + child_attrs.push(KeyValue::new(attributes::SERVER_ADDRESS, addr)); + } + if let Some(activity_id) = req.activity_id() { + child_attrs.push(KeyValue::new( + attributes::ACTIVITY_ID, + activity_id.as_str().to_string(), + )); + } + if req_failed { + child_attrs.push(KeyValue::new( + attributes::ERROR_TYPE, + u16::from(req_status.status_code()).to_string(), + )); + } + + let child_builder = tracer + .span_builder(REQUEST_SPAN_NAME) + .with_kind(SpanKind::Client) + .with_start_time(child_start) + .with_attributes(child_attrs); + let mut child = tracer.build_with_context(child_builder, &parent_cx); + if req_failed { + child.set_status(Status::error("request failed")); + } + child.end_with_timestamp(child_end); + } + + root.end_with_timestamp(op_end); +} + +/// Extracts the host portion of a request's endpoint URL for `server.address`. +fn server_address(req: &RequestDiagnostics) -> Option { + url::Url::parse(req.endpoint()) + .ok() + .and_then(|url| url.host_str().map(str::to_string)) +} diff --git a/sdk/cosmos/azure_data_cosmos/src/feed/change_feed_iterator.rs b/sdk/cosmos/azure_data_cosmos/src/feed/change_feed_iterator.rs index cacdfbc043c..2407c97ca49 100644 --- a/sdk/cosmos/azure_data_cosmos/src/feed/change_feed_iterator.rs +++ b/sdk/cosmos/azure_data_cosmos/src/feed/change_feed_iterator.rs @@ -14,6 +14,7 @@ use futures::future::BoxFuture; use futures::Stream; use serde::de::DeserializeOwned; +use crate::diagnostics::{CosmosOperationContext, DiagnosticsHandlerChain}; use crate::{driver_bridge, feed::page::FeedBody, feed::FeedPage, models::CosmosResponse}; type DriverPageFuture = BoxFuture<'static, (OperationPlan, crate::Result>)>; @@ -27,6 +28,11 @@ struct LiveState { plan: Option, in_flight: Option, errored: bool, + /// Emission handler chain plus the paged operation's identity, dispatched + /// once per page fetch (both success and failure) so change-feed pagination + /// honors the same once-per-operation contract as singleton operations. + diagnostics: DiagnosticsHandlerChain, + op_context: CosmosOperationContext, } impl LiveState { @@ -35,6 +41,8 @@ impl LiveState { container: Option, plan: OperationPlan, options: OperationOptions, + diagnostics: DiagnosticsHandlerChain, + op_context: CosmosOperationContext, ) -> Self { Self { driver, @@ -43,6 +51,8 @@ impl LiveState { plan: Some(plan), in_flight: None, errored: false, + diagnostics, + op_context, } } @@ -107,6 +117,14 @@ impl LiveState { } Err(err) => { *this.errored = true; + if let Some(page_diagnostics) = err.diagnostics() { + crate::feed::dispatch_page_diagnostics( + this.diagnostics, + this.op_context, + &page_diagnostics, + None, + ); + } task::Poll::Ready(Some(Err(err))) } Ok(Some(driver_response)) => { @@ -119,6 +137,12 @@ impl LiveState { // Return an empty page — do not try to deserialize the // (potentially empty) body. if status.status_code() == azure_core::http::StatusCode::NotModified { + crate::feed::dispatch_page_diagnostics( + this.diagnostics, + this.op_context, + &diagnostics, + Some(0), + ); let page = FeedPage::new(Vec::new(), headers, diagnostics); return task::Poll::Ready(Some(Ok(page))); } @@ -128,11 +152,23 @@ impl LiveState { // `ChangeFeedItem`) without stripping any fields. match deserialize_change_feed_items::(response) { Ok(items) => { + crate::feed::dispatch_page_diagnostics( + this.diagnostics, + this.op_context, + &diagnostics, + Some(items.len() as u64), + ); let page = FeedPage::new(items, headers, diagnostics); task::Poll::Ready(Some(Ok(page))) } Err(err) => { *this.errored = true; + crate::feed::dispatch_page_diagnostics( + this.diagnostics, + this.op_context, + &diagnostics, + None, + ); task::Poll::Ready(Some(Err(err))) } } @@ -227,9 +263,18 @@ impl ChangeFeedPageIterator { container: Option, plan: OperationPlan, options: OperationOptions, + diagnostics: DiagnosticsHandlerChain, + op_context: CosmosOperationContext, ) -> Self { Self { - state: Box::pin(LiveState::new(driver, container, plan, options)), + state: Box::pin(LiveState::new( + driver, + container, + plan, + options, + diagnostics, + op_context, + )), _marker: PhantomData, } } diff --git a/sdk/cosmos/azure_data_cosmos/src/feed/iterator.rs b/sdk/cosmos/azure_data_cosmos/src/feed/iterator.rs index d9cf096e94e..72cb40aa725 100644 --- a/sdk/cosmos/azure_data_cosmos/src/feed/iterator.rs +++ b/sdk/cosmos/azure_data_cosmos/src/feed/iterator.rs @@ -17,6 +17,7 @@ use futures::future::BoxFuture; use futures::Stream; use serde::de::DeserializeOwned; +use crate::diagnostics::{CosmosOperationContext, DiagnosticsHandlerChain}; use crate::{driver_bridge, feed::query_page::QueryFeedPage}; type DriverPageFuture = BoxFuture<'static, (OperationPlan, crate::Result>)>; @@ -32,6 +33,11 @@ struct LiveState { /// `Some` while a page fetch is pending. in_flight: Option, exhausted: bool, + /// Emission handler chain plus the paged operation's identity, dispatched + /// once per page fetch (both success and failure) so query pagination + /// honors the same once-per-operation contract as singleton operations. + diagnostics: DiagnosticsHandlerChain, + op_context: CosmosOperationContext, } impl LiveState { @@ -40,6 +46,8 @@ impl LiveState { container: Option, plan: OperationPlan, options: OperationOptions, + diagnostics: DiagnosticsHandlerChain, + op_context: CosmosOperationContext, ) -> Self { Self { driver, @@ -48,6 +56,8 @@ impl LiveState { plan: Some(plan), in_flight: None, exhausted: false, + diagnostics, + op_context, } } @@ -116,15 +126,45 @@ impl LiveState { Err(err) => { // An error from the driver indicates a failure to fetch the next page. Mark ourselves as exhausted and return the error. *this.exhausted = true; + if let Some(page_diagnostics) = err.diagnostics() { + crate::feed::dispatch_page_diagnostics( + this.diagnostics, + this.op_context, + &page_diagnostics, + None, + ); + } task::Poll::Ready(Some(Err(err))) } Ok(Some(driver_response)) => { // Successfully got a response from the driver. Convert it into a QueryFeedPage and yield it. let response = driver_bridge::driver_response_to_cosmos_response(driver_response); + // Only take a diagnostics handle when a handler is registered: the + // clone is a per-page atomic op the empty-chain default skips. + let page_diagnostics = + (!this.diagnostics.is_empty()).then(|| response.diagnostics()); match QueryFeedPage::from_response(response) { - Ok(page) => task::Poll::Ready(Some(Ok(page))), + Ok(page) => { + if let Some(page_diagnostics) = &page_diagnostics { + crate::feed::dispatch_page_diagnostics( + this.diagnostics, + this.op_context, + page_diagnostics, + Some(page.items().len() as u64), + ); + } + task::Poll::Ready(Some(Ok(page))) + } Err(err) => { *this.exhausted = true; + if let Some(page_diagnostics) = &page_diagnostics { + crate::feed::dispatch_page_diagnostics( + this.diagnostics, + this.op_context, + page_diagnostics, + None, + ); + } task::Poll::Ready(Some(Err(err))) } } @@ -195,9 +235,18 @@ impl QueryItemIterator { container: Option, plan: OperationPlan, options: OperationOptions, + diagnostics: DiagnosticsHandlerChain, + op_context: CosmosOperationContext, ) -> Self { Self { - source: PageSource::Live(Box::pin(LiveState::new(driver, container, plan, options))), + source: PageSource::Live(Box::pin(LiveState::new( + driver, + container, + plan, + options, + diagnostics, + op_context, + ))), current: None, _marker: PhantomData, } diff --git a/sdk/cosmos/azure_data_cosmos/src/feed/mod.rs b/sdk/cosmos/azure_data_cosmos/src/feed/mod.rs index e2123c7d9d3..961b8931301 100644 --- a/sdk/cosmos/azure_data_cosmos/src/feed/mod.rs +++ b/sdk/cosmos/azure_data_cosmos/src/feed/mod.rs @@ -4,6 +4,10 @@ //! Types related to Cosmos DB feed operations, including query and change feed //! iteration, pagination and related models. +use azure_core::http::Context; + +use crate::diagnostics::{CosmosOperationContext, DiagnosticsContext, DiagnosticsHandlerChain}; + // ========================================================================= // Public API // ========================================================================= @@ -31,3 +35,35 @@ mod iterator; mod page; mod query; mod query_page; + +/// Dispatches a completed feed page's diagnostics to the registered handler +/// chain, carrying the paged operation's identity ([`query_items`] / +/// [`query_change_feed`] etc. plus database/container). +/// +/// This is the per-page completion seam for query and change-feed pagination, +/// mirroring the singleton completion seam on `ClientContext`: it lets metrics, +/// tracing, and sampled-logging handlers observe each page fetch, on both +/// success and failure. It is skipped when no handler is registered, so the +/// default paginated path does no diagnostics work per page. +/// +/// `returned_item_count`, when supplied, feeds the `returned_rows` development +/// metric with the number of items the page yielded. +/// +/// [`query_items`]: crate::clients::ContainerClient::query_items +/// [`query_change_feed`]: crate::clients::ContainerClient::query_change_feed +pub(crate) fn dispatch_page_diagnostics( + diagnostics: &DiagnosticsHandlerChain, + op_context: &CosmosOperationContext, + page_diagnostics: &DiagnosticsContext, + returned_item_count: Option, +) { + if diagnostics.is_empty() { + return; + } + let mut op = op_context.clone(); + if let Some(count) = returned_item_count { + op = op.with_returned_item_count(count); + } + let cx = Context::new().with_value(op); + diagnostics.dispatch(page_diagnostics, &cx); +} diff --git a/sdk/cosmos/azure_data_cosmos/src/options/client.rs b/sdk/cosmos/azure_data_cosmos/src/options/client.rs index 9b4bab4d324..6df181f3e2f 100644 --- a/sdk/cosmos/azure_data_cosmos/src/options/client.rs +++ b/sdk/cosmos/azure_data_cosmos/src/options/client.rs @@ -3,8 +3,12 @@ //! [`CosmosClientOptions`] — options for [`CosmosClient`](crate::CosmosClient) construction. +use std::sync::Arc; + use azure_data_cosmos_driver::options::{OperationOptions, UserAgentSuffix}; +use crate::diagnostics::{DiagnosticsHandler, DiagnosticsHandlerChain}; + /// Options used when creating a [`CosmosClient`](crate::CosmosClient). /// /// This struct is used internally by [`CosmosClientBuilder`](crate::CosmosClientBuilder). @@ -17,6 +21,8 @@ pub struct CosmosClientOptions { /// unless overridden by per-request options. pub operation: OperationOptions, pub(crate) user_agent_suffix: Option, + /// Diagnostics emission handlers invoked once per operation at completion. + pub(crate) diagnostics_handlers: DiagnosticsHandlerChain, } impl CosmosClientOptions { @@ -29,4 +35,14 @@ impl CosmosClientOptions { self.operation = operation; self } + + /// Registers a [`DiagnosticsHandler`](crate::diagnostics::DiagnosticsHandler) + /// invoked once per operation at completion. + /// + /// Handlers run in registration order. Call this multiple times to build an + /// ordered chain. + pub fn with_diagnostics_handler(mut self, handler: Arc) -> Self { + self.diagnostics_handlers = self.diagnostics_handlers.with_handler(handler); + self + } } 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 new file mode 100644 index 00000000000..319b8149c2f --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/handler_propagation.rs @@ -0,0 +1,310 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! End-to-end coverage that a registered [`DiagnosticsHandler`] actually receives +//! a completed context through a **real** SDK operation — driven by the in-memory +//! emulator. This guards the completion seams (singleton success + failure, and +//! paginated success) against wiring regressions that still compile. + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; + +use azure_core::http::Context; +use azure_data_cosmos::diagnostics::{ + CosmosOperationContext, DiagnosticsContext, DiagnosticsHandler, +}; +use azure_data_cosmos::options::Region; +use azure_data_cosmos::{ + AccountEndpoint, AccountReference, CosmosClient, CosmosClientBuilder, CosmosRuntimeBuilder, + FeedScope, Query, RoutingStrategy, +}; +use azure_data_cosmos_driver::in_memory_emulator::{ + ConsistencyLevel, InMemoryEmulatorHttpClient, VirtualAccountConfig, VirtualRegion, +}; +use futures::TryStreamExt; +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +const EMULATOR_GATEWAY_URL: &str = "https://eastus.emulator.local"; + +#[derive(Debug, Serialize, Deserialize, PartialEq)] +struct TestDoc { + id: String, + pk: String, + value: i64, +} + +/// The operation-scope identity (`CosmosOperationContext`) a handler observed on +/// its most recent invocation. +#[derive(Clone, Debug, Default, PartialEq)] +struct ObservedOp { + operation_name: Option, + database_name: Option, + container_name: Option, +} + +/// A [`DiagnosticsHandler`] that counts invocations, how many carried a failed +/// context, and the operation-scope identity of the latest invocation — so a test +/// can assert the chain fired on both success and failure *and* that the correct +/// operation/database/container identity (WS8) was propagated. +#[derive(Default)] +struct CountingHandler { + total: AtomicUsize, + failures: AtomicUsize, + last_op: Mutex>, +} + +impl CountingHandler { + fn total(&self) -> usize { + self.total.load(Ordering::SeqCst) + } + + fn failures(&self) -> usize { + self.failures.load(Ordering::SeqCst) + } + + /// The operation identity observed on the most recent invocation, or `None` + /// when the handler was invoked without a `CosmosOperationContext`. + fn last_op(&self) -> Option { + self.last_op.lock().unwrap().clone() + } +} + +impl DiagnosticsHandler for CountingHandler { + fn handle(&self, diagnostics: &DiagnosticsContext, cx: &Context<'_>) { + self.total.fetch_add(1, Ordering::SeqCst); + if diagnostics.is_failure() { + self.failures.fetch_add(1, Ordering::SeqCst); + } + // Capture the SDK-supplied operation identity carried on the pipeline + // context so a test can assert the WS8 `db.*` wiring (operation name, + // database, container) actually reaches handlers. Store the observed + // value verbatim (including `None`) so missing wiring is detectable. + let observed = cx.value::().map(|op| ObservedOp { + operation_name: op.operation_name().map(str::to_owned), + database_name: op.database_name().map(str::to_owned), + container_name: op.container_name().map(str::to_owned), + }); + *self.last_op.lock().unwrap() = observed; + } +} + +/// Builds an emulator-backed SDK client with `handler` registered and provisions +/// an empty `(database, container)` keyed on `/pk`. +async fn setup() -> (CosmosClient, Arc, String, String) { + let config = VirtualAccountConfig::new(vec![VirtualRegion::new( + "East US", + azure_core::http::Url::parse(EMULATOR_GATEWAY_URL).unwrap(), + )]) + .unwrap() + .with_consistency(ConsistencyLevel::Session); + + let emulator = Arc::new(InMemoryEmulatorHttpClient::new(config)); + let store = emulator.store(); + + let db = format!("diag-{}", &Uuid::new_v4().to_string()[..8]); + let container = "items".to_string(); + store.create_database(&db); + store.create_container( + &db, + &container, + serde_json::from_value(serde_json::json!({ + "paths": ["/pk"], + "kind": "Hash", + "version": 2 + })) + .unwrap(), + ); + + let account = AccountReference::with_authentication_key( + EMULATOR_GATEWAY_URL.parse::().unwrap(), + azure_core::credentials::Secret::new("dGVzdGtleQ=="), + ); + + let handler = Arc::new(CountingHandler::default()); + let client = CosmosClientBuilder::new() + .with_runtime( + CosmosRuntimeBuilder::from(emulator.runtime_builder()) + .build() + .await + .unwrap(), + ) + .with_diagnostics_handler(handler.clone()) + .build(account, RoutingStrategy::ProximityTo(Region::EAST_US)) + .await + .unwrap(); + + (client, handler, db, container) +} + +#[tokio::test] +async fn handler_receives_singleton_success() { + let (client, handler, db, container) = setup().await; + let c = client + .database_client(&db) + .container_client(&container) + .await + .unwrap(); + + let before = handler.total(); + c.create_item( + "pkA", + "doc-1", + &TestDoc { + id: "doc-1".into(), + pk: "pkA".into(), + value: 1, + }, + None, + ) + .await + .unwrap(); + + assert_eq!( + handler.total(), + before + 1, + "a singleton success must dispatch exactly one completion callback" + ); + assert_eq!( + handler.last_op(), + Some(ObservedOp { + operation_name: Some("create_item".to_string()), + database_name: Some(db.clone()), + container_name: Some(container.clone()), + }), + "the create_item operation must propagate its db.* identity to handlers" + ); +} + +#[tokio::test] +async fn handler_receives_singleton_failure() { + let (client, handler, db, container) = setup().await; + let c = client + .database_client(&db) + .container_client(&container) + .await + .unwrap(); + + let before_total = handler.total(); + let before_failures = handler.failures(); + + let result = c.read_item("pkMissing", "does-not-exist", None).await; + assert!(result.is_err(), "reading a missing item must fail"); + + assert_eq!( + handler.total(), + before_total + 1, + "a singleton failure must dispatch exactly one completion callback" + ); + assert_eq!( + handler.failures(), + before_failures + 1, + "the failed operation must be dispatched exactly once with a failed context" + ); + assert_eq!( + handler.last_op().and_then(|op| op.operation_name), + Some("read_item".to_string()), + "the failed read_item must still propagate its operation identity" + ); +} + +#[tokio::test] +async fn handler_receives_paginated_success() { + let (client, handler, db, container) = setup().await; + let c = client + .database_client(&db) + .container_client(&container) + .await + .unwrap(); + + for i in 0..3 { + let id = format!("q-{i}"); + c.create_item( + "pkQ", + &id, + &TestDoc { + id: id.clone(), + pk: "pkQ".into(), + value: i, + }, + None, + ) + .await + .unwrap(); + } + + 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(); + + assert_eq!(items.len(), 3); + assert_eq!( + handler.total(), + before + 1, + "a single query page must dispatch exactly one completion callback" + ); + assert_eq!( + handler.last_op(), + Some(ObservedOp { + operation_name: Some("query_items".to_string()), + database_name: Some(db.clone()), + container_name: Some(container.clone()), + }), + "the query_items operation must propagate its db.* identity to handlers" + ); +} + +/// The paginated **failure** dispatch seam — a page fetch that errors — must fire +/// the handler exactly once with a failed context. This is a distinct completion +/// seam from the paginated-success and singleton paths and can regress +/// independently, so it gets its own emulator-driven coverage. +#[tokio::test] +async fn handler_receives_paginated_failure() { + let (client, handler, db, container) = setup().await; + let c = client + .database_client(&db) + .container_client(&container) + .await + .unwrap(); + + let before_total = handler.total(); + let before_failures = handler.failures(); + + // 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; + assert!( + result.is_err(), + "the invalid query must surface as a terminal page-fetch error" + ); + + assert_eq!( + handler.total(), + before_total + 1, + "a failed query page must dispatch exactly one completion callback" + ); + assert_eq!( + handler.failures(), + before_failures + 1, + "the failed page fetch must be dispatched exactly once with a failed context" + ); +} diff --git a/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/mod.rs b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/mod.rs index e2b905900f5..003e65b0ed2 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/mod.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/mod.rs @@ -15,6 +15,7 @@ pub mod dtx_live_comparison; pub mod dtx_sdk_validation; pub mod dual_backend; pub mod end_to_end; +pub mod handler_propagation; pub mod hpk; pub mod partition_key_equality; pub mod query_comparison; diff --git a/sdk/cosmos/azure_data_cosmos_driver/ARCHITECTURE.md b/sdk/cosmos/azure_data_cosmos_driver/ARCHITECTURE.md index 1d18e9c982b..5566d2216e5 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/ARCHITECTURE.md +++ b/sdk/cosmos/azure_data_cosmos_driver/ARCHITECTURE.md @@ -369,6 +369,43 @@ let options = DiagnosticsOptions::builder() If output exceeds the limit, it's truncated with an indicator. +### Record Bounding (Retry Storms) + +Independently of the summary-string size limit, the finalized per-attempt +`RequestDiagnostics` list itself is bounded so a retry storm cannot grow the +context — or any artifact derived from it — without limit: + +```rust +use azure_data_cosmos_driver::options::DiagnosticsOptions; + +let options = DiagnosticsOptions::builder() + .with_max_request_diagnostics(512) // record cap (default) + .build()?; +``` + +- **Default**: 512 records +- **Minimum**: 16 records +- **Environment Variable**: `AZURE_COSMOS_DIAGNOSTICS_MAX_REQUEST_DIAGNOSTICS` + +When an operation makes more attempts than the cap, the finalized list is +compacted at completion (never mid-operation, so in-flight request handles stay +valid): + +1. **Run-collapse** — runs of near-identical consecutive retries (same region, + endpoint, status incl. sub-status, and execution context) are collapsed to + their first and last record. +2. **Exact per-run rollup** — each run keeps an exact count, total RU, and + min/max/P50 duration, so no attempt is lost from the aggregate. +3. **Global-bucket fallback** — a region ping-pong (`A→B→A→B`) or a high-cardinality + `410` fan-out across many physical-partition endpoints is bounded by keeping + only the largest key buckets. +4. **Explicit truncation** — every drop is marked on `DiagnosticsContext::compaction()` + (`retained_truncated`, `omitted_runs`, `omitted_request_count`), never silent. + +`DiagnosticsContext::request_count()` and `total_request_charge()` remain exact +(the true totals across all attempts); `retained_request_count()` reports how many +records survived in `requests()`. + ### Compaction & Deduplication (Summary Mode) Summary mode applies intelligent compaction: diff --git a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md index fa1d74748a4..85527aaa9c6 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md @@ -14,6 +14,8 @@ ### Features Added +- Added `DiagnosticsOptions::max_request_diagnostics` (default 512, minimum 16), which bounds the finalized per-attempt `RequestDiagnostics` list under a `429`/`410` retry storm by compacting near-identical retry runs at completion; `DiagnosticsContext::request_count()`/`total_request_charge()` stay exact, and `retained_request_count()`/`compaction()` surface every truncation. ([#4789](https://github.com/Azure/azure-sdk-for-rust/pull/4789)) +- Added `DiagnosticsThresholds` (Java-like defaults) and the completed-operation predicates `DiagnosticsContext::is_failure`/`is_completed`/`is_threshold_violated` (plus an `operation_name` accessor) that power the SDK's tail-based sampling handlers. `DiagnosticsContext::threshold_breach_for` and the `ThresholdBreach` enum additionally report *which* bound (point/non-point latency or request charge) an operation crossed. ([#4789](https://github.com/Azure/azure-sdk-for-rust/pull/4789)) - Added runtime diagnostics output configuration via `CosmosDriverRuntimeBuilder::with_diagnostics_options`, including the env-backed `AZURE_COSMOS_DIAGNOSTICS_DEFAULT_VERBOSITY` option. The built-in default is now summary diagnostics JSON. ([#4733](https://github.com/Azure/azure-sdk-for-rust/pull/4733)) - Added Gateway 2.0 transport (a regional proxy forwarding RNTBD-over-HTTP/2). Gateway 2.0 is used automatically when the account advertises thin-client endpoints, the connectivity probe confirms them, and the runtime has not opted out. ([#4319](https://github.com/Azure/azure-sdk-for-rust/pull/4319), [#4803](https://github.com/Azure/azure-sdk-for-rust/pull/4803)) - Added `ConnectionPoolOptions::gateway_v2_disabled` and `ConnectionPoolOptionsBuilder::with_gateway_v2_disabled`, with `AZURE_COSMOS_CONNECTION_POOL_GATEWAY_V2_DISABLED` configuration and an environment-only `_OVERRIDE` incident switch. ([#4763](https://github.com/Azure/azure-sdk-for-rust/pull/4763)) diff --git a/sdk/cosmos/azure_data_cosmos_driver/DIAGNOSTICS-CONTRACT.md b/sdk/cosmos/azure_data_cosmos_driver/DIAGNOSTICS-CONTRACT.md new file mode 100644 index 00000000000..824cec0b3e9 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos_driver/DIAGNOSTICS-CONTRACT.md @@ -0,0 +1,509 @@ + + +# The driver↔SDK diagnostics contract & its OpenTelemetry mapping + +> **Status: contract for the observability layer landing in this PR.** This document defines the +> **contract** between the `azure_data_cosmos_driver` and its consumers — the Rust +> `azure_data_cosmos` SDK today, and the next-major Go/Python SDKs over FFI — plus the +> **OpenTelemetry customer experience**. The emission layer it describes (the handler chain and +> the built-in metrics/tracing/log handlers) is **implemented in this same PR**, on top of the +> already-shipping `DiagnosticsContext` foundation; every addition is **additive, non-breaking, +> and off-by-default**. This doc file itself adds no code. +> +> **Emission model.** The **driver produces and materializes** one canonical, always-collected +> `DiagnosticsContext`; the **SDK decides what telemetry to emit** from it, through a chain of +> pluggable [`DiagnosticsHandler`](#4-the-emission-model--the-diagnosticshandler-chain)s with +> **tail-based (late-bound) sampling** and **cross-operation rate limiting**. Rule of thumb: +> *the driver produces and materializes; the SDK emits.* +> +> **Priority.** The contract, the OpenTelemetry metrics/traces mapping, and the emission model +> come **first**. A separately prototyped append-only *capture engine* is a **deferred +> optimization** ([WS7], salvaged from PR #4619) — it stays OFF by default and is an +> implementation detail *under* this contract, not part of it. + +## 0. What exists today vs. what this PR adds + +To keep the design honest, references below are tagged by where each piece lands, rather than +implying everything already exists: + +- **[main]** — already on `upstream/main` (`sdk/cosmos/azure_data_cosmos_driver`): the + always-collected [`DiagnosticsContext`][ctx] and its `requests()` / `duration()` / `status()` / + `total_request_charge()` / `regions_contacted()` / `machine_id()` / + `to_json_string(verbosity)` surface (plus the crate-private `is_completed()`), + [`RequestDiagnostics`][ctx], the `DiagnosticsOptions` / `DiagnosticsVerbosity` config, and the + non-optional `CosmosResponse::diagnostics() -> Arc`. **This is the + foundation this contract builds on, and it already exists** — so there is no separate + foundation workstream (the earlier "WS0" was dropped and its stand-alone PR #4785 closed as + redundant). +- **[WS2]** — the SDK-side [`DiagnosticsHandler`](#4-the-emission-model--the-diagnosticshandler-chain) + trait + chain (the emission extension point). +- **[WS3] / [WS4] / [WS5]** — the three built-in handlers: OTel metrics, tail-sampled backdated + tracing, and the rate-limited sampling-log handler. +- **[WS6]** — the retry-storm **bounded-size** compaction (salvaged from the closed PR #4683), + built over the [main] context. +- **[WS8]** — SDK **op-scope wiring**: the SDK supplies each operation's identity (operation + name, database, container) to the handler chain via a `CosmosOperationContext` on the pipeline + `Context`, so the emitted metrics/spans carry `db.operation.name` / `db.namespace` / + `db.collection.name` in production. +- **[WS4/WS5]** — the small driver additions the emission layer needs that upstream lacked: + [`DiagnosticsThresholds`][thresholds] and the public `is_failure()` / `is_completed()` / + `is_threshold_violated()` predicates and the `operation_name()` accessor. +- **[WS4c]** — the upstream span-backdating trait additions to `typespec_client_core` + + `azure_core_opentelemetry`, shipped as the **separate** core-crate PR #4784 (not in this PR). +- **[WS7]** — the deferred, OFF-by-default append-only **capture engine** (salvaged from the + closed PR #4619): the `event.rs` `Span`/`Attr` model, the encode path, and the per-op / + per-attempt wall-clock capture optimization. Not in this PR. +- **[design]** — a surface this contract sketches that is **not yet implemented** (e.g. the + SDK-facing `DiagnosticsLevel { Minimal, Standard, Full }`, which maps onto the [main] + `DiagnosticsVerbosity`). + +> **[WS2]–[WS8] all ship together in one combined observability PR (#4789)** — this contract doc +> is folded into it. They are additive and off-by-default (the OTel handlers are behind the +> `metrics` / `distributed_tracing` features) and are built entirely on the existing [main] +> foundation. +> +> ⚠ **Correction to the earlier draft of this doc (PR #4684).** An earlier draft assumed the +> `DiagnosticsContext` foundation was **not yet on `main`** and would land in a blocking-root +> workstream first. That is wrong: the full foundation — `DiagnosticsContext`, +> `CosmosResponse::diagnostics()`, `RequestDiagnostics`, `DiagnosticsVerbosity`, and +> `DiagnosticsOptions` — is **already on `upstream/main`**. The one diagnostics type that is +> **not** upstream is [`DiagnosticsThresholds`][thresholds]; this PR adds it alongside the +> emission handlers. The tags above reflect that reality. + +## 1. Requirements this contract satisfies + +Each requirement below traces to a section here (and to the workstream that implements it). + +| # | Requirement | Where addressed in this doc | +| -- | ----------- | --------------------------- | +| R1 | Long-running (10–12 h) benchmark observability: **quiet at steady state, rich on error**, enough to root-cause a 5–10 min error window | The whole design; realized by §4 (chain) + §5 (tail-sampling + rate-limit) + §10.2 (always-on metrics) | +| R2 | Emit **OTel metrics** (top priority) — stable semconv first | §10.2 (operation-level metrics) | +| R3 | Emit **OTel tracing** with rich, Java-like features | §10.4 (span tree) | +| R4 | **Tail-based / late-bound sampling**: decide span emission *after* an op completes, by latency/outcome | §5.2 (the emit/skip decision) | +| R5 | **Rate-limit under error storms**: cap emissions so 10k errors/sec don't peg CPU | §5.3 (cross-operation rate limiting) + §8 (bounded per-artifact size) | +| R6 | **Driver → context only** (≤ debug level); **SDK decides emission** from the context | §4 (handler chain) + §11 (SDK↔driver split) | +| R7 | Metrics power **client-side Grafana dashboards** (account/region SLA) | §10.2 (metric instruments + dimensions) | +| R8 | Metrics carry **dimensions** (operation, status, consistency, region…) → per-combination series | §10.2 + §10.3 (attribute tiers) | +| R9 | Don't lock/log a line for **every fast op** — steady state must be near-zero cost | §5 defaults (tail-sampling skips fast successes) + §10.2 (low-cardinality always-on) | + +## 2. Why a contract first + +The driver already produces rich per-operation diagnostics and hands them back through a +cheap handle. Two things were underspecified, and the SDK requirements must drive the driver +design (not the other way around): + +1. **One implicit shape.** Materialization is JSON-only and implicit + ([`to_json_string`][ctx] [main]). Different consumers want different shapes: metrics want a + structured object, tracing wants spans, logs want a string. JSON is the single costliest + step (bench: struct build ~3.7 µs → +detailed JSON ~6.6 µs), so forcing it on everyone is + wasteful — and lossy/expensive across an FFI boundary. +2. **No agreed emission decision.** *Who* decides whether an operation is even worth a span or + a log line, and *where*? The driver must not emit eagerly (R6); the decision belongs to the + SDK, computed **after** the operation completes (R4). + +So we fix the **contract**, the **OpenTelemetry mapping**, and the **emission model** first; +the hot-path capture engine [WS7] is a deferrable optimization. + +## 3. The contract in one picture + +```text +azure_data_cosmos_driver azure_data_cosmos (SDK) +──────────────────────── ─────────────────────── +operation completes + │ + ▼ + Arc ── cheap handle ──▶ DiagnosticsHandler chain [WS2] + [main] ALWAYS collected, (atomic incr) ├─ CosmosMetricsHandler [WS3] always-on, low-cardinality + ALWAYS returned; the ├─ CosmosTracingHandler [WS4] tail-sampled, backdated spans + driver emits only └─ SamplingLogHandler [WS5] rate-limited log lines + debug-level `tracing` ▲ each handler consumes the COMPLETED context + itself — never eager │ and maps it to its own OTel/semconv view + OTel spans/metrics. +``` + +Four invariants: + +- **P1 — The handle is cheap and unconditional.** In Rust it *is* + `CosmosResponse::diagnostics() -> Arc` [main]. Cloning it is an atomic + increment. It is **always** returned; `diagnostics()` is **non-optional**. +- **P2 — The driver produces & materializes; the SDK emits.** The driver never starts spans at + operation-start and never talks to an OTel exporter on the default path. Emission is an SDK + concern, realized by the handler chain (§4). This is R6. +- **P3 — Emission is decided from a *completed* context.** Handlers run once, at operation + completion, so the SDK can decide *after the fact* — by latency/outcome — whether an op is + worth a span or a log line (tail-based sampling, §5.2). This is R4. +- **P4 — The level/threshold governs EXPOSURE and DEPTH, never COLLECTION.** A diagnostics + *level* bounds the high-cardinality **transport-level** detail a materialization includes; it + never removes operation-level diagnostics and never disables collection. + +There is **no parallel diagnostics model**: the driver owns one canonical +[`DiagnosticsContext`][ctx] [main] and every representation is a view over it. + +## 4. The emission model — the `DiagnosticsHandler` chain + +**The SDK owns an ordered chain of `DiagnosticsHandler`s; each consumes the completed context +and decides, independently, what (if anything) to emit.** This is the extension point that +realizes R6 (driver produces, SDK decides) and hosts the sampling (R4) and rate-limiting (R5) +policies. + +```rust +// azure_data_cosmos (SDK) — [WS2] +pub trait DiagnosticsHandler: Send + Sync { + /// Called once, at operation completion, with the completed driver context and the + /// active trace `Context`. Handlers decide independently whether/what to emit. + fn handle(&self, ctx: &DiagnosticsContext, cx: &Context); +} + +// Built-in handlers registered by the app: +struct CosmosMetricsHandler { /* [WS3] always-on, low-cardinality */ } +struct CosmosTracingHandler { /* [WS4] tail-sampled, backdated spans */ } +struct SamplingLogHandler { /* [WS5] rate-limited log lines */ } +``` + +Properties: + +- **Ordered, `Arc`, no-op when empty.** Zero registered handlers ⇒ zero + emission overhead (the driver still collects the context, but nothing is emitted). Ordering is + deterministic. +- **Cosmos-local now, promotable later (D1).** The chain is built **inside `azure_data_cosmos`** + as an **additive, swappable** surface — deliberately *not* over-fitted to the Java + `CosmosDiagnosticsHandler` shape. If other Azure SDKs want the same extension point, the + abstraction can be **promoted into `azure_core` later** without breaking Cosmos callers. +- **SDK maps to its own semconv view.** Each handler translates the driver's + `DiagnosticsContext` into its **own** OTel/semantic-convention representation (§10); the + driver's internal model never leaks into the emitted telemetry. + +This replaces the earlier "materializer-only" framing: materialization (§6) is still how a +handler *reads* the context, but the **decision + emission** now has an explicit home. + +## 5. Three orthogonal controls — do not conflate them + +The earlier draft partly merged these. They are **independent** knobs: + +### 5.1 `DiagnosticsLevel` — static **DEPTH** + +`DiagnosticsLevel { Minimal, Standard, Full }` [design] is a **static config** that bounds *how +much transport-level detail a materialization includes* (§7). It answers **"how deep?"** — not +"whether to emit" and not "how many". It never disables collection (P4). + +### 5.2 Tail-based sampling — dynamic **WHETHER** (R4) + +A **per-operation, late-bound decision**, computed from the *completed* context, that answers +**"emit a span/log for this op at all?"**: + +```rust +fn should_emit_span(ctx: &DiagnosticsContext, thresholds: &DiagnosticsThresholds) -> bool { + ctx.is_completed() + && (ctx.is_failure() || ctx.is_threshold_violated(thresholds) /* || level == Full */) +} +``` + +- **Default (D4):** a fast success (e.g. a 5 ms point read) emits **no span**; a slow op + (≥ threshold) or a failure **does**. This is exactly R4/R9 — no per-fast-op cost. +- **Thresholds use [`DiagnosticsThresholds`][thresholds] [WS4/WS5]** (added by this PR) and are + **configurable via the standard Rust options chain**; start from Java-like defaults, tunable before GA. +- This is the piece the earlier draft under-specified. It is distinct from `DiagnosticsLevel` + (§5.1): the level sets *depth of what a span contains*; tail-sampling sets *whether a span + exists*. + +### 5.3 Cross-operation rate limiting — bounded **HOW MANY** (R5) + +A **cross-operation cap** on emission frequency during a storm, so 10k errors/sec don't all get +logged/traced and peg the CPU: + +- A reusable **token-bucket / count-per-interval** limiter (default e.g. **≤ N lines/min**), + used by `SamplingLogHandler` [WS5] (and optionally by the tracing handler under storms). +- **Always allow a bounded number of failures** through, then emit **one** `"suppressed N until + reset"` line per window — so a storm is visible but bounded. +- **On by default (D5)** — production-ready by default. +- This is **independent** of §8's per-artifact size bound: §8 caps the size of a *single* + diagnostics artifact; §5.3 caps the *number of artifacts* emitted per interval. + +## 6. The three representations (views over one model) + +| Materializer | Consumer intent | Backed by | +| --- | --- | --- | +| **Structured object** | metrics | [`DiagnosticsContext::requests()`][ctx] [main], reduced into an operation roll-up | +| **OTel span tree** | traces | reconstructed from [`DiagnosticsContext`][ctx] + per-attempt [`RequestDiagnostics`][ctx] [main]; a `Span`/`Attr` in-memory form is the OTel-aligned shape (kept in [WS7]) | +| **String** | logs | [`DiagnosticsContext::to_json_string(verbosity)`][ctx] [main] | + +Materialization is **explicit, lazy, and per-representation**: each is paid only when a handler +asks for it, so the expensive JSON step is never paid on a metrics-only or span-only path. + +## 7. Where the level / threshold applies + +**Gating bounds high-cardinality TRANSPORT-level telemetry — it never eliminates +diagnostics.** + +| Tier | Examples | Cardinality | Gating | +| --- | --- | --- | --- | +| **Operation-level** | operation name, final status, request/retry/throttled counts, total RU, total duration, regions contacted | low | **Always on.** Never gated away. | +| **Transport-level** | per-replica / per-partition (partition key range, feed range), endpoint, direct-mode channel, transport kind/security/http-version, per-attempt RU/latency | high | **Gated by `DiagnosticsLevel` / threshold.** Included on error / slow / high level; summarized or elided on a fast-success low level. | + +`DiagnosticsLevel { Minimal, Standard, Full }` [design] maps onto the internal +`DiagnosticsVerbosity` [main]: + +- `Minimal` — operation-level only. +- `Standard` — + region-grouped/deduplicated transport summary (`Verbosity::Summary`). +- `Full` — + every per-attempt transport record (`Verbosity::Detailed`). + +> **Collection is not gated.** Operation-level metrics are *computed by iterating the +> per-attempt records* ([`requests()`][ctx] [main]), so "cheap op-level only" is not achievable +> by dropping transport collection. The driver **always collects** the full per-attempt records +> and the level gates only *materialization + exposure* (P4). + +## 8. Bounded-size guarantee (retry storms) + +**Every materialized representation has an upper bound on size that is independent of attempt +count** [WS6], so a `410`/`429` retry storm or a large fan-out query cannot produce an +unbounded object, span tree, or string. + +- **Mechanism.** `DiagnosticsVerbosity::Summary` [main] groups requests by region, keeps first + + last per region in full, and deduplicates the middle by `(endpoint, status, sub_status, + execution_context)` with count + min/max/P50. +- **Configurable caps [WS6] (salvaged from PR #4683).** The per-attempt record list is bounded + by a `DiagnosticsOptions::max_request_diagnostics` knob (default **512**, min 16), keyed on + `(region, endpoint, status, sub_status, execution_context)`. Consecutive near-identical + retries collapse to **first + last of each run + a count**; an order-robust global key-bucket + fallback holds the bound under a region ping-pong (`A→B→A→B`); the per-run rollup is itself + bounded so a high-cardinality `410` fan-out cannot grow the artifact by topology. +- **The span tree derives from the bounded record list.** The span emitter reconstructs one + child span per *retained* `RequestDiagnostics` attempt, so the span count is bounded by the + same `DiagnosticsOptions::max_request_diagnostics` cap (default **512**) — there is no + independent span cap or run-collapse today. Collapsing each run into a single count-bearing + span, with its own span-count target, is a possible future optimization and is not yet + implemented. +- **Truncation is marked, never silent** — the structured object carries explicit compaction + metadata (true vs retained count, per-run rollup, and a truncation marker when a cap is hit). + +This is the **per-artifact size bound**; it is distinct from the **cross-operation rate limit** +(§5.3), which caps how *many* artifacts are emitted per interval. The append-only capture +engine that realizes this cheaply is the deferred optimization [WS7]. + +## 9. The FFI boundary (Go / Python next-major) + +**How diagnostics flow over FFI: an opaque handle + explicit materialize calls. Never force +full JSON at the boundary.** + +```text +Rust driver FFI (C ABI) Go / Python SDK +─────────── ─────────── ─────────────── +Arc ── boxed → DiagnosticsHandle (opaque ptr) hold cheaply; free later + cosmos_diag_materialize( + handle, representation, caller picks the shape + level, out_buf) + ├ Metrics → packed struct / flatbuffer + ├ Spans → span-tree buffer (or per-span callback) + └ Log → UTF-8 bytes (JSON / compact / encoded) + cosmos_diag_free(handle) explicit lifetime +``` + +Contract rules for the boundary: + +1. **Opaque handle.** The FFI surfaces an opaque pointer wrapping the `Arc`. Crossing the + boundary is a pointer move + refcount, **not** a serialization. +2. **Explicit materialize, consumer-chosen shape.** One `materialize(handle, representation, + level)` entry point where `representation ∈ {Metrics, Spans, Log}`. +3. **No forced JSON.** JSON is one representation requested explicitly; a Go metrics exporter + never triggers it. +4. **Explicit lifetime (opaque handle + `free`).** The handle is freed by an explicit + `cosmos_diag_free`; refcounting stays on the Rust side. A scoped materialize-then-drop + callback may wrap it later. +5. **Bounded output.** Every materialization honors the bounded-size guarantee (§8), so an FFI + buffer can be sized predictably even under a retry storm (read the configured cap, not a + hardcoded value). + +The exact wire encoding of each representation (packed struct vs flatbuffer vs span callback) +is a follow-up implementation detail; the contract fixes only "opaque handle + explicit +per-representation materialize + bounded output". + +## 10. OpenTelemetry mapping + +This mapping targets the **OpenTelemetry semantic conventions for Azure Cosmos DB** +(`open-telemetry/semantic-conventions`, `docs/db/cosmosdb.md`). The names below are the **real** +semconv names. + +> ⚠ **Naming correction (was wrong in PR #4684).** The earlier draft used non-existent +> `db.cosmosdb.*` names (`db.cosmosdb.operation.duration`, `db.cosmosdb.status_code`, +> `db.cosmosdb.sub_status_code`, `db.cosmosdb.request_charge`, `db.operation`). The correct +> names are `db.client.operation.duration`, `db.response.status_code`, `db.operation.name`, +> with Cosmos-specifics under the `azure.cosmosdb.*` namespace. This section uses the correct +> names throughout. + +### 10.1 System, span name, span kind + +- `db.system.name = "azure.cosmosdb"` (set at span creation). +- **Span name** follows the DB span-name convention (e.g. `"read_item orders"` — operation + + target). +- **Span kind** = `CLIENT`. Span status follows the OTel *Recording Errors* guidance. + +### 10.2 Operation-level metrics (always-on, low cardinality) + +Source: an operation roll-up over [`DiagnosticsContext::requests()`][ctx] [main], emitted by +`CosmosMetricsHandler` [WS3]. **Emit the *stable* instruments first**, with only low-cardinality +attributes. This powers client-side Grafana dashboards (R7) with per-combination series (R8). + +| Instrument | Stability | Kind | Unit | Description | +| --- | --- | --- | --- | --- | +| `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).)* | + +**Always-on metric attributes (low cardinality, D7):** `db.operation.name`, +`db.response.status_code`, `db.collection.name`, `db.namespace`, `error.type`, `server.address`, +`db.system.name`. **High-cardinality attributes are OFF by default** (opt-in) — see §10.3. + +> **Metrics transport gap (D2).** Neither `typespec_client_core` nor `azure_core_opentelemetry` +> currently exposes a `Meter`/`Histogram`/`Counter` abstraction. `CosmosMetricsHandler` [WS3] +> therefore emits via the **raw `opentelemetry` metrics API behind an off-by-default feature** +> (Cosmos owns a `Meter`) now, and we drive a first-class `Meter` abstraction into `azure_core` +> as a follow-up, then migrate. **Develop Cosmos-local; move to `azure_core` if there is shared +> interest.** + +### 10.3 Attributes — stable vs. development tiers + +Cosmos spans (and, where low-cardinality, metrics) use these semconv attributes. The **stable** +tier is safe as always-on metric dimensions; the **development** tier is span-first / opt-in as +a metric dimension to control time-series cardinality (D7). + +**Stable** (safe as always-on metric dimensions) + +| Attribute | Notes | +| --- | --- | +| `db.operation.name` | Canonical snake_case op name — use verbatim (see §10.3.1). | +| `db.collection.name` | Cosmos container name. | +| `db.namespace` | Database name. | +| `db.response.status_code` | Cosmos status code (string). 4xx/5xx are errors. | +| `error.type` | Present iff the operation failed; matches status or exception type. | +| `server.address` / `server.port` | Host / port (port only when not default 443). | +| `db.operation.batch.size` | Number of ops in a batch. | +| `db.query.text` | Query text (parameterized; sanitized per DB span rules). | +| `db.stored_procedure.name` | Stored-procedure name, when applicable. | +| `user_agent.original` | SDK-generated user-agent string. | + +**Development** (span-first; opt-in as metric dimensions) + +| Attribute | Notes | +| --- | --- | +| `azure.cosmosdb.connection.mode` | `gateway` or `direct`. *(Deferred — not yet populated by the automatic operation path; [#4789](https://github.com/Azure/azure-sdk-for-rust/pull/4789).)* | +| `azure.cosmosdb.consistency.level` | `Eventual` / `Session` / `Strong` / `BoundedStaleness` / `ConsistentPrefix`. *(Deferred — not yet populated by the automatic operation path; [#4789](https://github.com/Azure/azure-sdk-for-rust/pull/4789).)* | +| `azure.cosmosdb.operation.contacted_regions` | Ordered list of contacted regions (cross-region call if > 1). | +| `azure.cosmosdb.operation.request_charge` | RU consumed (double). | +| `azure.cosmosdb.response.sub_status_code` | Cosmos sub-status code (int). | +| `db.response.returned_rows` | Row count in the result set (int). | +| `azure.cosmosdb.request.body.size` | Request payload size in bytes. | +| `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.** + +#### 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`, … . +If none applies, use a language-agnostic snake_case method name. + +### 10.4 Traces (span tree) — reconstructed & backdated + +Emitted by `CosmosTracingHandler` [WS4], **only when tail-sampling (§5.2) says so**. + +| `DiagnosticsContext` element | OTel span | +| --- | --- | +| operation (root) | root span, kind `CLIENT`, name ` `; window backdated over [`duration()`][ctx] [main] | +| each retained `RequestDiagnostics` (attempt / hedge leg) | child span, kind `CLIENT` | +| request-event timeline | timed span **events** on the attempt span | +| hedging | a hedge span with terminal state + regions | +| aggregated multi-run op | a **single synthetic operation root** spanning the first run's start to the last run's end, each run's attempts as children | + +- **Backdating (Gap A / D3).** A *completed* `DiagnosticsContext` is reconstructed into a + **backdated** span tree — because the current `typespec_client_core` `Tracer`/`Span` traits + build spans at "now" and have **no** hook to set explicit start/end timestamps. Two-track + plan: **(a) upstream** — add `start_span_at`/`end_at` (or a `SpanBuilder`-style API) to the + `typespec_client_core` traits + `azure_core_opentelemetry` impl [WS4c]; **(b) Cosmos-local + fallback** — the raw `opentelemetry` `SpanBuilder` (`with_start_time` / `end_with_timestamp`), + proven by the salvaged PR #4685 spike. **Ship (b) now behind one internal seam; migrate to (a) + when it lands.** +- **Build from the bounded list.** The emitter builds the tree from the retained attempt list + ([`requests()`][ctx], bounded by §8), emitting **one child span per retained attempt** — never + one span per pre-compaction attempt (the retained list is already bounded by + `max_request_diagnostics`). Collapsing each run into a single count-bearing span is a possible + future optimization and is not yet implemented. + +### 10.5 `azure_core` attribute alignment & the two upstream gaps + +Where a Cosmos span overlaps `azure_core`-emitted HTTP spans, reuse the `azure_core` +span-attribute names so the spans correlate: + +| Diagnostics field | `azure_core` attribute name | +| --- | --- | +| operation activity id | `az.client_request_id` | +| per-attempt service request id | `az.service_request.id` ⚠ **dot**, not underscore | +| HTTP status | `http.response.status_code` | +| retry index | `http.request.resend_count` | +| endpoint | `server.address` / `url.full` | +| namespace | `azure.resource_provider.namespace` (`Microsoft.DocumentDB`) | +| error | `error.type` | + +> ⚠ **Two `azure_core` gotchas (D2/D3).** (1) The constant is `az.service_request.id` (a **dot** +> before `id`), while `az.client_request_id` uses an underscore. (2) These constants are +> **module-private** in `azure_core` +> ([`http/policies/instrumentation/mod.rs`][azcore]) and cannot be imported. Until `azure_core` +> exposes them, centralize identical string literals in one Cosmos-local module and file an +> `azure_core` issue to expose public constants + fix the naming inconsistency. **Cosmos-specific +> attributes use the real `azure.cosmosdb.*` namespace** (not the earlier draft's `db.cosmosdb.*`). + +## 11. The OTel-aligned event model & the SDK↔driver split + +- **Keep an OTel-aligned in-memory span model.** `Span { kind, parent: Option, + start/end: TimeOffset }` + typed `Attr`s is an OpenTelemetry span tree in all but the emitter, + and is worth retaining as the canonical in-memory shape — a data shape, not a hot-path + commitment. It rides with the deferred capture engine [WS7]; the default path does not require + it. +- **SDK-vs-driver split (D8): default SDK-side emission.** The **driver produces and + materializes** the `DiagnosticsContext` (and may offer an *opt-in* exporter); the **SDK + emits** the public operation span + operation metrics via the handler chain (§4). Default + emission is SDK-side so there is exactly **one public span per operation** (no double-count). + Rule of thumb: *the driver produces and materializes; the SDK emits.* + +## 12. Resolved decisions + +These map to the plan's decisions D1–D10. + +| # | Decision | Resolution | +| --- | --- | --- | +| **D1** | Emission model | **`DiagnosticsHandler` chain** in the SDK (not a materializer-only API). Built Cosmos-local, additive/swappable, promotable to `azure_core` later. (§4) | +| **D2** | Metrics transport (Gap B) | No `Meter` abstraction exists in `azure_core`/`typespec_client_core`. Emit via the **raw `opentelemetry` metrics API behind an off-by-default feature** now; drive an `azure_core` `Meter` follow-up, then migrate. (§10.2) | +| **D3** | Span backdating (Gap A) | Two-track: **(a) upstream** trait additions to `typespec_client_core` + `azure_core_opentelemetry` [WS4c]; **(b) Cosmos-local** raw-`opentelemetry` `SpanBuilder` now. Ship (b) behind one seam; migrate to (a). (§10.4) | +| **D4** | Tail-sampling policy | Default: emit a span only when `is_completed` and (`is_failure` or `is_threshold_violated`). Fast successes emit no span; thresholds configurable via the standard options chain, Java-like defaults. (§5.2) | +| **D5** | Rate limiting | Token-bucket / count-per-interval limiter used by `SamplingLogHandler`; always allow a bounded number of failures + one "suppressed N" line per window. **On by default.** (§5.3) | +| **D6** | Semconv naming | Real names only: `db.client.operation.duration`; `db.operation.name`, `db.response.status_code`, `azure.cosmosdb.*`, `db.system.name="azure.cosmosdb"`. Centralize attribute literals Cosmos-local. (§10) | +| **D7** | Metric cardinality | Operation-scope tags always-on; high-cardinality tags (consistency level, per-request region, partition-key-range, endpoint, replica) off by default, opt-in. (§10.2/§10.3) | +| **D8** | SDK vs driver emission | Default **SDK-side** emission → exactly one public span per op; driver may expose an opt-in exporter. (§11) | +| **D9** | Bounded size vs capture engine | Land the **bounded-size guarantee** [WS6] as a customer-facing property; keep the append-only capture engine [WS7] as an OFF-by-default internal optimization. (§8) | +| **D10** | Client instance id | Stable per-client id feeds the active-instance metric + `azure.client.id`. Prefer `vmId`, fall back to a static GUID; check whether `DiagnosticsContext` already carries it. (§10.3) | + +## 13. Scope & guardrails + +- **Additive / non-breaking.** The public boundary is [`diagnostics::DiagnosticsContext`][ctx] + [main], consumed by `azure_data_cosmos`. `CosmosResponse::diagnostics()` stays non-optional. No + SemVer break. **This contract doc itself adds no code or public API; it is folded into the + combined observability PR (#4789), which carries the implementation and its CHANGELOG entries.** +- **Off-by-default / feature-gated.** New telemetry (the metrics feature, the capture engine) + is OFF by default (behind a feature or an unset exporter), so merging the implementation + workstreams cannot change existing behavior. +- **Diagnostics are always collected** — there is no full-disable mode; the level governs + exposure/depth, not collection (P4). +- **No `azure_core` / `typespec_client_core` change** is *required* by this contract; the + upstream `Meter` (D2) and span-backdating (D3) additions are proposals pursued in parallel. + +[thresholds]: https://github.com/Azure/azure-sdk-for-rust/blob/main/sdk/cosmos/azure_data_cosmos_driver/src/options/diagnostics_thresholds.rs +[ctx]: https://github.com/Azure/azure-sdk-for-rust/blob/main/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/diagnostics_context.rs +[azcore]: https://github.com/Azure/azure-sdk-for-rust/blob/main/sdk/core/azure_core/src/http/policies/instrumentation/mod.rs diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/compaction.rs b/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/compaction.rs new file mode 100644 index 00000000000..dddab7c698c --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/compaction.rs @@ -0,0 +1,402 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Bounding of an operation's per-attempt diagnostics under a retry storm. +//! +//! A single Cosmos operation appends one [`RequestDiagnostics`] record per +//! attempt to its [`DiagnosticsContext`](super::DiagnosticsContext). A pathological +//! retry storm (a partition hammered with `429`/`410` for the whole retry budget, +//! or a `410` fanned out across thousands of physical-partition endpoints) would +//! otherwise grow that list — and every artifact derived from it — without bound. +//! +//! This module bounds the *finalized* per-attempt list to at most +//! [`DiagnosticsOptions::max_request_diagnostics`](crate::options::DiagnosticsOptions::max_request_diagnostics) +//! records, while preserving the storm's shape and exact aggregates: +//! +//! - **Run-collapse (Phase 1):** consecutive near-identical retries — same +//! region, endpoint, status (incl. sub-status) and execution context — are +//! collapsed to their first and last record plus an exact per-run rollup +//! (count, total request charge, min/max/P50 duration). +//! - **Global-bucket fallback (Phase 2):** when Phase 1 alone does not fit within +//! the cap (e.g. a region ping-pong `A→B→A→B` where every consecutive run is +//! length one), attempts are grouped by key regardless of order and only the +//! `cap` largest buckets are kept for BOTH the retained records and the per-run +//! rollup, so they stay coherent and the serialized artifact is bounded by the +//! cap rather than by storm cardinality. +//! +//! Every drop is surfaced **explicitly** on [`CompactionInfo`] +//! (`retained_truncated`, `omitted_runs`, `omitted_request_count`) — never silent. +//! +//! Compaction runs only at finalization +//! ([`DiagnosticsContextBuilder::complete`](super::DiagnosticsContextBuilder)), never +//! mid-operation, so it never invalidates the index-based `RequestHandle`s handed +//! out while the operation is in flight. The bound is on the finalized serialized +//! artifact, not on live mid-operation memory. + +use std::collections::HashMap; + +use serde::Serialize; + +use super::diagnostics_context::{percentile_sorted, ExecutionContext, RequestDiagnostics}; +use crate::models::{CosmosStatus, RequestCharge}; +use crate::options::Region; + +/// Metadata describing how an operation's per-attempt diagnostics were compacted +/// under a retry storm. +/// +/// Present on a [`DiagnosticsContext`](super::DiagnosticsContext) only when the +/// number of attempts exceeded the configured +/// [`max_request_diagnostics`](crate::options::DiagnosticsOptions::max_request_diagnostics) +/// cap. It records the true attempt count, how many records were retained, and a +/// per-run rollup so the storm's shape (which region/endpoint/status repeated, +/// and how many times) is preserved even though the middle of each run was +/// dropped. +#[non_exhaustive] +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +pub struct CompactionInfo { + /// True total number of attempts before compaction. + pub original_request_count: usize, + /// Number of per-attempt records retained after compaction. + pub retained_request_count: usize, + /// Number of runs whose middle records were dropped (run length > 2). + pub collapsed_runs: usize, + /// Total number of distinct runs detected. Equal to `runs.len()` unless the + /// per-run rollup was itself bounded under a high-cardinality storm, in + /// which case `runs` holds only the largest ones and `omitted_runs` the rest. + pub total_runs: usize, + /// `true` when the retained per-attempt list hit the configured cap and + /// later records were dropped (the global-bucket fallback under an + /// order-ping-pong storm with more than `cap / 2` distinct keys). The + /// dropped attempts are still counted in `original_request_count` and the + /// aggregate rollup; this flag makes that truncation **explicit**, never + /// silent. + #[serde(default, skip_serializing_if = "bool_is_false")] + pub retained_truncated: bool, + /// Number of runs omitted from `runs` because the per-run rollup was bounded + /// to keep the serialized artifact size independent of storm *cardinality* + /// (e.g. a `410` fan-out across thousands of physical-partition endpoints). + /// `0` when every run is present. + #[serde(default, skip_serializing_if = "usize_is_zero")] + pub omitted_runs: usize, + /// Total attempt count represented by the omitted runs (see `omitted_runs`). + /// These attempts remain reflected in `original_request_count`; only their + /// per-run rollup rows were elided. + #[serde(default, skip_serializing_if = "usize_is_zero")] + pub omitted_request_count: usize, + /// Per-run rollup, in operation order (or first-seen order under the + /// global-bucket fallback). Bounded to the largest runs by attempt count + /// under a high-cardinality storm; see `omitted_runs` for the remainder. + pub runs: Vec, +} + +/// A single collapsed run of near-identical retries. +/// +/// Groups attempts that share the same region, endpoint, status (including +/// sub-status) and execution context. The first and last attempt of each run +/// are retained in full in +/// [`DiagnosticsContext::requests`](super::DiagnosticsContext::requests); this +/// rollup carries the count and duration/charge statistics for the run so the +/// elided middle is still accounted for. +#[non_exhaustive] +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +pub struct CompactedRun { + /// Normalized region name for the run, if the attempts carried one. + #[serde(skip_serializing_if = "Option::is_none")] + pub region: Option, + /// Endpoint the run targeted. + pub endpoint: String, + /// HTTP status and Cosmos sub-status shared by the run. + #[serde(flatten)] + pub status: CosmosStatus, + /// Execution context shared by the run. + pub execution_context: ExecutionContext, + /// Number of attempts in the run. + pub count: usize, + /// Total request charge (RU) across the run. + pub total_request_charge: RequestCharge, + /// Minimum per-attempt duration in the run, in milliseconds. + pub min_duration_ms: u64, + /// Maximum per-attempt duration in the run, in milliseconds. + pub max_duration_ms: u64, + /// Median (P50) per-attempt duration in the run, in milliseconds. + pub p50_duration_ms: u64, +} + +/// The key that defines a "near-identical" run: same region, endpoint, status +/// (incl. sub-status) and execution context. +#[derive(Clone, PartialEq, Eq, Hash)] +struct CompactionKey { + region: Option, + endpoint: String, + status: CosmosStatus, + execution_context: ExecutionContext, +} + +impl CompactionKey { + fn of(req: &RequestDiagnostics) -> Self { + Self { + region: req.region().cloned(), + endpoint: req.endpoint().to_string(), + status: *req.status(), + execution_context: req.execution_context(), + } + } +} + +/// The retained per-attempt records plus the per-run rollup produced by +/// compaction. +pub(super) struct CompactionResult { + pub(super) retained: Vec, + pub(super) runs: Vec, + pub(super) collapsed_runs: usize, + /// `true` when `retained` was truncated to the cap (records dropped beyond + /// the first+last-per-run set); surfaced on `CompactionInfo::retained_truncated`. + pub(super) retained_truncated: bool, + /// Total number of distinct runs detected, before the per-run rollup was + /// bounded; surfaced on `CompactionInfo::total_runs`. + pub(super) total_runs: usize, + /// Number of runs dropped from `runs` because the rollup was bounded to the + /// cap; surfaced on `CompactionInfo::omitted_runs`. + pub(super) omitted_runs: usize, + /// Total attempt count represented by the omitted runs; surfaced on + /// `CompactionInfo::omitted_request_count`. + pub(super) omitted_request_count: usize, +} + +/// `skip_serializing_if` helper: a `0` count is omitted from compaction output. +fn usize_is_zero(n: &usize) -> bool { + *n == 0 +} + +/// `skip_serializing_if` helper: a `false` flag is omitted from compaction output. +fn bool_is_false(b: &bool) -> bool { + !*b +} + +/// Builds a [`CompactedRun`] rollup from a run/bucket of attempts. +fn compacted_run(reqs: &[&RequestDiagnostics]) -> CompactedRun { + let count = reqs.len(); + let mut durations: Vec = reqs.iter().map(|r| r.duration_ms()).collect(); + durations.sort_unstable(); + let total_request_charge: RequestCharge = reqs.iter().map(|r| r.request_charge()).sum(); + let first = reqs[0]; + CompactedRun { + region: first.region().map(|r| r.as_str().to_string()), + endpoint: first.endpoint().to_string(), + status: *first.status(), + execution_context: first.execution_context(), + count, + total_request_charge, + min_duration_ms: durations.first().copied().unwrap_or(0), + max_duration_ms: durations.last().copied().unwrap_or(0), + p50_duration_ms: percentile_sorted(&durations, 50), + } +} + +/// Pushes the first and (when the run has more than one attempt) last record of +/// a run into `retained`. +fn push_first_last(retained: &mut Vec, run: &[&RequestDiagnostics]) { + if let Some(first) = run.first() { + retained.push((*first).clone()); + } + if run.len() > 1 { + if let Some(last) = run.last() { + retained.push((*last).clone()); + } + } +} + +/// Bounds an operation's per-attempt diagnostics to at most `cap` records. +/// +/// Phase 1 collapses runs of **consecutive** near-identical retries in order +/// (the common same-region storm). If that alone does not fit within `cap` +/// (e.g. a region ping-pong `A→B→A→B` where every consecutive run is length +/// one), Phase 2 falls back to a global key-bucket collapse that groups all +/// attempts by key regardless of order and keeps only the `cap` largest buckets +/// (by attempt count) for BOTH the retained records and the per-run rollup, so +/// the two stay coherent and the serialized artifact is bounded by `cap` rather +/// than by storm cardinality. Every drop is surfaced explicitly +/// (`retained_truncated`, `omitted_runs`, `omitted_request_count`), never silent. +pub(super) fn compact_requests(requests: Vec, cap: usize) -> CompactionResult { + let run_length = run_length_compact(&requests); + if run_length.retained.len() <= cap { + return run_length; + } + global_bucket_compact(&requests, cap) +} + +/// Phase 1: order-preserving run-length collapse of consecutive equal-key runs. +fn run_length_compact(requests: &[RequestDiagnostics]) -> CompactionResult { + let mut retained = Vec::new(); + let mut runs = Vec::new(); + let mut collapsed_runs = 0usize; + + let mut i = 0; + while i < requests.len() { + let key = CompactionKey::of(&requests[i]); + let mut j = i + 1; + while j < requests.len() && CompactionKey::of(&requests[j]) == key { + j += 1; + } + let run: Vec<&RequestDiagnostics> = requests[i..j].iter().collect(); + runs.push(compacted_run(&run)); + push_first_last(&mut retained, &run); + if run.len() > 2 { + collapsed_runs += 1; + } + i = j; + } + + let total_runs = runs.len(); + CompactionResult { + retained, + runs, + collapsed_runs, + retained_truncated: false, + total_runs, + omitted_runs: 0, + omitted_request_count: 0, + } +} + +/// Phase 2: order-robust global key-bucket collapse, bounded by `cap`. +/// +/// Groups every attempt by [`CompactionKey`] (first-seen order preserved for +/// output), then bounds BOTH the retained records and the per-run rollup with a +/// single ranking so they stay coherent. When there are more than `cap` distinct +/// keys (a high-cardinality `410` fan-out across physical-partition endpoints), +/// the buckets holding the operation-wide **first** and **final** attempts are +/// reserved first, then the remaining slots go to the buckets with the highest +/// attempt count (tie-break first-seen order); the rest are rolled into an +/// explicit `(omitted_runs, omitted_request_count)` marker. Both the per-run +/// rollup and the retained first+last records are drawn from that SAME kept set, +/// so every retained record has a matching run in the rollup — a downstream span +/// emitter never sees an attempt whose run was omitted. +/// +/// The operation-wide first attempt is emitted as the first retained record and +/// the operation-wide terminal attempt as the last, and both are preserved +/// through the final truncation that keeps `retained.len() <= cap` (marked via +/// `retained_truncated`, never silent). This matters because downstream fallbacks +/// treat the last retained request as terminal — `effective_status()` / +/// `is_failure()` when no operation status was stamped, and the reconstructed +/// span's end time — so the true terminal must never be ranked or truncated out. +fn global_bucket_compact(requests: &[RequestDiagnostics], cap: usize) -> CompactionResult { + let mut order: Vec = Vec::new(); + let mut groups: HashMap> = HashMap::new(); + for req in requests { + let key = CompactionKey::of(req); + match groups.get_mut(&key) { + Some(bucket) => bucket.push(req), + None => { + order.push(key.clone()); + groups.insert(key, vec![req]); + } + } + } + + let total_runs = order.len(); + let counts: Vec = order.iter().map(|key| groups[key].len()).collect(); + + // `collapsed_runs` counts every run whose middle was dropped (length > 2), + // whether or not it survives the rollup bound. + let collapsed_runs = counts.iter().filter(|&&c| c > 2).count(); + + // The buckets holding the operation-wide FIRST and FINAL attempts must always + // survive, independent of bucket ranking: downstream fallbacks treat the last + // retained request as terminal — `effective_status()` / `is_failure()` fall + // back to it when no operation status was stamped, and the reconstructed + // span's end time is anchored to it. The first-seen key `order[0]` owns the + // first attempt; the globally-last attempt's key owns the final attempt. + let last_key = CompactionKey::of(requests.last().expect("requests is non-empty")); + let last_index = order + .iter() + .position(|key| *key == last_key) + .expect("last attempt's key is present in first-seen order"); + + // One ranking drives the kept records and the kept rollup so they stay + // coherent, but the first and final buckets are reserved first so the op-wide + // first/terminal attempts are never ranked out. `keep[i]` marks whether + // `order[i]` survives. + let keep: Vec = if total_runs > cap { + let mut kept = vec![false; total_runs]; + kept[0] = true; + kept[last_index] = true; + let reserved = kept.iter().filter(|k| **k).count(); + // Fill the remaining slots with the largest non-reserved buckets + // (tie-break first-seen index), staying within `cap`. + let remaining = cap.saturating_sub(reserved); + if remaining > 0 { + let mut ranked: Vec = (0..total_runs).filter(|&i| !kept[i]).collect(); + ranked.sort_by(|&a, &b| counts[b].cmp(&counts[a]).then(a.cmp(&b))); + for &i in ranked.iter().take(remaining) { + kept[i] = true; + } + } + kept + } else { + vec![true; total_runs] + }; + + let mut omitted_runs = 0usize; + let mut omitted_request_count = 0usize; + for (i, &count) in counts.iter().enumerate() { + if !keep[i] { + omitted_runs += 1; + omitted_request_count += count; + } + } + + // Per-run rollup for every kept bucket, in first-seen order. Because + // `retained` is drawn only from kept buckets, its keys are a subset of the + // run rollup keys (coherent by construction). + let mut runs = Vec::new(); + for (i, key) in order.iter().enumerate() { + if keep[i] { + runs.push(compacted_run(&groups[key])); + } + } + + // Retained first+last records for kept buckets, in first-seen order — but + // hold the operation-wide terminal attempt back so it can be appended as the + // final record. That keeps it terminal even when its bucket isn't the last in + // first-seen order, and even after the truncation below. + let mut retained = Vec::new(); + for (i, key) in order.iter().enumerate() { + if !keep[i] { + continue; + } + let bucket = &groups[key]; + if *key == last_key { + // Push the bucket's first record (when the run has more than one + // attempt); its last record is the deferred terminal, appended below. + if bucket.len() > 1 { + if let Some(first) = bucket.first() { + retained.push((*first).clone()); + } + } + } else { + push_first_last(&mut retained, bucket); + } + } + let terminal = requests.last().expect("requests is non-empty").clone(); + + // Even limited to kept buckets, first+last records (plus the terminal) can + // exceed the record cap. Truncate the middle, always preserving the op-wide + // first (index 0) and the terminal (appended last); the drop is surfaced via + // `retained_truncated` and the omitted attempts stay exact in `CompactionInfo`. + let retained_truncated = retained.len() + 1 > cap; + if retained_truncated { + retained.truncate(cap.saturating_sub(1)); + } + retained.push(terminal); + + CompactionResult { + retained, + runs, + collapsed_runs, + retained_truncated, + total_runs, + omitted_runs, + omitted_request_count, + } +} 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 49f67a7e7fb..94cf02a77e8 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 @@ -10,7 +10,7 @@ use crate::{ driver::{pipeline::hedging_diagnostics::HedgeDiagnostics, routing::CosmosEndpoint}, models::{ActivityId, CosmosResponseHeaders, CosmosStatus, RequestCharge, SubStatusCode}, - options::{DiagnosticsOptions, DiagnosticsVerbosity, Region}, + options::{DiagnosticsOptions, DiagnosticsThresholds, DiagnosticsVerbosity, Region}, system::CpuMemoryMonitor, }; use azure_core::http::StatusCode; @@ -21,6 +21,29 @@ use std::{ time::{Duration, Instant}, }; +use super::compaction::{compact_requests, CompactedRun, CompactionInfo}; + +// ============================================================================= +// Threshold breach classification +// ============================================================================= + +/// The specific sampling threshold a completed operation crossed. +/// +/// Returned by +/// [`DiagnosticsContext::threshold_breach_for`](DiagnosticsContext::threshold_breach_for) +/// so emission handlers can record *why* a diagnostic was sampled (which bound +/// was breached), not just that one was. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[non_exhaustive] +pub enum ThresholdBreach { + /// Latency exceeded the point-operation latency threshold. + PointLatency, + /// Latency exceeded the non-point-operation latency threshold. + NonPointLatency, + /// Total request charge exceeded the request-charge (RU) threshold. + RequestCharge, +} + // ============================================================================= // Execution Context // ============================================================================= @@ -498,6 +521,53 @@ impl RequestDiagnostics { } } + /// **Internal test helper — do not call.** + /// + /// Builds a completed [`RequestDiagnostics`] entry with explicit endpoint, + /// region, status, charge, and start/completion instants, so emission-layer + /// tests can synthesize realistic (and backdated) attempt spans. Gated + /// behind the `__internal_test_diagnostics_construction` Cargo feature. + #[cfg(feature = "__internal_test_diagnostics_construction")] + #[doc(hidden)] + pub fn for_testing( + endpoint: impl Into, + region: Option, + status: CosmosStatus, + request_charge: RequestCharge, + started_at: Instant, + completed_at: Instant, + ) -> Self { + let duration_ms = completed_at + .saturating_duration_since(started_at) + .as_millis() as u64; + Self { + execution_context: ExecutionContext::Initial, + pipeline_type: PipelineType::DataPlane, + transport_security: TransportSecurity::Secure, + transport_kind: TransportKind::Gateway, + transport_http_version: TransportHttpVersion::Http2, + region, + endpoint: endpoint.into(), + status, + request_charge, + activity_id: None, + session_token: None, + server_duration_ms: None, + started_at, + completed_at: Some(completed_at), + duration_ms, + events: Vec::new(), + transport_shard: None, + failed_transport_shards: Vec::new(), + local_shard_retry_count: 0, + timed_out: false, + request_sent: RequestSentStatus::Sent, + error: None, + #[cfg(feature = "fault_injection")] + fault_injection_evaluations: Vec::new(), + } + } + /// Records completion of this request. /// /// Since we received a response, the request was definitely sent. @@ -1055,6 +1125,10 @@ struct DiagnosticsOutput<'a> { system_usage: Option, #[serde(skip_serializing_if = "Option::is_none")] machine_id: Option<&'a str>, + /// Present only when the per-attempt list was compacted under a retry storm; + /// absent (and thus byte-identical to prior output) otherwise. + #[serde(skip_serializing_if = "Option::is_none")] + compaction: Option<&'a CompactionInfo>, #[serde(flatten)] payload: DiagnosticsPayload<'a>, } @@ -1116,9 +1190,45 @@ struct TruncatedOutput<'a> { total_duration_ms: u64, request_count: usize, truncated: bool, + /// Present only when the per-attempt list was compacted under a retry storm. + /// + /// Counts-only: the per-run rollup (`CompactionInfo::runs`, up to `cap` + /// endpoint-bearing entries) is deliberately omitted here — it is the + /// unbounded part that can blow the size budget, and re-serializing it in the + /// size-limited fallback is exactly what would keep the "truncated" summary + /// oversized. + #[serde(skip_serializing_if = "Option::is_none")] + compaction: Option, message: &'static str, } +/// Counts-only projection of [`CompactionInfo`] for the size-limited truncated +/// summary. Carries the scalar counts (always tiny) but never the per-run rollup. +#[derive(Serialize)] +struct CompactionSummary { + original_request_count: usize, + retained_request_count: usize, + collapsed_runs: usize, + total_runs: usize, + retained_truncated: bool, + omitted_runs: usize, + omitted_request_count: usize, +} + +impl From<&CompactionInfo> for CompactionSummary { + fn from(info: &CompactionInfo) -> Self { + Self { + original_request_count: info.original_request_count, + retained_request_count: info.retained_request_count, + collapsed_runs: info.collapsed_runs, + total_runs: info.total_runs, + retained_truncated: info.retained_truncated, + omitted_runs: info.omitted_runs, + omitted_request_count: info.omitted_request_count, + } + } +} + /// Status of the CPU sample history in a [`SystemUsageSnapshot`]. #[derive(Clone, Copy, Debug, Serialize)] #[serde(rename_all = "snake_case")] @@ -1562,19 +1672,70 @@ impl DiagnosticsContextBuilder { /// shared via `Arc` without any locking overhead. pub(crate) fn complete(self) -> DiagnosticsContext { let duration = self.started_at.elapsed(); + + // Exact operation-level total charge, summed from the FULL attempt list + // before any compaction so it stays exact even when the retained list is + // bounded under a retry storm. + let total_request_charge: RequestCharge = + self.requests.iter().map(|r| r.request_charge).sum(); + + // Capture the contacted regions in first-contact order from the FULL + // attempt list, before compaction can drop whole buckets: a region whose + // only attempts are elided must still surface at the operation level. + let regions_contacted = ordered_unique_regions(&self.requests); + + // Bound the finalized per-attempt list under a retry storm. + // + // Common path (attempts <= cap): the list is retained verbatim, no + // `CompactionInfo` is attached, and the serialized output is + // byte-identical to the pre-compaction behavior. + // + // Storm path (attempts > cap): run-length collapse (with a global + // key-bucket fallback for order ping-pong) bounds the retained records + // and per-run rollup to `cap`, while a `CompactionInfo` marker records + // the true attempt count, the exact per-run aggregates, and every drop. + // + // Compaction runs here at finalization, never mid-operation, so any + // outstanding `RequestHandle` indices are never invalidated. The bound + // is on the finalized serialized artifact, not on live mid-operation + // memory: `self.requests` still grows one entry per attempt while the + // operation is in flight. + let cap = self.options.max_request_diagnostics(); + let original_count = self.requests.len(); + let (requests, compaction) = if original_count > cap { + let compacted = compact_requests(self.requests, cap); + let info = CompactionInfo { + original_request_count: original_count, + retained_request_count: compacted.retained.len(), + collapsed_runs: compacted.collapsed_runs, + total_runs: compacted.total_runs, + retained_truncated: compacted.retained_truncated, + omitted_runs: compacted.omitted_runs, + omitted_request_count: compacted.omitted_request_count, + runs: compacted.runs, + }; + (compacted.retained, Some(info)) + } else { + (self.requests, None) + }; + DiagnosticsContext { activity_id: self.activity_id, duration, - requests: Arc::new(self.requests), + requests: Arc::new(requests), + total_request_charge, + regions_contacted, status: self.status, options: self.options, cpu_monitor: self.cpu_monitor, machine_id: self.machine_id, + operation_name: None, #[cfg(feature = "fault_injection")] fault_injection_enabled: self.fault_injection_enabled, #[cfg(not(feature = "fault_injection"))] fault_injection_enabled: false, hedge_diagnostics: self.hedge_diagnostics, + compaction, #[cfg(test)] test_system_usage: self.test_system_usage, cached_json_detailed: OnceLock::new(), @@ -1629,9 +1790,28 @@ pub struct DiagnosticsContext { /// All request diagnostics (shared via `Arc` for efficient multi-read). /// /// `Vec` in Rust guarantees insertion order, so requests are stored in - /// the order they were added. + /// the order they were added. Under a retry storm this list is compacted at + /// finalization to at most + /// [`max_request_diagnostics`](crate::options::DiagnosticsOptions::max_request_diagnostics) + /// records; see [`compaction`](Self::compaction). requests: Arc>, + /// Total request charge (RU) across **all** attempts. + /// + /// Computed from the full attempt list at finalization, before any + /// compaction, so it stays exact even when `requests` was bounded under a + /// retry storm. + total_request_charge: RequestCharge, + + /// Regions contacted during the operation, in first-contact order. + /// + /// Captured at finalization from the **full** attempt list — before any + /// retry-storm compaction — so a region whose only attempts were dropped + /// from `requests` is still reported. Duplicates are removed while + /// preserving the order in which each region was first contacted, which the + /// Cosmos semantic conventions require (it conveys failover order). + regions_contacted: Vec, + /// Operation-level combined HTTP status and sub-status (final status after retries). status: Option, @@ -1644,6 +1824,16 @@ pub struct DiagnosticsContext { /// Machine identifier (VM ID on Azure, generated UUID otherwise). machine_id: Option>, + /// 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 + /// [`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. + operation_name: Option>, + /// Whether fault injection was enabled when this operation executed. fault_injection_enabled: bool, @@ -1660,6 +1850,13 @@ pub struct DiagnosticsContext { #[cfg(test)] test_system_usage: Option, + /// Compaction metadata, present only when the per-attempt list exceeded the + /// configured `max_request_diagnostics` cap under a retry storm and was + /// compacted. `None` for normal operations, where `requests` is the full, + /// unmodified set of attempts (and the serialized output is byte-identical + /// to the pre-compaction behavior). + compaction: Option, + /// Cached JSON string for detailed verbosity. cached_json_detailed: OnceLock, @@ -1691,6 +1888,80 @@ impl DiagnosticsContext { .complete() } + /// **Internal escape hatch — do not call.** + /// + /// Synthesizes a completed [`DiagnosticsContext`] carrying an explicit + /// operation `duration` and final `status`, so the wrapper SDK + /// (`azure_data_cosmos`) can exercise emission-layer code paths — such as + /// the metrics handler — against a realistic operation-level rollup without + /// standing up the full driver pipeline. Per-request diagnostics remain + /// empty; only the operation-scope fields the emission layer reads are set. + /// + /// Gated behind the `__internal_test_diagnostics_construction` Cargo feature + /// (enabled only by the wrapper SDK's own test build) and `#[doc(hidden)]`, + /// so it never appears on the public surface. Mirrors [`for_testing`](Self::for_testing). + #[cfg(feature = "__internal_test_diagnostics_construction")] + #[doc(hidden)] + pub fn for_testing_completed( + activity_id: ActivityId, + duration: Duration, + status: Option, + ) -> Self { + let mut context = + DiagnosticsContextBuilder::new(activity_id, Arc::new(DiagnosticsOptions::default())) + .complete(); + context.duration = duration; + context.status = status; + context + } + + /// **Internal test helper — do not call.** + /// + /// Builds a fully-populated [`DiagnosticsContext`] from explicit parts so + /// the SDK's emission-layer tests (tracing, sampled logging) can exercise + /// realistic operations — including backdated per-request timestamps — + /// without a live driver pipeline. Gated behind the + /// `__internal_test_diagnostics_construction` Cargo feature and + /// `#[doc(hidden)]`, mirroring [`for_testing`](Self::for_testing). + #[cfg(feature = "__internal_test_diagnostics_construction")] + #[doc(hidden)] + pub fn for_testing_with_requests( + activity_id: ActivityId, + duration: Duration, + status: Option, + operation_name: Option<&str>, + requests: Vec, + ) -> Self { + // Mirror the pipeline's exact-at-finalization rollup by summing the + // per-attempt charges supplied by the test. + let total_request_charge = RequestCharge::new( + requests + .iter() + .map(|r| r.request_charge().value()) + .sum::(), + ); + let regions_contacted = ordered_unique_regions(&requests); + DiagnosticsContext { + activity_id, + duration, + requests: Arc::new(requests), + total_request_charge, + regions_contacted, + status, + options: Arc::new(DiagnosticsOptions::default()), + cpu_monitor: None, + machine_id: None, + operation_name: operation_name.map(Arc::from), + fault_injection_enabled: false, + hedge_diagnostics: None, + #[cfg(test)] + test_system_usage: None, + compaction: None, + cached_json_detailed: OnceLock::new(), + cached_json_summary: OnceLock::new(), + } + } + /// Concatenates the per-request diagnostics from a sequence of /// sub-operation contexts into a single aggregated [`DiagnosticsContext`]. /// @@ -1720,16 +1991,96 @@ impl DiagnosticsContext { .iter() .map(|c| c.duration) .fold(Duration::ZERO, |a, b| a.saturating_add(b)); + // Sum each source's exact total charge (which already accounts for any + // per-sub-op compaction) rather than re-summing the possibly-compacted + // concatenated records. + let total_request_charge: RequestCharge = + sources.iter().map(|c| c.total_request_charge).sum(); + + // Exact total attempts across all sub-ops. Each source's own count is + // already exact — even if that source was individually compacted — so + // summing the per-source counts keeps `request_count()` exact on the + // aggregate instead of collapsing to the retained-record count. + let original_request_count: usize = sources.iter().map(|c| c.request_count()).sum(); + let cap = last.options.max_request_diagnostics(); + + // Re-bound the concatenated retained records so the aggregate artifact + // stays within the cap regardless of how many sub-ops contributed, and + // attach a `CompactionInfo` whenever the retained records under-count the + // true attempts (from re-bounding here or from a sub-op's own + // compaction) so the exact original count is never lost. + let (requests, compaction) = if aggregated_requests.len() > cap { + let compacted = compact_requests(aggregated_requests, cap); + let retained_request_count = compacted.retained.len(); + let info = CompactionInfo { + original_request_count, + retained_request_count, + collapsed_runs: compacted.collapsed_runs, + total_runs: compacted.total_runs, + retained_truncated: compacted.retained_truncated, + omitted_runs: compacted.omitted_runs, + omitted_request_count: original_request_count + .saturating_sub(retained_request_count), + runs: compacted.runs, + }; + (compacted.retained, Some(info)) + } else if original_request_count > aggregated_requests.len() { + // The concatenation fits the cap, but at least one sub-op was itself + // compacted, so the retained records under-count the true attempts. + // Attach a counts-only marker (carrying the sub-ops' per-run rollup + // entries) so `request_count()` stays exact and the storm shape is + // preserved. + let retained_request_count = aggregated_requests.len(); + let source_infos = || sources.iter().filter_map(|c| c.compaction.as_ref()); + let runs: Vec = source_infos() + .flat_map(|info| info.runs.iter().cloned()) + .collect(); + let info = CompactionInfo { + original_request_count, + retained_request_count, + collapsed_runs: source_infos().map(|info| info.collapsed_runs).sum(), + total_runs: source_infos().map(|info| info.total_runs).sum(), + retained_truncated: false, + omitted_runs: source_infos().map(|info| info.omitted_runs).sum(), + omitted_request_count: original_request_count + .saturating_sub(retained_request_count), + runs, + }; + (aggregated_requests, Some(info)) + } else { + // No sub-op was compacted and the concatenation fits the cap: the + // aggregate is exact and verbatim. + (aggregated_requests, None) + }; + + // First-contact-ordered union of the sub-ops' contacted regions. Each + // source already captured its exact ordered regions from its full + // attempt list, so concatenating them in sub-op order (dedup preserving + // order) yields the operation-level failover order without re-deriving + // from the possibly-compacted concatenation. + let mut regions_contacted: Vec = Vec::new(); + for source in sources { + for region in &source.regions_contacted { + if !regions_contacted.contains(region) { + regions_contacted.push(region.clone()); + } + } + } + Some(DiagnosticsContext { activity_id: last.activity_id.clone(), duration: aggregated_duration, - requests: Arc::new(aggregated_requests), + requests: Arc::new(requests), + total_request_charge, + regions_contacted, status: last.status, options: Arc::clone(&last.options), cpu_monitor: last.cpu_monitor.clone(), machine_id: last.machine_id.clone(), + operation_name: last.operation_name.clone(), fault_injection_enabled: sources.iter().any(|c| c.fault_injection_enabled), hedge_diagnostics: None, + compaction, #[cfg(test)] test_system_usage: last.test_system_usage.clone(), cached_json_detailed: OnceLock::new(), @@ -1756,26 +2107,74 @@ impl DiagnosticsContext { self.status.as_ref() } + /// Returns the **effective** final status for the operation. + /// + /// This is the operation-level [`status`](Self::status) when one was recorded, + /// otherwise the terminal attempt's status — mirroring the fallback used by + /// [`is_failure`](Self::is_failure). Some driver error-finalization paths graft + /// diagnostics onto the returned error without stamping an operation-level + /// status; this accessor lets emitters (metrics/tracing) report an accurate + /// status/`error.type` for those failures instead of treating them as unknown. + /// `None` only when there is neither an operation status nor any attempt. + pub fn effective_status(&self) -> Option { + self.status + .or_else(|| self.requests.last().map(|request| *request.status())) + } + /// Returns the total request charge (RU) across all requests. + /// + /// This stays exact even under a retry storm: it is summed from the full + /// attempt list at finalization, before any compaction of + /// [`requests`](Self::requests). pub fn total_request_charge(&self) -> RequestCharge { - self.requests.iter().map(|r| r.request_charge).sum() + self.total_request_charge } /// Returns the number of requests made during this operation. + /// + /// This is always the **true** total number of attempts, even when the + /// per-attempt list was compacted under a retry storm. Use + /// [`retained_request_count`](Self::retained_request_count) for the number + /// of records actually retained in [`requests`](Self::requests). pub fn request_count(&self) -> usize { + self.compaction + .as_ref() + .map(|c| c.original_request_count) + .unwrap_or(self.requests.len()) + } + + /// Returns the number of per-attempt records retained in + /// [`requests`](Self::requests). + /// + /// Equal to [`request_count`](Self::request_count) for normal operations; + /// bounded by the configured + /// [`max_request_diagnostics`](crate::options::DiagnosticsOptions::max_request_diagnostics) + /// cap under a retry storm. + pub fn retained_request_count(&self) -> usize { self.requests.len() } - /// Returns all regions contacted during this operation. + /// Returns compaction metadata when a retry storm exceeded the configured + /// [`max_request_diagnostics`](crate::options::DiagnosticsOptions::max_request_diagnostics) + /// cap and the per-attempt list was compacted. + /// + /// `None` for normal operations, where the retained list is the full, + /// unmodified set of attempts. + pub fn compaction(&self) -> Option<&CompactionInfo> { + self.compaction.as_ref() + } + + /// Returns all regions contacted during this operation, in first-contact + /// order. + /// + /// The list is captured at finalization from the **full** attempt list — + /// before any retry-storm compaction — so a region whose only attempts were + /// dropped from [`requests`](Self::requests) is still reported here. + /// Duplicates are removed while preserving the order in which each region was + /// first contacted, as the Cosmos semantic conventions require for the + /// contacted-regions attribute (it conveys failover order). pub fn regions_contacted(&self) -> Vec { - let mut regions: Vec = self - .requests - .iter() - .filter_map(|r| r.region.clone()) - .collect(); - regions.sort(); - regions.dedup(); - regions + self.regions_contacted.clone() } /// Returns a shared reference to all request diagnostics. @@ -1822,6 +2221,115 @@ impl DiagnosticsContext { self.fault_injection_enabled } + /// 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). + pub fn operation_name(&self) -> Option<&str> { + self.operation_name.as_deref() + } + + /// Returns `true` when this context represents a finished operation. + /// + /// A [`DiagnosticsContext`] is immutable and finalized at construction, so + /// any context with a recorded final status or at least one request is + /// complete. This is the gate the emission handlers check before deciding + /// whether to emit. + pub fn is_completed(&self) -> bool { + self.status.is_some() || !self.requests.is_empty() + } + + /// Returns `true` when the operation completed with a non-success status. + /// + /// Derived from [`status`](Self::status): an operation is a failure when its + /// final [`CosmosStatus`] is not a success. When no operation-level status was + /// recorded — some driver error-finalization paths graft diagnostics onto the + /// returned error without first stamping the operation status — this falls + /// back to the terminal attempt's status, so a genuine failure still gates the + /// tail-sampled handlers. A context with neither a status nor any request is + /// not treated as a failure. + pub fn is_failure(&self) -> bool { + match self.status.as_ref() { + Some(status) => !status.is_success(), + None => self + .requests + .last() + .is_some_and(|request| !request.status().is_success()), + } + } + + /// Returns `true` when the operation crossed one of the sampling + /// [`thresholds`](DiagnosticsThresholds) — the tail-based sampling signal. + /// + /// The latency check uses the point-operation threshold when + /// [`operation_name`](Self::operation_name) identifies a single-item + /// operation, the non-point threshold when it identifies a query/batch/etc., + /// and — when the operation name is unknown — falls back to the (stricter) + /// point-operation threshold so genuinely slow operations are still caught. + /// + /// The request-charge threshold is compared against the operation's total + /// RU. The payload-size threshold is not evaluated yet (the context does not + /// carry body sizes). + pub fn is_threshold_violated(&self, thresholds: &DiagnosticsThresholds) -> bool { + self.is_threshold_violated_for(thresholds, None) + } + + /// 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. + pub fn is_threshold_violated_for( + &self, + thresholds: &DiagnosticsThresholds, + operation_name: Option<&str>, + ) -> bool { + self.threshold_breach_for(thresholds, operation_name) + .is_some() + } + + /// Like [`is_threshold_violated_for`](Self::is_threshold_violated_for), but + /// reports *which* threshold was crossed rather than just whether one was. + /// + /// Latency is checked before request charge, so when an operation is both + /// slow and expensive the latency breach is reported. Returns `None` when no + /// threshold was crossed. The point/non-point latency threshold is chosen + /// from `operation_name` exactly as in + /// [`is_threshold_violated_for`](Self::is_threshold_violated_for). + pub fn threshold_breach_for( + &self, + thresholds: &DiagnosticsThresholds, + operation_name: Option<&str>, + ) -> Option { + let operation_name = operation_name.or_else(|| self.operation_name()); + let (latency_threshold, latency_breach) = match operation_name { + Some(name) if crate::options::is_point_operation(name) => ( + thresholds.point_operation_latency(), + ThresholdBreach::PointLatency, + ), + Some(_) => ( + thresholds.non_point_operation_latency(), + ThresholdBreach::NonPointLatency, + ), + None => ( + thresholds.point_operation_latency(), + ThresholdBreach::PointLatency, + ), + }; + if self.duration > latency_threshold { + return Some(latency_breach); + } + if self.total_request_charge().value() > thresholds.request_charge() { + return Some(ThresholdBreach::RequestCharge); + } + None + } + /// Serializes diagnostics to a JSON string. /// /// The result is lazily cached - the first call computes the JSON, @@ -1866,10 +2374,11 @@ impl DiagnosticsContext { let output = DiagnosticsOutput { activity_id: &self.activity_id, total_duration_ms, - total_request_charge: self.requests.iter().map(|r| r.request_charge).sum(), - request_count: self.requests.len(), + total_request_charge: self.total_request_charge(), + request_count: self.request_count(), system_usage, machine_id: self.machine_id.as_ref().map(|s| s.as_str()), + compaction: self.compaction.as_ref(), payload: DiagnosticsPayload::Requests { requests: &self.requests, }, @@ -1902,10 +2411,11 @@ impl DiagnosticsContext { let output = DiagnosticsOutput { activity_id: &self.activity_id, total_duration_ms, - total_request_charge: self.requests.iter().map(|r| r.request_charge).sum(), - request_count: self.requests.len(), + total_request_charge: self.total_request_charge(), + request_count: self.request_count(), system_usage: self.resolve_system_usage(), machine_id: self.machine_id.as_ref().map(|s| s.as_str()), + compaction: self.compaction.as_ref(), payload: DiagnosticsPayload::Summary { regions: region_summaries, }, @@ -1922,8 +2432,9 @@ impl DiagnosticsContext { let truncated = TruncatedOutput { activity_id: &self.activity_id, total_duration_ms, - request_count: self.requests.len(), + request_count: self.request_count(), truncated: true, + compaction: self.compaction.as_ref().map(CompactionSummary::from), message: "Output truncated to fit size limit. Use Detailed verbosity for full diagnostics.", }; @@ -1939,12 +2450,16 @@ impl Clone for 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: self.operation_name.clone(), fault_injection_enabled: self.fault_injection_enabled, hedge_diagnostics: self.hedge_diagnostics.clone(), + compaction: self.compaction.clone(), #[cfg(test)] test_system_usage: self.test_system_usage.clone(), // OnceLock does not implement Clone, so we propagate any cached @@ -1968,12 +2483,20 @@ impl Clone for DiagnosticsContext { impl PartialEq for DiagnosticsContext { fn eq(&self, other: &Self) -> bool { // Compare semantic data only; cached JSON is derived and excluded. + // `total_request_charge` IS compared: after compaction it is no longer + // derivable from `requests` (dropped runs still carry charge), so + // excluding it would let two contexts with a different public + // `total_request_charge()` result compare equal. self.activity_id == other.activity_id && self.duration == other.duration && self.requests == other.requests + && self.total_request_charge == other.total_request_charge + && self.regions_contacted == other.regions_contacted && self.status == other.status && self.options == other.options + && self.operation_name == other.operation_name && self.hedge_diagnostics == other.hedge_diagnostics + && self.compaction == other.compaction } } @@ -2013,6 +2536,23 @@ impl std::fmt::Display for DiagnosticsContext { } } +/// Collects the regions contacted across `requests` in first-contact order. +/// +/// Duplicates are removed while preserving the order in which each region was +/// first seen. Unlike a sort+dedup this keeps the failover order intact, which +/// the Cosmos semantic conventions require for the contacted-regions attribute. +fn ordered_unique_regions(requests: &[RequestDiagnostics]) -> Vec { + let mut regions: Vec = Vec::new(); + for request in requests { + if let Some(region) = request.region() { + if !regions.contains(region) { + regions.push(region.clone()); + } + } + } + regions +} + /// Builds a summary for requests in a single region. fn build_region_summary( region: Option, @@ -2095,7 +2635,7 @@ fn deduplicate_requests(requests: Vec<&RequestDiagnostics>) -> Vec u64 { +pub(crate) fn percentile_sorted(values: &[u64], p: u8) -> u64 { if values.is_empty() { return 0; } @@ -2356,6 +2896,189 @@ mod tests { assert_eq!(regions.len(), 2); } + #[test] + fn regions_contacted_preserves_order_and_survives_compaction() { + // Regions must surface in first-contact (failover) order — not sorted — + // and a region contacted only by a bucket that global-bucket compaction + // drops from the retained per-attempt list must still be reported, + // because the set is captured from the full attempt list before + // compaction. Covers the metrics/tracing contacted-region attribute + // contract and the compaction bucket-drop gap. + let cap = 16; + let region_count = 20usize; + let mut b = DiagnosticsContextBuilder::new( + ActivityId::from_string("regions-storm".to_string()), + options_with_cap(cap), + ); + // Contact `region_count` distinct single-attempt region buckets in + // reverse-numeric order, so first-contact order differs from sorted + // order. Global compaction reserves the operation-wide first ("Region + // 19") and terminal ("Region 00") buckets and drops middle ones. + for i in (0..region_count).rev() { + record_run( + &mut b, + ExecutionContext::RegionFailover, + &format!("Region {i:02}"), + "https://acct/", + CosmosStatus::new(StatusCode::Gone), + 1.0, + 1, + ); + } + b.set_operation_status(StatusCode::Ok, None); + let ctx = b.complete(); + + let info = ctx + .compaction() + .expect("more distinct buckets than the cap must compact"); + assert!( + info.omitted_runs >= 1, + "at least one whole bucket must have been dropped from the retained list" + ); + + // First-contact order across all contacted regions, none lost. + let expected: Vec = (0..region_count) + .rev() + .map(|i| Region::new(format!("Region {i:02}"))) + .collect(); + let regions = ctx.regions_contacted(); + assert_eq!( + regions, expected, + "all contacted regions must survive compaction, in first-contact order" + ); + + // Prove the order is first-contact, not sorted. + let mut sorted = regions.clone(); + sorted.sort(); + assert_ne!( + regions, sorted, + "first-contact order must differ from a sorted list for these inputs" + ); + + // A region whose only attempt was a dropped MIDDLE bucket must still + // appear at the operation level even though it's gone from the retained + // list. "Region 01" is neither the first nor the last attempt, so it is + // eligible to be ranked out. + let retained_regions: Vec = ctx + .requests() + .iter() + .filter_map(|r| r.region().cloned()) + .collect(); + let dropped = Region::new("Region 01".to_string()); + assert!( + !retained_regions.contains(&dropped), + "a low-ranked middle bucket should have been dropped from the retained list" + ); + assert!( + regions.contains(&dropped), + "a region dropped from the retained list must still be reported" + ); + + // The operation-wide first and terminal attempts are reserved regardless + // of bucket ranking: their regions must survive in the retained list, and + // the terminal attempt ("Region 00") must be the LAST retained record so + // downstream status/span-end fallbacks see the true terminal. + let first = Region::new("Region 19".to_string()); + let terminal = Region::new("Region 00".to_string()); + assert!( + retained_regions.contains(&first), + "the operation-wide first attempt's bucket must be retained" + ); + assert_eq!( + retained_regions.last(), + Some(&terminal), + "the operation-wide terminal attempt must be the last retained record" + ); + } + + #[test] + fn effective_status_falls_back_to_terminal_attempt() { + // A context with no operation-level status but a failed terminal attempt + // must report that attempt's status as the effective status, so metrics + // and tracing can classify status-less error-finalization paths + // accurately instead of dropping the status / using the _OTHER catch-all. + let ctx = make_context_with(ActivityId::new_uuid(), |builder| { + let h = builder.start_test_request( + ExecutionContext::Initial, + Some(Region::WEST_US_2), + "https://test.westus2.documents.azure.com", + ); + builder.complete_request(h, StatusCode::TooManyRequests, None); + // Intentionally no set_operation_status: the operation-level status + // stays `None`, mirroring the graft-onto-error finalization paths. + }); + + assert!( + ctx.status().is_none(), + "no operation-level status should be set" + ); + assert_eq!( + ctx.effective_status().map(|s| s.status_code()), + Some(StatusCode::TooManyRequests), + "effective status must fall back to the terminal attempt" + ); + assert!(ctx.is_failure()); + } + + #[test] + fn size_limited_summary_fallback_omits_run_rollup() { + // Under a storm whose full summary exceeds the size budget, the truncated + // fallback must NOT re-serialize the (large) per-run rollup — re-emitting + // it is exactly what would keep the "truncated" summary oversized. It + // keeps counts only. + let cap = 64; + let options = Arc::new( + DiagnosticsOptions::builder() + .with_max_request_diagnostics(cap) + .with_max_summary_size_bytes(4096) + .build() + .expect("valid options"), + ); + let mut b = DiagnosticsContextBuilder::new( + ActivityId::from_string("storm-size".to_string()), + options, + ); + // Many distinct long-endpoint single-attempt buckets, so the per-run + // rollup (bounded to `cap`) alone is large enough that the full summary + // blows the size budget. + for i in 0..(cap + 4) { + let endpoint = format!( + "https://very-long-endpoint-{i:04}.documents.azure.com:443/padding/segment/path" + ); + record_run( + &mut b, + ExecutionContext::RegionFailover, + &format!("Region {i:04}"), + &endpoint, + CosmosStatus::new(StatusCode::Gone), + 1.0, + 1, + ); + } + b.set_operation_status(StatusCode::ServiceUnavailable, None); + let ctx = b.complete(); + assert!(ctx.compaction().is_some(), "the storm must compact"); + + let summary = ctx.to_json_string(Some(DiagnosticsVerbosity::Summary)); + assert!( + summary.contains("\"truncated\":true"), + "summary must have fallen back to the truncated form:\n{summary}" + ); + assert!( + !summary.contains("\"runs\""), + "the truncated summary must omit the per-run rollup:\n{summary}" + ); + assert!( + summary.contains("original_request_count"), + "the counts-only compaction summary must still be present:\n{summary}" + ); + assert!( + summary.len() <= 4096, + "the truncated summary must respect the size budget, was {} bytes", + summary.len() + ); + } + #[test] fn aggregate_sub_operations_concatenates_request_diagnostics() { // Concatenates sub-op RequestDiagnostics in input order, inherits @@ -2408,6 +3131,76 @@ mod tests { assert!(regions.contains(&Region::EAST_US_2)); } + #[test] + fn aggregate_sub_operations_preserves_exact_counts_when_sources_compacted() { + // Two sub-ops, each individually compacted past a small cap (the PATCH + // Read + Replace shape under a retry storm). The aggregate must report + // the exact total attempt count — the sum of the sub-ops' true counts, + // not just the retained records — and keep the combined artifact within + // the cap, rather than dropping the compaction metadata. + let cap = 16; + + let mut read = DiagnosticsContextBuilder::new( + ActivityId::from_string("agg-read".to_string()), + options_with_cap(cap), + ); + record_run( + &mut read, + ExecutionContext::Retry, + "East US", + "https://east/", + CosmosStatus::new(StatusCode::TooManyRequests), + 2.0, + 600, + ); + read.set_operation_status(StatusCode::Ok, None); + let read_ctx = Arc::new(read.complete()); + assert!( + read_ctx.compaction().is_some(), + "source must be individually compacted" + ); + + let mut replace = DiagnosticsContextBuilder::new( + ActivityId::from_string("agg-replace".to_string()), + options_with_cap(cap), + ); + record_run( + &mut replace, + ExecutionContext::Retry, + "West US", + "https://west/", + CosmosStatus::new(StatusCode::Gone), + 3.0, + 400, + ); + replace.set_operation_status(StatusCode::Created, None); + let replace_ctx = Arc::new(replace.complete()); + + let aggregated = + DiagnosticsContext::aggregate_sub_operations(&[read_ctx.clone(), replace_ctx.clone()]) + .expect("aggregation must succeed"); + + // Exact total across sub-ops, not just the retained records. + assert_eq!( + aggregated.request_count(), + read_ctx.request_count() + replace_ctx.request_count() + ); + assert_eq!(aggregated.request_count(), 1000); + // The combined artifact respects the cap. + assert!( + aggregated.retained_request_count() <= cap, + "retained {} exceeds cap {cap}", + aggregated.retained_request_count() + ); + // Exact total charge preserved (600*2.0 + 400*3.0). + assert!((aggregated.total_request_charge().value() - 2400.0).abs() < f64::EPSILON); + // Compaction metadata reports the exact original attempt count. + let info = aggregated + .compaction() + .expect("aggregate of compacted sources must carry compaction metadata"); + assert_eq!(info.original_request_count, 1000); + } + #[test] fn aggregate_sub_operations_returns_none_for_empty_input() { // Edge case: defensive None for callers that don't pre-check — @@ -3348,4 +4141,545 @@ mod tests { let ctx = builder.complete(); assert_eq!(ctx.machine_id(), None); } + + // ---- Compaction (retry-storm bounding, WS6) ------------------------------------------- + + fn options_with_cap(cap: usize) -> Arc { + Arc::new( + DiagnosticsOptions::builder() + .with_max_request_diagnostics(cap) + .build() + .expect("valid cap"), + ) + } + + /// Records `count` attempts sharing one (region, endpoint, status, exec-ctx) + /// key, each charged `charge` RU, so per-run aggregates are deterministic. + fn record_run( + builder: &mut DiagnosticsContextBuilder, + exec: ExecutionContext, + region: &str, + endpoint: &str, + status: CosmosStatus, + charge: f64, + count: usize, + ) { + for _ in 0..count { + let h = + builder.start_test_request(exec, Some(Region::new(region.to_string())), endpoint); + builder.update_request(h, |req| req.with_charge(RequestCharge::new(charge))); + builder.complete_request(h, status.status_code(), status.sub_status()); + } + } + + #[test] + fn retry_storm_429_is_bounded_and_lossless() { + // A single partition hammered with 429 for the whole retry budget: one + // run of many near-identical attempts. The retained list must be bounded + // by the cap while the true count and exact aggregates survive. + let cap = 16; + let mut b = DiagnosticsContextBuilder::new( + ActivityId::from_string("storm-429".to_string()), + options_with_cap(cap), + ); + record_run( + &mut b, + ExecutionContext::Retry, + "East US", + "https://east/", + CosmosStatus::new(StatusCode::TooManyRequests), + 2.0, + 1000, + ); + b.set_operation_status(StatusCode::TooManyRequests, None); + let ctx = b.complete(); + + // True total preserved; retained records bounded by the cap. + assert_eq!(ctx.request_count(), 1000); + assert!( + ctx.retained_request_count() <= cap, + "retained {} exceeds cap {cap}", + ctx.retained_request_count() + ); + + let info = ctx.compaction().expect("a storm past the cap must compact"); + assert_eq!(info.original_request_count, 1000); + assert_eq!(info.retained_request_count, ctx.retained_request_count()); + assert_eq!(info.runs.len(), 1); + assert_eq!(info.runs[0].count, 1000); + + // Aggregates stay EXACT despite the bounded retained list. + assert_eq!( + info.runs[0].total_request_charge, + RequestCharge::new(2000.0) + ); + assert_eq!(ctx.total_request_charge(), RequestCharge::new(2000.0)); + assert!(info.runs[0].min_duration_ms <= info.runs[0].p50_duration_ms); + assert!(info.runs[0].p50_duration_ms <= info.runs[0].max_duration_ms); + + // Truncation is visible in the serialized output and the output is bounded. + let json = ctx.to_json_string(Some(DiagnosticsVerbosity::Detailed)); + assert!(json.contains("\"compaction\""), "compaction marker missing"); + assert!(json.contains("\"original_request_count\":1000")); + assert!( + json.len() < 16 * 1024, + "detailed json {} bytes is not bounded", + json.len() + ); + + // First and last of the run are retained in full. + let requests = ctx.requests(); + assert_eq!( + u16::from(requests.first().unwrap().status().status_code()), + 429 + ); + assert_eq!( + u16::from(requests.last().unwrap().status().status_code()), + 429 + ); + } + + #[test] + fn mixed_429_410_runs_preserve_boundaries_and_exact_counts() { + // A 429 storm that escalates to a 410/1002 (PartitionKeyRangeGone) storm + // then recovers with a 200. Order-preserving run-length collapse keeps + // each run's boundaries and exact per-run aggregates. + let cap = 16; + let mut b = DiagnosticsContextBuilder::new( + ActivityId::from_string("mixed-429-410".to_string()), + options_with_cap(cap), + ); + record_run( + &mut b, + ExecutionContext::Retry, + "East US", + "https://east/", + CosmosStatus::new(StatusCode::TooManyRequests), + 1.0, + 100, + ); + record_run( + &mut b, + ExecutionContext::Retry, + "East US", + "https://east/", + CosmosStatus::new(StatusCode::Gone).with_sub_status(1002), + 1.0, + 50, + ); + record_run( + &mut b, + ExecutionContext::Retry, + "West US", + "https://west/", + CosmosStatus::new(StatusCode::Ok), + 1.0, + 1, + ); + b.set_operation_status(StatusCode::Ok, None); + let ctx = b.complete(); + + assert_eq!(ctx.request_count(), 151); + assert!(ctx.retained_request_count() <= cap); + let info = ctx.compaction().expect("compacted"); + assert_eq!(info.runs.len(), 3); + assert_eq!(info.runs[0].count, 100); + assert_eq!(info.runs[1].count, 50); + assert_eq!(info.runs[2].count, 1); + + // The 410 run carries its sub-status exactly. + assert_eq!( + info.runs[1].status, + CosmosStatus::new(StatusCode::Gone).with_sub_status(1002) + ); + + // Exact, lossless totals across all three runs. + let run_attempts: usize = info.runs.iter().map(|r| r.count).sum(); + assert_eq!(run_attempts, 151); + assert_eq!(ctx.total_request_charge(), RequestCharge::new(151.0)); + + // Onset (429) and terminal (200) boundaries retained; order preserved. + let requests = ctx.requests(); + assert_eq!( + u16::from(requests.first().unwrap().status().status_code()), + 429 + ); + assert_eq!( + u16::from(requests.last().unwrap().status().status_code()), + 200 + ); + } + + #[test] + fn region_ping_pong_is_bounded_via_global_fallback() { + // Alternating regions: every consecutive run is length one, defeating the + // order-preserving run-length collapse and forcing the order-robust + // global key-bucket fallback. The artifact must still stay bounded. + let cap = 16; + let mut b = DiagnosticsContextBuilder::new( + ActivityId::from_string("pingpong".to_string()), + options_with_cap(cap), + ); + for _ in 0..200 { + let he = b.start_test_request( + ExecutionContext::Retry, + Some(Region::new("East US")), + "https://east/", + ); + b.complete_request(he, StatusCode::ServiceUnavailable, None); + let hw = b.start_test_request( + ExecutionContext::Retry, + Some(Region::new("West US")), + "https://west/", + ); + b.complete_request(hw, StatusCode::ServiceUnavailable, None); + } + b.set_operation_status(StatusCode::ServiceUnavailable, None); + let ctx = b.complete(); + + assert_eq!(ctx.request_count(), 400); + assert!( + ctx.retained_request_count() <= cap, + "retained {} exceeds cap {cap}", + ctx.retained_request_count() + ); + let info = ctx.compaction().expect("ping-pong storm must compact"); + // Two distinct keys -> two runs, each covering 200 attempts. + assert_eq!(info.runs.len(), 2); + assert!(info.runs.iter().all(|r| r.count == 200)); + + // Both regions still reported (normalized: lowercase, no spaces). + let mut regions: Vec = ctx + .regions_contacted() + .iter() + .map(|r| r.as_str().to_string()) + .collect(); + regions.sort(); + assert_eq!(regions, ["eastus".to_string(), "westus".to_string()]); + } + + #[test] + fn distinct_endpoint_410_fanout_is_bounded() { + // A 410 fan-out across thousands of physical-partition endpoints is + // high-cardinality exactly when the storm is worst: every attempt is a + // distinct key. Both the retained records AND the per-run rollup must + // stay bounded by the cap, the omission explicit, the total exact. + let cap = 16; + let distinct = 5000usize; + let mut b = DiagnosticsContextBuilder::new( + ActivityId::from_string("fanout".to_string()), + options_with_cap(cap), + ); + for i in 0..distinct { + record_run( + &mut b, + ExecutionContext::Retry, + "East US", + &format!("https://pkrange-{i}/"), + CosmosStatus::new(StatusCode::Gone).with_sub_status(1002), + 1.0, + 1, + ); + } + b.set_operation_status(StatusCode::Gone, Some(SubStatusCode::new(1002))); + let ctx = b.complete(); + + assert_eq!(ctx.request_count(), distinct); + let info = ctx + .compaction() + .expect("a high-cardinality storm must compact"); + assert_eq!(info.original_request_count, distinct); + assert_eq!(info.total_runs, distinct); + + // Both the retained records and the per-run rollup are bounded by the cap. + assert!( + ctx.retained_request_count() <= cap, + "retained {} exceeds cap {cap}", + ctx.retained_request_count() + ); + assert!( + info.runs.len() <= cap, + "runs {} not bounded by cap {cap}", + info.runs.len() + ); + + // Every drop is explicit, never silent, and lossless in aggregate. + assert!(info.omitted_runs > 0, "run omission must be marked"); + assert_eq!(info.omitted_runs, info.total_runs - info.runs.len()); + let retained_run_attempts: usize = info.runs.iter().map(|r| r.count).sum(); + assert_eq!(retained_run_attempts + info.omitted_request_count, distinct); + + // Detailed JSON size is bounded by the cap, independent of the topology. + let json = ctx.to_json_string(Some(DiagnosticsVerbosity::Detailed)); + assert!(json.contains("\"omitted_runs\""), "omission not surfaced"); + assert!( + json.len() < 32 * 1024, + "detailed json {} bytes grows with topology (distinct={distinct})", + json.len() + ); + } + + #[test] + fn phase2_heterogeneous_runs_keep_largest_and_stay_coherent() { + // Phase 2 with heterogeneous run counts: a few "hot" keys (introduced + // LAST) each retry many times, interleaved with many "cold" single-attempt + // keys (introduced FIRST). The rollup must keep the largest runs, the + // retained records must be drawn from the SAME kept set (so a span emitter + // never sees an attempt whose run was omitted), and totals stay exact. + let cap = 16; + let cold = 40usize; + let hot = 3usize; + let hot_retries = 100usize; + let mut b = DiagnosticsContextBuilder::new( + ActivityId::from_string("heterogeneous".to_string()), + options_with_cap(cap), + ); + + for i in 0..cold { + let h = b.start_test_request( + ExecutionContext::Retry, + Some(Region::new("East US")), + &format!("https://cold-{i}/"), + ); + b.complete_request(h, StatusCode::TooManyRequests, None); + } + // Interleaved so no two consecutive attempts share a key (forces Phase 2). + for _round in 0..hot_retries { + for j in 0..hot { + let h = b.start_test_request( + ExecutionContext::Retry, + Some(Region::new("West US")), + &format!("https://hot-{j}/"), + ); + b.complete_request(h, StatusCode::ServiceUnavailable, None); + } + } + b.set_operation_status(StatusCode::ServiceUnavailable, None); + let ctx = b.complete(); + + let total = cold + hot * hot_retries; + assert_eq!(ctx.request_count(), total); + let info = ctx + .compaction() + .expect("a heterogeneous storm must compact"); + assert_eq!(info.total_runs, cold + hot); + + assert!( + info.runs.len() <= cap, + "runs {} exceed cap {cap}", + info.runs.len() + ); + assert!( + ctx.retained_request_count() <= cap, + "retained {} exceed cap {cap}", + ctx.retained_request_count() + ); + + // The largest (hot) runs are the ones kept in the rollup. + let kept_hot = info.runs.iter().filter(|r| r.count == hot_retries).count(); + assert_eq!( + kept_hot, hot, + "all hot runs must survive the by-count rollup" + ); + + // Coherence: every retained record's key is represented by a run. + let run_keys: std::collections::HashSet<( + Option, + String, + CosmosStatus, + ExecutionContext, + )> = info + .runs + .iter() + .map(|r| { + ( + r.region.clone(), + r.endpoint.clone(), + r.status, + r.execution_context, + ) + }) + .collect(); + for rec in ctx.requests().iter() { + let id = ( + rec.region().map(|r| r.as_str().to_string()), + rec.endpoint().to_string(), + *rec.status(), + rec.execution_context(), + ); + assert!( + run_keys.contains(&id), + "retained record {id:?} has no matching run in the bounded rollup" + ); + } + + // Exact, lossless totals despite the bounded rollup + retained list. + let kept_run_attempts: usize = info.runs.iter().map(|r| r.count).sum(); + assert_eq!(kept_run_attempts + info.omitted_request_count, total); + assert_eq!(info.omitted_runs, info.total_runs - info.runs.len()); + } + + #[test] + fn under_cap_is_not_compacted_and_output_has_no_marker() { + // The default cap (512) is far above a normal operation's attempts, so + // the list is retained verbatim and the output is byte-identical to the + // pre-compaction behavior (no compaction marker). + let ctx = make_context_with(ActivityId::from_string("normal".to_string()), |b| { + for _ in 0..3 { + let h = b.start_test_request( + ExecutionContext::Retry, + Some(Region::new("East US")), + "https://east/", + ); + b.complete_request(h, StatusCode::TooManyRequests, None); + } + }); + + assert!(ctx.compaction().is_none()); + assert_eq!(ctx.request_count(), 3); + assert_eq!(ctx.retained_request_count(), 3); + for verbosity in [ + DiagnosticsVerbosity::Detailed, + DiagnosticsVerbosity::Summary, + ] { + let json = ctx.to_json_string(Some(verbosity)); + assert!( + !json.contains("compaction"), + "{verbosity} output must not carry a compaction marker: {json}" + ); + } + } + + #[test] + fn operation_name_defaults_to_none() { + let ctx = make_context_with(ActivityId::new_uuid(), |_| {}); + assert_eq!(ctx.operation_name(), None); + } + + #[test] + fn is_failure_reflects_operation_status() { + let ok = make_context_with(ActivityId::new_uuid(), |b| { + b.set_operation_status(StatusCode::Ok, None); + }); + assert!(!ok.is_failure()); + + let failed = make_context_with(ActivityId::new_uuid(), |b| { + b.set_operation_status(StatusCode::TooManyRequests, None); + }); + assert!(failed.is_failure()); + + // No recorded status is not a failure. + let no_status = make_context_with(ActivityId::new_uuid(), |_| {}); + assert!(!no_status.is_failure()); + } + + #[test] + fn is_failure_falls_back_to_terminal_request_status() { + // Some driver error-finalization paths graft diagnostics onto the error + // without stamping the operation status. The terminal attempt's failed + // status must still surface as a failure so tail-sampled handlers emit. + let failed = make_context_with(ActivityId::new_uuid(), |b| { + let handle = b.start_test_request( + ExecutionContext::Initial, + Some(Region::WEST_US_2), + "https://test.documents.azure.com", + ); + b.complete_request(handle, StatusCode::TooManyRequests, None); + }); + assert!(failed.is_failure()); + + // A successful terminal attempt with no operation status is not a failure. + let ok = make_context_with(ActivityId::new_uuid(), |b| { + let handle = b.start_test_request( + ExecutionContext::Initial, + Some(Region::WEST_US_2), + "https://test.documents.azure.com", + ); + b.complete_request(handle, StatusCode::Ok, None); + }); + assert!(!ok.is_failure()); + } + + #[test] + fn is_completed_requires_status_or_request() { + let empty = make_context_with(ActivityId::new_uuid(), |_| {}); + assert!(!empty.is_completed()); + + let with_status = make_context_with(ActivityId::new_uuid(), |b| { + b.set_operation_status(StatusCode::Ok, None); + }); + assert!(with_status.is_completed()); + + let with_request = make_context_with(ActivityId::new_uuid(), |b| { + let handle = b.start_test_request( + ExecutionContext::Initial, + Some(Region::WEST_US_2), + "https://test.documents.azure.com", + ); + b.complete_request(handle, StatusCode::Ok, None); + }); + assert!(with_request.is_completed()); + } + + #[test] + fn is_threshold_violated_on_request_charge() { + let thresholds = DiagnosticsThresholds::default().with_request_charge(100.0); + + let cheap = make_context_with(ActivityId::new_uuid(), |b| { + let handle = b.start_test_request( + ExecutionContext::Initial, + Some(Region::WEST_US_2), + "https://test.documents.azure.com", + ); + b.update_request(handle, |req| req.request_charge = RequestCharge::new(10.0)); + b.complete_request(handle, StatusCode::Ok, None); + }); + assert!(!cheap.is_threshold_violated(&thresholds)); + + let expensive = make_context_with(ActivityId::new_uuid(), |b| { + let handle = b.start_test_request( + ExecutionContext::Initial, + Some(Region::WEST_US_2), + "https://test.documents.azure.com", + ); + b.update_request(handle, |req| req.request_charge = RequestCharge::new(150.0)); + b.complete_request(handle, StatusCode::Ok, None); + }); + assert!(expensive.is_threshold_violated(&thresholds)); + } + + #[test] + fn threshold_breach_reports_the_specific_bound() { + let thresholds = DiagnosticsThresholds::default().with_request_charge(100.0); + + // Cheap, fast success: no breach. + let cheap = make_context_with(ActivityId::new_uuid(), |b| { + let handle = b.start_test_request( + ExecutionContext::Initial, + Some(Region::WEST_US_2), + "https://test.documents.azure.com", + ); + b.update_request(handle, |req| req.request_charge = RequestCharge::new(10.0)); + b.complete_request(handle, StatusCode::Ok, None); + }); + assert_eq!( + cheap.threshold_breach_for(&thresholds, Some("read_item")), + None + ); + + // Over the RU bound: the breach names the request charge. + let expensive = make_context_with(ActivityId::new_uuid(), |b| { + let handle = b.start_test_request( + ExecutionContext::Initial, + Some(Region::WEST_US_2), + "https://test.documents.azure.com", + ); + b.update_request(handle, |req| req.request_charge = RequestCharge::new(150.0)); + b.complete_request(handle, StatusCode::Ok, None); + }); + assert_eq!( + expensive.threshold_breach_for(&thresholds, Some("read_item")), + Some(ThresholdBreach::RequestCharge) + ); + } } diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/mod.rs b/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/mod.rs index f8b62e65700..8581941402d 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/mod.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/diagnostics/mod.rs @@ -18,14 +18,17 @@ //! When an operation completes, the builder is consumed to create an immutable //! `DiagnosticsContext` which is safe to share via `Arc` without locking. +mod compaction; mod diagnostics_context; mod proxy_configuration; +pub use compaction::{CompactedRun, CompactionInfo}; pub(crate) use diagnostics_context::DiagnosticsContextBuilder; pub use diagnostics_context::{ DiagnosticsContext, ExecutionContext, FailedTransportShardDiagnostics, PipelineType, RequestDiagnostics, RequestEvent, RequestEventType, RequestHandle, RequestSentStatus, - TransportHttpVersion, TransportKind, TransportSecurity, TransportShardDiagnostics, + ThresholdBreach, TransportHttpVersion, TransportKind, TransportSecurity, + TransportShardDiagnostics, }; pub use proxy_configuration::ProxyConfiguration; diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/lib.rs b/sdk/cosmos/azure_data_cosmos_driver/src/lib.rs index cd94c0ef497..34d7313bca8 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/lib.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/lib.rs @@ -64,4 +64,4 @@ pub use diagnostics::{DiagnosticsContext, ExecutionContext, RequestDiagnostics, pub use driver::{CosmosDriver, CosmosDriverRuntime, CosmosDriverRuntimeBuilder, OperationPlan}; pub use error::{CosmosError, CosmosErrorBuilder, CosmosStatus, Result, SubStatusCode}; pub use models::{ActivityId, CosmosResponse, RequestCharge, ResponseBody}; -pub use options::{DiagnosticsOptions, DiagnosticsVerbosity, DriverOptions}; +pub use options::{DiagnosticsOptions, DiagnosticsThresholds, DiagnosticsVerbosity, DriverOptions}; diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/options/diagnostics_options.rs b/sdk/cosmos/azure_data_cosmos_driver/src/options/diagnostics_options.rs index 20ebf7557f8..3a0e872cd5f 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/options/diagnostics_options.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/options/diagnostics_options.rs @@ -13,6 +13,16 @@ const DEFAULT_MAX_SUMMARY_SIZE_BYTES: usize = 8 * 1024; /// Minimum allowed size for diagnostic summary output (4 KB). const MIN_MAX_SUMMARY_SIZE_BYTES: usize = 4 * 1024; +/// Default maximum number of per-attempt `RequestDiagnostics` records retained +/// on a finalized `DiagnosticsContext` before the list is compacted. +const DEFAULT_MAX_REQUEST_DIAGNOSTICS: usize = 512; + +/// Minimum allowed value for `max_request_diagnostics`. +/// +/// A cap this low still leaves room for the first and last record of several +/// distinct runs while keeping a retry storm's artifact bounded. +const MIN_MAX_REQUEST_DIAGNOSTICS: usize = 16; + /// Controls the verbosity level of diagnostic output. /// /// Diagnostics can be output in different levels of detail depending on @@ -95,6 +105,16 @@ pub struct DiagnosticsOptions { /// will be truncated to fit within this limit. Default is 8 KB. pub(crate) max_summary_size_bytes: usize, + /// Maximum number of per-attempt `RequestDiagnostics` records retained on a + /// finalized `DiagnosticsContext`. + /// + /// When an operation makes more attempts than this (a retry storm), the + /// finalized per-attempt list is compacted — runs of near-identical + /// consecutive retries are collapsed to their first and last record plus an + /// exact per-run rollup (count, request charge, min/max/P50 duration), and + /// the compaction is marked explicitly. Default is 512, minimum 16. + pub(crate) max_request_diagnostics: usize, + /// Default verbosity level when none is specified. /// /// Used when `to_json_string` is called with `None` for verbosity. @@ -120,6 +140,12 @@ impl DiagnosticsOptions { self.max_summary_size_bytes } + /// Returns the maximum number of per-attempt `RequestDiagnostics` records + /// retained before the finalized list is compacted under a retry storm. + pub fn max_request_diagnostics(&self) -> usize { + self.max_request_diagnostics + } + /// Returns the default verbosity level. pub fn default_verbosity(&self) -> DiagnosticsVerbosity { self.default_verbosity @@ -142,6 +168,9 @@ impl DiagnosticsOptions { /// /// - `AZURE_COSMOS_DIAGNOSTICS_MAX_SUMMARY_SIZE_BYTES`: Maximum size in bytes for /// summary mode output (default: `8192`, min: `4096`) +/// - `AZURE_COSMOS_DIAGNOSTICS_MAX_REQUEST_DIAGNOSTICS`: Maximum number of per-attempt +/// request records retained before the finalized list is compacted under a retry +/// storm (default: `512`, min: `16`) /// - `AZURE_COSMOS_DIAGNOSTICS_DEFAULT_VERBOSITY`: Default verbosity level. /// Valid values: `default`, `summary`, `detailed` (default: `summary`) /// @@ -162,6 +191,8 @@ impl DiagnosticsOptions { pub struct DiagnosticsOptionsBuilder { #[option(env = "AZURE_COSMOS_DIAGNOSTICS_MAX_SUMMARY_SIZE_BYTES")] max_summary_size_bytes: Option, + #[option(env = "AZURE_COSMOS_DIAGNOSTICS_MAX_REQUEST_DIAGNOSTICS")] + max_request_diagnostics: Option, #[option(env = "AZURE_COSMOS_DIAGNOSTICS_DEFAULT_VERBOSITY")] default_verbosity: Option, } @@ -181,6 +212,15 @@ impl DiagnosticsOptionsBuilder { self } + /// Sets the maximum number of per-attempt request records retained before + /// the finalized diagnostics list is compacted under a retry storm. + /// + /// Must be at least 16. Default: 512. + pub fn with_max_request_diagnostics(mut self, max: usize) -> Self { + self.max_request_diagnostics = Some(max); + self + } + /// Sets the default verbosity level. /// /// Default: `DiagnosticsVerbosity::Summary`. @@ -197,6 +237,7 @@ impl DiagnosticsOptionsBuilder { /// /// Returns an error if: /// - `max_summary_size_bytes` is less than 4096 + /// - `max_request_diagnostics` is less than 16 pub fn build(self) -> crate::error::Result { // The builder doubles as the env source: `Self::from_env()` reads each // `#[option(env = ...)]` field from its `AZURE_COSMOS_*` variable @@ -213,6 +254,14 @@ impl DiagnosticsOptionsBuilder { ValidationBounds::min(MIN_MAX_SUMMARY_SIZE_BYTES), )?; + let max_request_diagnostics = resolve_from_env( + self.max_request_diagnostics, + env.max_request_diagnostics, + "AZURE_COSMOS_DIAGNOSTICS_MAX_REQUEST_DIAGNOSTICS", + DEFAULT_MAX_REQUEST_DIAGNOSTICS, + ValidationBounds::min(MIN_MAX_REQUEST_DIAGNOSTICS), + )?; + let default_verbosity = match self.default_verbosity.or(env.default_verbosity) { Some(DiagnosticsVerbosity::Default) | None => DiagnosticsVerbosity::Summary, Some(verbosity) => verbosity, @@ -220,6 +269,7 @@ impl DiagnosticsOptionsBuilder { Ok(DiagnosticsOptions { max_summary_size_bytes, + max_request_diagnostics, default_verbosity, }) } @@ -233,6 +283,7 @@ mod tests { fn defaults() { let options = DiagnosticsOptions::default(); assert_eq!(options.max_summary_size_bytes, 8 * 1024); + assert_eq!(options.max_request_diagnostics, 512); assert_eq!(options.default_verbosity, DiagnosticsVerbosity::Summary); } @@ -243,11 +294,13 @@ mod tests { // each `#[option(env = ...)]` string to the right field. let cfg = DiagnosticsOptionsBuilder::from_env_vars(|key| match key { "AZURE_COSMOS_DIAGNOSTICS_MAX_SUMMARY_SIZE_BYTES" => Ok("16384".to_string()), + "AZURE_COSMOS_DIAGNOSTICS_MAX_REQUEST_DIAGNOSTICS" => Ok("128".to_string()), "AZURE_COSMOS_DIAGNOSTICS_DEFAULT_VERBOSITY" => Ok("summary".to_string()), _ => Err(std::env::VarError::NotPresent), }); assert_eq!(cfg.max_summary_size_bytes, Some(16384)); + assert_eq!(cfg.max_request_diagnostics, Some(128)); assert_eq!(cfg.default_verbosity, Some(DiagnosticsVerbosity::Summary)); } @@ -257,10 +310,12 @@ mod tests { // to the default instead of erroring. let cfg = DiagnosticsOptionsBuilder::from_env_vars(|key| match key { "AZURE_COSMOS_DIAGNOSTICS_MAX_SUMMARY_SIZE_BYTES" => Ok("huge".to_string()), + "AZURE_COSMOS_DIAGNOSTICS_MAX_REQUEST_DIAGNOSTICS" => Ok("lots".to_string()), "AZURE_COSMOS_DIAGNOSTICS_DEFAULT_VERBOSITY" => Ok("nonsense".to_string()), _ => Err(std::env::VarError::NotPresent), }); assert!(cfg.max_summary_size_bytes.is_none()); + assert!(cfg.max_request_diagnostics.is_none()); assert!(cfg.default_verbosity.is_none()); } @@ -268,11 +323,13 @@ mod tests { fn custom_values() { let options = DiagnosticsOptionsBuilder::new() .with_max_summary_size_bytes(16 * 1024) + .with_max_request_diagnostics(64) .with_default_verbosity(DiagnosticsVerbosity::Detailed) .build() .unwrap(); assert_eq!(options.max_summary_size_bytes, 16 * 1024); + assert_eq!(options.max_request_diagnostics, 64); assert_eq!(options.default_verbosity, DiagnosticsVerbosity::Detailed); } @@ -299,6 +356,28 @@ mod tests { .contains("must be at least 4096")); } + #[test] + fn max_request_diagnostics_too_small() { + let result = DiagnosticsOptionsBuilder::new() + .with_max_request_diagnostics(8) // below minimum of 16 + .build(); + + assert!(result.is_err()); + assert!(result + .unwrap_err() + .to_string() + .contains("must be at least 16")); + } + + #[test] + fn max_request_diagnostics_at_minimum_is_accepted() { + let options = DiagnosticsOptionsBuilder::new() + .with_max_request_diagnostics(16) + .build() + .unwrap(); + assert_eq!(options.max_request_diagnostics, 16); + } + #[test] fn verbosity_from_str() { assert_eq!( diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/options/diagnostics_thresholds.rs b/sdk/cosmos/azure_data_cosmos_driver/src/options/diagnostics_thresholds.rs new file mode 100644 index 00000000000..df8b21a12fa --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos_driver/src/options/diagnostics_thresholds.rs @@ -0,0 +1,237 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Tail-based sampling thresholds for diagnostics emission. +//! +//! [`DiagnosticsThresholds`] captures the per-operation limits that the SDK's +//! emission handlers (tracing, sampled logging) use to decide *whether* a +//! completed operation is "interesting" enough to emit. A completed operation +//! is interesting when it fails or when it crosses one of these thresholds — +//! see [`DiagnosticsContext::is_threshold_violated`](crate::diagnostics::DiagnosticsContext::is_threshold_violated). +//! +//! The defaults mirror the Java SDK's `CosmosDiagnosticsThresholds` so the two +//! stacks behave consistently out of the box, and every threshold is +//! configurable through the usual builder-style options chain. + +use std::time::Duration; + +/// Default point-operation latency threshold (1 second), matching Java. +const DEFAULT_POINT_OPERATION_LATENCY: Duration = Duration::from_secs(1); + +/// Default non-point-operation latency threshold (3 seconds), matching Java. +const DEFAULT_NON_POINT_OPERATION_LATENCY: Duration = Duration::from_secs(3); + +/// Default request-charge threshold in Request Units (1000 RU), matching Java. +const DEFAULT_REQUEST_CHARGE_RU: f64 = 1000.0; + +/// Default payload-size threshold in bytes (1 MiB), matching Java. +const DEFAULT_PAYLOAD_SIZE_BYTES: u64 = 1024 * 1024; + +/// Per-operation thresholds that gate diagnostics emission (tail-based sampling). +/// +/// An operation is considered "interesting" — and therefore eligible for a span +/// or a sampled log line — when it *fails* or when it crosses one of these +/// thresholds. Latency is checked against +/// [`point_operation_latency`](Self::point_operation_latency) for point +/// operations (single-item reads/writes) and +/// [`non_point_operation_latency`](Self::non_point_operation_latency) for +/// everything else (queries, batches, bulk, …). +/// +/// # Defaults +/// +/// The defaults match the Java SDK: +/// +/// | Threshold | Default | +/// | --------- | ------- | +/// | Point-operation latency | 1 s | +/// | Non-point-operation latency | 3 s | +/// | Request charge | 1000 RU | +/// | Payload size | 1 MiB | +/// +/// # Examples +/// +/// ``` +/// use std::time::Duration; +/// use azure_data_cosmos_driver::options::DiagnosticsThresholds; +/// +/// // Java-like defaults. +/// let thresholds = DiagnosticsThresholds::default(); +/// assert_eq!(thresholds.point_operation_latency(), Duration::from_secs(1)); +/// +/// // Tune via the options chain. +/// let strict = DiagnosticsThresholds::default() +/// .with_point_operation_latency(Duration::from_millis(250)) +/// .with_request_charge(50.0); +/// assert_eq!(strict.point_operation_latency(), Duration::from_millis(250)); +/// assert_eq!(strict.request_charge(), 50.0); +/// ``` +#[derive(Clone, Copy, Debug, PartialEq)] +#[non_exhaustive] +pub struct DiagnosticsThresholds { + point_operation_latency: Duration, + non_point_operation_latency: Duration, + request_charge: f64, + payload_size: u64, +} + +impl DiagnosticsThresholds { + /// Creates thresholds with the Java-like defaults. + /// + /// Equivalent to [`DiagnosticsThresholds::default`]. + pub fn new() -> Self { + Self::default() + } + + /// Sets the latency threshold for point operations (single-item CRUD). + /// + /// Operations slower than this are eligible for emission. Default: 1 second. + #[must_use] + pub fn with_point_operation_latency(mut self, latency: Duration) -> Self { + self.point_operation_latency = latency; + self + } + + /// Sets the latency threshold for non-point operations (queries, batches, …). + /// + /// Operations slower than this are eligible for emission. Default: 3 seconds. + #[must_use] + pub fn with_non_point_operation_latency(mut self, latency: Duration) -> Self { + self.non_point_operation_latency = latency; + self + } + + /// Sets the total request-charge threshold, in Request Units (RU). + /// + /// Operations charging more than this are eligible for emission. + /// Default: 1000 RU. + /// + /// Non-finite (`NaN`/±∞) or negative values are **ignored** and leave the + /// current threshold unchanged. Such values would otherwise corrupt + /// RU-based sampling: `NaN` makes every `charge > threshold` comparison + /// false (silently disabling RU sampling), and a negative bound makes it + /// always true (sampling every operation). + #[must_use] + pub fn with_request_charge(mut self, request_charge: f64) -> Self { + if request_charge.is_finite() && request_charge >= 0.0 { + self.request_charge = request_charge; + } + self + } + + /// Sets the payload-size threshold, in bytes. + /// + /// Default: 1 MiB. + /// + /// **Note:** this threshold is retained for parity with the Java SDK and + /// forward compatibility, but is not evaluated yet because + /// [`DiagnosticsContext`](crate::diagnostics::DiagnosticsContext) does not + /// currently carry request/response body sizes. + #[must_use] + pub fn with_payload_size(mut self, payload_size: u64) -> Self { + self.payload_size = payload_size; + self + } + + /// Returns the point-operation latency threshold. + pub fn point_operation_latency(&self) -> Duration { + self.point_operation_latency + } + + /// Returns the non-point-operation latency threshold. + pub fn non_point_operation_latency(&self) -> Duration { + self.non_point_operation_latency + } + + /// Returns the request-charge threshold, in Request Units (RU). + pub fn request_charge(&self) -> f64 { + self.request_charge + } + + /// Returns the payload-size threshold, in bytes. + pub fn payload_size(&self) -> u64 { + self.payload_size + } +} + +impl Default for DiagnosticsThresholds { + fn default() -> Self { + Self { + point_operation_latency: DEFAULT_POINT_OPERATION_LATENCY, + non_point_operation_latency: DEFAULT_NON_POINT_OPERATION_LATENCY, + request_charge: DEFAULT_REQUEST_CHARGE_RU, + payload_size: DEFAULT_PAYLOAD_SIZE_BYTES, + } + } +} + +/// Returns `true` when `operation_name` denotes a single-item ("point") +/// operation, which is latency-gated by +/// [`DiagnosticsThresholds::point_operation_latency`]. +/// +/// Anything not recognized as a point operation (queries, batches, bulk, +/// change feed, …) is treated as a non-point operation. Names are the +/// canonical semantic-convention `db.operation.name` values. +pub(crate) fn is_point_operation(operation_name: &str) -> bool { + matches!( + operation_name, + "read_item" | "create_item" | "upsert_item" | "replace_item" | "delete_item" | "patch_item" + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn defaults_match_java() { + let t = DiagnosticsThresholds::default(); + assert_eq!(t.point_operation_latency(), Duration::from_secs(1)); + assert_eq!(t.non_point_operation_latency(), Duration::from_secs(3)); + assert_eq!(t.request_charge(), 1000.0); + assert_eq!(t.payload_size(), 1024 * 1024); + } + + #[test] + fn builder_overrides_apply() { + let t = DiagnosticsThresholds::new() + .with_point_operation_latency(Duration::from_millis(100)) + .with_non_point_operation_latency(Duration::from_millis(500)) + .with_request_charge(10.0) + .with_payload_size(4096); + assert_eq!(t.point_operation_latency(), Duration::from_millis(100)); + assert_eq!(t.non_point_operation_latency(), Duration::from_millis(500)); + assert_eq!(t.request_charge(), 10.0); + assert_eq!(t.payload_size(), 4096); + } + + #[test] + fn point_operation_classification() { + assert!(is_point_operation("read_item")); + assert!(is_point_operation("create_item")); + assert!(is_point_operation("patch_item")); + assert!(!is_point_operation("query_items")); + assert!(!is_point_operation("execute_batch")); + assert!(!is_point_operation("unknown")); + } + + #[test] + fn request_charge_ignores_non_finite_and_negative() { + // A valid override applies. + let base = DiagnosticsThresholds::default().with_request_charge(250.0); + assert_eq!(base.request_charge(), 250.0); + + // NaN would make `charge > threshold` false for every operation, + // silently disabling RU sampling; it must be ignored and leave the + // prior value intact. + assert_eq!(base.with_request_charge(f64::NAN).request_charge(), 250.0); + // A negative bound would sample every operation; ignored. + assert_eq!(base.with_request_charge(-1.0).request_charge(), 250.0); + // Infinities are non-finite; ignored. + assert_eq!( + base.with_request_charge(f64::INFINITY).request_charge(), + 250.0 + ); + // Zero is a valid (finite, non-negative) threshold. + assert_eq!(base.with_request_charge(0.0).request_charge(), 0.0); + } +} diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/options/mod.rs b/sdk/cosmos/azure_data_cosmos_driver/src/options/mod.rs index 4c12ea43b45..418a146cff1 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/options/mod.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/options/mod.rs @@ -18,6 +18,7 @@ mod availability_strategy; mod connection_pool; mod diagnostics_options; +mod diagnostics_thresholds; mod driver_options; pub(crate) mod env_parsing; mod identity; @@ -34,6 +35,8 @@ pub use connection_pool::{ConnectionPoolOptions, ConnectionPoolOptionsBuilder}; pub use diagnostics_options::{ DiagnosticsOptions, DiagnosticsOptionsBuilder, DiagnosticsVerbosity, }; +pub(crate) use diagnostics_thresholds::is_point_operation; +pub use diagnostics_thresholds::DiagnosticsThresholds; pub use driver_options::{DriverOptions, DriverOptionsBuilder}; pub(crate) use env_parsing::parse_duration_millis_from_env; pub use identity::{CorrelationId, UserAgentSuffix, WorkloadId};