Cosmos: Microbenchmark and some optimizations

Cargo.toml

@@ -13,6 +13,7 @@ members = [
   "sdk/core/azure_core_test_macros",
   "sdk/core/azure_core_opentelemetry",
   "sdk/cosmos/azure_data_cosmos",
+  "sdk/cosmos/azure_data_cosmos_benchmarks",
   "sdk/cosmos/azure_data_cosmos_driver",
   "sdk/cosmos/azure_data_cosmos_macros",
   "sdk/cosmos/azure_data_cosmos_perf",
@@ -175,6 +176,13 @@ uuid = { version = "1.20", features = ["v4"] }
 windows = { version = "0.62", default-features = false }
 zip = { version = "8.2", default-features = false, features = ["deflate"] }
 
+# Profile for running benchmarks
+# Inherits release optimizations but keeps debug symbols so that tools
+# can produce annotated output.
+[profile.bench]
+inherits = "release"
+debug = 1
+
 [workspace.lints.clippy]
 large_futures = "deny"
 uninlined_format_args = "allow"

sdk/core/typespec/src/http/headers.rs

@@ -134,6 +134,11 @@ impl Headers {
         Self::default()
     }
 
+    /// Create a new headers collection with at least the specified capacity.
+    pub fn with_capacity(n: usize) -> Self {
+        Self(std::collections::HashMap::with_capacity(n))
+    }
+
     /// Gets the headers represented by `H`, or return an error if the header is not found.
     pub fn get<H: FromHeaders>(&self) -> crate::Result<H> {
         match H::from_headers(self) {

sdk/cosmos/.cspell.json

@@ -4,6 +4,7 @@
   ],
   "ignoreWords": [
     "accountname",
+    "ALPN",
     "apacsoutheast",
     "Appleby's",
     "australiacentral",
@@ -42,6 +43,7 @@
     "failback",
     "failovers",
     "FILETIME",
+    "flamegraph",
     "fract",
     "francecentral",
     "francesouth",
@@ -94,6 +96,7 @@
     "PPCB",
     "pushback",
     "qname",
+    "RAII",
     "readfeed",
     "replicaset",
     "reqs",

sdk/cosmos/azure_data_cosmos/src/driver_bridge.rs

@@ -24,18 +24,20 @@ use crate::{
 /// Converts a driver [`DriverResponse`] into the SDK's typed [`CosmosResponse<T>`].
 ///
 /// This reconstructs an `azure_core::Response<T>` from the driver's raw bytes,
-/// status code, and headers, then wraps it in the SDK's response type.
+/// status code, and headers, then wraps it in the SDK's response type using
+/// the pre-parsed headers from the driver to avoid a redundant parse.
 pub(crate) fn driver_response_to_cosmos_response<T>(
     driver_response: DriverResponse,
 ) -> CosmosResponse<T> {
     let status_code: StatusCode = driver_response.status().status_code();
-    let headers = driver_response_headers_to_headers(driver_response.headers());
+    let cosmos_headers = driver_response.headers().clone();
+    let headers = driver_response_headers_to_headers(&cosmos_headers);
     let body = driver_response.into_body();
 
     let raw_response = RawResponse::from_bytes(status_code, headers, Bytes::from(body));
     let typed_response: Response<T> = raw_response.into();
 
-    CosmosResponse::from_response(typed_response)
+    CosmosResponse::from_driver_response(typed_response, cosmos_headers)
 }
 
 /// Converts driver [`CosmosResponseHeaders`] into raw [`Headers`] for the SDK response.

sdk/cosmos/azure_data_cosmos/src/fault_injection/http_client.rs

@@ -1,6 +1,8 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+// cSpell:ignore evals
+
 use super::result::FaultInjectionResult;
 use super::rule::FaultInjectionRule;
 use super::FaultInjectionErrorType;

sdk/cosmos/azure_data_cosmos/src/models/cosmos_response.rs

@@ -50,11 +50,14 @@ impl<T> CosmosResponse<T> {
         }
     }
 
-    /// Creates a `CosmosResponse` from a typed response without a request.
+    /// Creates a `CosmosResponse` from a typed response and pre-parsed headers.
     ///
-    /// Used for driver-routed operations where no `CosmosRequest` is available.
-    pub(crate) fn from_response(response: Response<T>) -> Self {
-        let cosmos_headers = CosmosResponseHeaders::from_headers(response.headers());
+    /// Used by the driver bridge to avoid re-parsing headers that were already
+    /// parsed by the driver pipeline.
+    pub(crate) fn from_driver_response(
+        response: Response<T>,
+        cosmos_headers: CosmosResponseHeaders,
+    ) -> Self {
         let diagnostics = CosmosDiagnostics::from_headers(&cosmos_headers);
         Self {
             response,

sdk/cosmos/azure_data_cosmos_benchmarks/Cargo.toml

@@ -0,0 +1,29 @@
+[package]
+name = "azure_data_cosmos_benchmarks"
+version = "0.1.0"
+description = "Benchmarks for the Azure Cosmos DB Rust driver"
+publish = false
+authors.workspace = true
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+rust-version.workspace = true
+
+[[bench]]
+name = "point_read"
+harness = false
+
+[dependencies]
+async-trait.workspace = true
+azure_core = { workspace = true }
+azure_data_cosmos_driver = { path = "../azure_data_cosmos_driver", features = [
+  "__internal_mocking",
+] }
+tokio = { workspace = true, features = ["rt-multi-thread", "time"] }
+url.workspace = true
+
+[dev-dependencies]
+criterion.workspace = true
+
+[lints]
+workspace = true

sdk/cosmos/azure_data_cosmos_benchmarks/README.md

@@ -0,0 +1,35 @@
+# azure_data_cosmos_benchmarks
+
+Criterion benchmarks for the Azure Cosmos DB Rust driver. All benchmarks
+replace the reqwest transport with an in-memory mock, so they measure driver
+overhead only — routing, signing, retry state, response parsing, and session
+token management — with no network I/O.
+
+## Running the benchmarks
+
+### Standard latency benchmark
+
+```sh
+cargo bench -p azure_data_cosmos_benchmarks --bench point_read
+```
+
+Results are written to `target/criterion/point_read/`.
+
+### CPU flamegraph (pprof)
+
+Pass `--profile-time <seconds>` to enable pprof sampling:
+
+```sh
+cargo bench -p azure_data_cosmos_benchmarks --bench point_read -- --profile-time 30
+```
+
+The flamegraph SVG is written to
+`target/criterion/point_read/profile/point_read.svg`.
+
+For readable symbol names, build with debug symbols:
+
+```sh
+cargo bench -p azure_data_cosmos_benchmarks --profile bench --bench point_read -- --profile-time 30
+```
+
+

sdk/cosmos/azure_data_cosmos_benchmarks/benches/point_read.rs

@@ -0,0 +1,82 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+//! Criterion benchmark for `CosmosDriver::execute_operation` — point read.
+//!
+//! By default the reqwest transport is replaced with an in-memory mock so that
+//! the benchmark measures driver overhead (routing, signing, retry state,
+//! response parsing, session token management) without any network I/O.
+//!
+//! Set `AZURE_BENCH_MODE=live` to run against a real Cosmos DB endpoint. See
+//! `azure_data_cosmos_benchmarks::setup_live` for the required environment
+//! variables.
+//!
+//! Cache priming (account metadata, container metadata) is performed in setup,
+//! outside the measured iteration loop.
+//!
+//! # CPU flamegraph profiling
+//!
+//! Run with `--profile-time` to generate a flamegraph SVG via pprof:
+//!
+//! ```text
+//! cargo bench -p azure_data_cosmos_benchmarks --bench point_read -- --profile-time 30
+//! ```
+//!
+//! Output: `target/criterion/point_read/profile/flamegraph.svg`
+
+use azure_data_cosmos_benchmarks::{self as common, BenchConfig};
+
+use std::time::Duration;
+
+use azure_data_cosmos_driver::{models::CosmosOperation, options::OperationOptions};
+use criterion::{criterion_group, criterion_main, Criterion, Throughput};
+use tokio::runtime::Builder;
+
+fn bench_point_read(c: &mut Criterion) {
+    let rt = Builder::new_current_thread()
+        .enable_all()
+        .build()
+        .expect("failed to create tokio runtime");
+
+    let mut group = c.benchmark_group("point_read");
+    group.throughput(Throughput::Elements(1));
+
+    match common::load_bench_config() {
+        BenchConfig::Mock => {
+            let (driver, item_ref) = rt.block_on(common::setup());
+            group.bench_function("mock", |b| {
+                b.to_async(&rt).iter(|| async {
+                    driver
+                        .execute_operation(
+                            CosmosOperation::read_item(item_ref.clone()),
+                            OperationOptions::default(),
+                        )
+                        .await
+                        .expect("execute_operation failed")
+                });
+            });
+        }
+        BenchConfig::Live => {
+            let (driver, item_ref) = rt.block_on(common::setup_live());
+            group
+                .sample_size(50)
+                .measurement_time(Duration::from_secs(30));
+            group.bench_function("live", |b| {
+                b.to_async(&rt).iter(|| async {
+                    driver
+                        .execute_operation(
+                            CosmosOperation::read_item(item_ref.clone()),
+                            OperationOptions::default(),
+                        )
+                        .await
+                        .expect("execute_operation failed")
+                });
+            });
+        }
+    }
+
+    group.finish();
+}
+
+criterion_group!(benches, bench_point_read);
+criterion_main!(benches);