Add Cosmos client-side observability layer

sdk/cosmos/.cspell.json

@@ -220,6 +220,7 @@
     "pgcosmos",
     "PIDS",
     "Pigw",
+    "pingpong",
     "pinning",
     "Pkrange",
     "pkranges",
@@ -256,12 +257,14 @@
     "rhost",
     "RNTBD",
     "rgby",
+    "rollup",
     "roundtripped",
     "roundtrips",
     "RTLD",
     "RUPM",
     "rwcache",
     "sbyte",
+    "semconv",
     "serviceunavailable",
     "serviceversion",
     "sess",

sdk/cosmos/azure_data_cosmos/CHANGELOG.md

@@ -4,6 +4,9 @@
 
 ### Features Added
 
+- Added a pluggable diagnostics emission layer on top of the always-collected `DiagnosticsContext`: the `DiagnosticsHandler` trait and the ordered `DiagnosticsHandlerChain`, registered via `CosmosClientBuilder::with_diagnostics_handler`. Each registered handler observes every operation's completed `DiagnosticsContext` exactly once at completion; with no handler registered the completion path is a zero-overhead no-op. The SDK attaches a `CosmosOperationContext` (operation name, database, and container identity the driver context does not carry) to the pipeline `Context` passed to the chain, so handlers can emit correct operation-scope telemetry.
+- Added an OpenTelemetry metrics handler, `CosmosMetricsHandler` (with `MetricsOptions`), behind the off-by-default `otel_metrics` feature. It records the stable `db.client.operation.duration` histogram from each completed operation, and — opt-in — development-tier metrics (`azure.cosmosdb.client.operation.request_charge`, `db.client.response.returned_rows`, `azure.cosmosdb.client.active_instance.count`). With the feature disabled the metrics code is not compiled; with it enabled but no meter provider registered, recording is a no-op.
+- Added tail-based sampling diagnostics emission handlers, registered via `CosmosClientBuilder::with_diagnostics_handler`: `SamplingLogHandler` logs a compact, rate-limited diagnostics line (via `tracing`) and, behind the off-by-default `otel_tracing` feature, `CosmosTracingHandler` emits a backdated OpenTelemetry span tree (operation span plus one child per retained attempt). Both emit only for operations that fail or breach a configurable `DiagnosticsThresholds` (re-exported from `azure_data_cosmos::diagnostics`); a fast, successful operation emits nothing. The log handler caps emissions per window (~100/min by default), always admits a bounded number of failures, and emits exactly one "suppressed N" notice per window.
 - Added runtime diagnostics output configuration via `CosmosRuntimeBuilder::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))
 - Gateway 2.0 transport (a regional proxy forwarding RNTBD-over-HTTP/2) is selected 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))
 - Added the re-exported `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))

sdk/cosmos/azure_data_cosmos/Cargo.toml

@@ -20,6 +20,7 @@ azure_core = { workspace = true, default-features = false }
 azure_data_cosmos_driver = { version = "0.6.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)
+otel_tracing = [
+  "dep:opentelemetry",
+] # Enables the OpenTelemetry 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
+otel_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.
@@ -100,6 +111,8 @@ features = [
   "hmac_rust",
   "key_auth",
   "native_tls",
+  "otel_metrics",
+  "otel_tracing",
   "rustls",
 ]
 

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, 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
@@ -92,7 +103,8 @@ impl ContainerClient {
             .await?;
 
         Ok(ResourceResponse::new(
-            crate::driver_bridge::driver_response_to_cosmos_response(driver_response),
+            self.context
+                .complete_operation(driver_response, || self.operation_context("read_container")),
         ))
     }
 
@@ -145,7 +157,9 @@ impl ContainerClient {
             .await?;
 
         Ok(ResourceResponse::new(
-            crate::driver_bridge::driver_response_to_cosmos_response(driver_response),
+            self.context.complete_operation(driver_response, || {
+                self.operation_context("replace_container")
+            }),
         ))
     }
 
@@ -231,7 +245,9 @@ impl ContainerClient {
             .await?;
 
         Ok(ResourceResponse::new(
-            crate::driver_bridge::driver_response_to_cosmos_response(driver_response),
+            self.context.complete_operation(driver_response, || {
+                self.operation_context("delete_container")
+            }),
         ))
     }
 
@@ -330,7 +346,8 @@ impl ContainerClient {
 
         // 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_operation(driver_response, || self.operation_context("create_item")),
         ))
     }
 
@@ -428,7 +445,8 @@ impl ContainerClient {
 
         // 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_operation(driver_response, || self.operation_context("replace_item")),
         ))
     }
 
@@ -534,7 +552,8 @@ impl ContainerClient {
             .await?;
 
         Ok(ItemResponse::new(
-            crate::driver_bridge::driver_response_to_cosmos_response(driver_response),
+            self.context
+                .complete_operation(driver_response, || self.operation_context("patch_item")),
         ))
     }
 
@@ -636,7 +655,8 @@ impl ContainerClient {
 
         // 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_operation(driver_response, || self.operation_context("upsert_item")),
         ))
     }
 
@@ -696,7 +716,8 @@ impl ContainerClient {
 
         // 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_operation(driver_response, || self.operation_context("read_item")),
         ))
     }
 
@@ -748,7 +769,8 @@ impl ContainerClient {
 
         // 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_operation(driver_response, || self.operation_context("delete_item")),
         ))
     }
 
@@ -1007,7 +1029,9 @@ impl ContainerClient {
             .await?;
 
         Ok(BatchResponse::new(
-            crate::driver_bridge::driver_response_to_cosmos_response(driver_response),
+            self.context.complete_operation(driver_response, || {
+                self.operation_context("execute_transactional_batch")
+            }),
         ))
     }
 

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,
@@ -227,9 +228,14 @@ impl CosmosClient {
             .execute_singleton_operation(operation, operation_options)
             .await?;
 
-        Ok(ResourceResponse::new(
-            crate::driver_bridge::driver_response_to_cosmos_response(driver_response),
-        ))
+        Ok(ResourceResponse::new(self.context.complete_operation(
+            driver_response,
+            || {
+                CosmosOperationContext::new()
+                    .with_operation_name("create_database")
+                    .with_database_name(id.to_string())
+            },
+        )))
     }
 }
 

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 is a
+    /// zero-overhead no-op.
+    ///
+    /// # Arguments
+    ///
+    /// * `handler` - The handler to append to this client's diagnostics chain.
+    pub fn with_diagnostics_handler(mut self, handler: Arc<dyn DiagnosticsHandler>) -> 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,
+            },
         })
     }
 }

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
@@ -90,7 +100,8 @@ impl DatabaseClient {
             .await?;
 
         Ok(ResourceResponse::new(
-            crate::driver_bridge::driver_response_to_cosmos_response(driver_response),
+            self.context
+                .complete_operation(driver_response, || self.operation_context("read_database")),
         ))
     }
 
@@ -180,7 +191,9 @@ impl DatabaseClient {
             .await?;
 
         Ok(ResourceResponse::new(
-            crate::driver_bridge::driver_response_to_cosmos_response(driver_response),
+            self.context.complete_operation(driver_response, || {
+                self.operation_context("create_container")
+            }),
         ))
     }
 
@@ -204,7 +217,9 @@ impl DatabaseClient {
             .await?;
 
         Ok(ResourceResponse::new(
-            crate::driver_bridge::driver_response_to_cosmos_response(driver_response),
+            self.context.complete_operation(driver_response, || {
+                self.operation_context("delete_database")
+            }),
         ))
     }