Cosmos: Microbenchmark and some optimizations (#4159)

I wrote up a quick microbenchmark of the driver's `execute_operation`
code path, which excludes the I/O entirely by using a fixed Mock
Transport (gated behind an internal feature flag). Using that, I was
able to get a 40+% improvement in performance (again, these times
**exclude** I/O, but should reflect CPU usage reductions).

Before:

```
point_read              time:   [16.691 µs 16.734 µs 16.783 µs]
```

After

```
point_read              time:   [9.7474 µs 9.7864 µs 9.8319 µs]
```

Obviously, these are tiny numbers, but I suspect they add up under high
concurrency. We won't know for sure until this is run in the benchmarks.
They should be safe, so I'm game to merge and see, but we could also try
running the tip of this branch in the perf infra (cc @tvaron3 ) to get
real results.

The optimizations I worked on were predominantly to reduce duplicate
parsing and `HashMap::insert`/`HashMap::get` calls. Rather than calling
`HashMap::get_optional_str(...)` for each header when materializing a
`CosmosResponseHeaders`, we now iterate the map. Rather than re-parsing
`url::Url`, the `AccountRegion` type now directly deserializes the
gateway response into an `AccountEndpoint`, allowing us to reuse the
same `url::Url` for each request. Rather than constantly `format!`-ing
the `CosmosResourceReference`, we now produce a single `ResourcePaths`
from it and allow getting slices from it for the various purposes.

The benchmarks include instructions for running both profiled and
unprofiled runs, and [criterion](https://criterion-rs.github.io/) does a
great job of tracking results and showing you improvements over time.
Having benchmark infra like this in place will be useful as we keep
going.

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: analogrelay

Cargo.lock

@@ -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"

Cargo.toml

@@ -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/core/typespec/src/http/headers.rs

@@ -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/.cspell.json

@@ -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/driver_bridge.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/fault_injection/http_client.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/src/models/cosmos_response.rs

@@ -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/Cargo.toml

@@ -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/README.md

@@ -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);