Improve Cosmos perf diagnostics

sdk/cosmos/azure_data_cosmos_driver/src/driver/transport/gateway_v2_dispatch.rs

@@ -423,6 +423,12 @@ pub(crate) fn unwrap_response_for_gateway_v2(
     if let Some(charge) = response.request_charge {
         headers.insert(response_header_names::REQUEST_CHARGE, charge.to_string());
     }
+    if let Some(duration_ms) = response.backend_request_duration_ms {
+        headers.insert(
+            response_header_names::SERVER_DURATION_MS,
+            duration_ms.to_string(),
+        );
+    }
     if let Some(token) = response.session_token {
         // GW2 surfaces only the vector portion of the session token, with the
         // partition key range id carried separately. Classic gateway emits
@@ -1931,6 +1937,7 @@ mod tests {
                 |tokens| {
                     write_u32_token(tokens, 0x001C, 1002);
                     write_double_token(tokens, 0x0015, 3.5);
+                    write_double_token(tokens, 0x0051, 12.75);
                     write_string_token(tokens, 0x003E, "1:2#3");
                     write_string_token(tokens, 0x0004, "\"etag\"");
                     write_string_token(tokens, 0x0003, "continuation");
@@ -1963,6 +1970,14 @@ mod tests {
                 .get_optional_str(&HeaderName::from_static("x-ms-request-charge")),
             Some("3.5")
         );
+        assert_eq!(
+            unwrapped.headers.get_optional_str(&HeaderName::from_static(
+                response_header_names::SERVER_DURATION_MS
+            )),
+            Some("12.75")
+        );
+        let parsed_headers = crate::models::CosmosResponseHeaders::from_headers(&unwrapped.headers);
+        assert_eq!(parsed_headers.server_duration_ms, Some(12.75));
         assert_eq!(
             unwrapped
                 .headers

sdk/cosmos/azure_data_cosmos_driver/src/driver/transport/rntbd/response.rs

@@ -34,6 +34,8 @@ pub(crate) struct RntbdResponse {
     pub(crate) lsn: Option<i64>,
     /// Request charge in request units.
     pub(crate) request_charge: Option<f64>,
+    /// Backend request processing duration in milliseconds.
+    pub(crate) backend_request_duration_ms: Option<f64>,
     /// Owner full name metadata.
     pub(crate) owner_full_name: Option<String>,
     /// Partition key range identifier.
@@ -87,6 +89,7 @@ impl RntbdResponse {
         let mut retry_after_ms = None;
         let mut lsn = None;
         let mut request_charge = None;
+        let mut backend_request_duration_ms = None;
         let mut owner_full_name = None;
         let mut sub_status = None;
         let mut partition_key_range_id = None;
@@ -116,6 +119,10 @@ impl RntbdResponse {
                 Ok(RntbdResponseToken::RequestCharge) => {
                     request_charge = Some(expect_f64(token, "RequestCharge")?);
                 }
+                Ok(RntbdResponseToken::BackendRequestDurationMilliseconds) => {
+                    backend_request_duration_ms =
+                        Some(expect_f64(token, "BackendRequestDurationMilliseconds")?);
+                }
                 Ok(RntbdResponseToken::OwnerFullName) => {
                     owner_full_name = Some(expect_string(token, "OwnerFullName")?);
                 }
@@ -164,6 +171,7 @@ impl RntbdResponse {
             retry_after_ms,
             lsn,
             request_charge,
+            backend_request_duration_ms,
             owner_full_name,
             partition_key_range_id,
             item_lsn,
@@ -229,6 +237,32 @@ mod tests {
         assert!(response.body.is_empty());
     }
 
+    #[test]
+    fn backend_request_duration_is_decoded() {
+        let mut frame = response_header(StatusCode::Ok);
+        Token::new(0x0051, TokenValue::Double(12.75))
+            .write_to(&mut frame)
+            .unwrap();
+        patch_total_len(&mut frame);
+
+        let response = RntbdResponse::read(&frame).unwrap();
+
+        assert_eq!(response.backend_request_duration_ms, Some(12.75));
+    }
+
+    #[test]
+    fn backend_request_duration_rejects_wrong_token_type() {
+        let mut frame = response_header(StatusCode::Ok);
+        Token::new(0x0051, TokenValue::ULong(12))
+            .write_to(&mut frame)
+            .unwrap();
+        patch_total_len(&mut frame);
+
+        let err = RntbdResponse::read(&frame).unwrap_err();
+
+        assert_eq!(*err.kind(), azure_core::error::ErrorKind::DataConversion);
+    }
+
     #[test]
     fn total_length_past_buffer_is_rejected() {
         let mut frame = response_header(StatusCode::Ok);

sdk/cosmos/azure_data_cosmos_driver/src/driver/transport/rntbd/tokens.rs

@@ -671,6 +671,8 @@ pub(super) enum RntbdResponseToken {
     TransportRequestId,
     /// Session token.
     SessionToken,
+    /// Backend request processing duration in milliseconds.
+    BackendRequestDurationMilliseconds,
 }
 
 impl TryFrom<u16> for RntbdResponseToken {
@@ -691,6 +693,7 @@ impl TryFrom<u16> for RntbdResponseToken {
             0x0029 => Ok(Self::GlobalCommittedLsn),
             0x0035 => Ok(Self::TransportRequestId),
             0x003E => Ok(Self::SessionToken),
+            0x0051 => Ok(Self::BackendRequestDurationMilliseconds),
             _ => Err(()),
         }
     }

sdk/cosmos/azure_data_cosmos_perf/README.md

@@ -70,6 +70,7 @@ cargo run -p azure_data_cosmos_perf -- \
 | `--throughput` | `100000` | Throughput (RU/s) when creating the container |
 | `--default-ttl` | `3600` | Default TTL in seconds for items (0 to disable) |
 | `--report-interval` | `300` | Stats reporting interval in seconds |
+| `--diagnostics-threshold-ms` | — | Log detailed diagnostics for successful operations slower than this threshold; disabled when omitted |
 | `--results-container` | `perfresults` | Container for storing perf results and error documents |
 | `--results-endpoint` | — | Cosmos DB endpoint for results (omit to use same account as `--endpoint`) |
 | `--results-database` | `perfdb` | Database name on the results account |
@@ -82,6 +83,9 @@ cargo run -p azure_data_cosmos_perf -- \
 | `--no-queries` | `false` | Disable query operations |
 | `--no-upserts` | `false` | Disable upsert operations |
 | `--no-creates` | `false` | Disable create operations |
+| `--no-feed-range-queries` | `false` | Disable per-feed-range query operations |
+| `--feed-range-query-max-pages` | `4` | Maximum pages consumed by each per-feed-range query |
+| `--feed-range-refresh-secs` | `60` | Interval between feed-range cache refreshes (0 disables refresh) |
 
 ### Examples
 

sdk/cosmos/azure_data_cosmos_perf/src/config.rs

@@ -66,11 +66,18 @@ pub struct Config {
     /// Disable the per-feed-range query operation.
     ///
     /// When enabled (the default), the harness fans out
-    /// `SELECT VALUE COUNT(1) FROM c` across the container's physical
-    /// partitions, round-robin one range per call.
+    /// `SELECT * FROM c` across the container's physical partitions,
+    /// round-robin one range per call.
     #[arg(long, default_value_t = false)]
     pub no_feed_range_queries: bool,
 
+    /// Maximum number of pages consumed by each per-feed-range query.
+    ///
+    /// Bounds the work performed by one query as the container grows while
+    /// still exercising continuation-token handling across multiple pages.
+    #[arg(long, default_value_t = 4)]
+    pub feed_range_query_max_pages: usize,
+
     /// Interval in seconds between background refreshes of the cached feed
     /// range list. Set to 0 to disable the background refresher (the cache
     /// then stays at the seed snapshot taken once at startup).
@@ -82,9 +89,35 @@ pub struct Config {
     pub feed_range_refresh_secs: u64,
 
     /// Number of concurrent operations (minimum: 1).
+    ///
+    /// Drives the default closed-loop mode: the harness spawns this many
+    /// persistent worker tasks, each serially issuing one operation at a
+    /// time (so this is also the maximum number of in-flight requests).
+    /// Ignored when `--target-rate` is set, which switches to open-loop
+    /// (fixed arrival rate) issuance instead.
     #[arg(long, default_value_t = 50)]
     pub concurrency: usize,
 
+    /// Target arrival rate in operations per second (open-loop mode).
+    ///
+    /// When set, requests are issued at this fixed rate regardless of how
+    /// fast prior requests complete, eliminating the coordinated-omission
+    /// bias of the closed-loop `--concurrency` model. In-flight requests
+    /// are bounded by `--max-in-flight`; once that bound is hit, excess
+    /// issuances are skipped and counted (not buffered) so a slow backend
+    /// surfaces as visible omission rather than unbounded memory growth.
+    /// Omit to use the default closed-loop `--concurrency` model.
+    #[arg(long)]
+    pub target_rate: Option<u64>,
+
+    /// Maximum number of in-flight requests in open-loop mode (minimum: 1).
+    ///
+    /// Only used when `--target-rate` is set. Bounds memory/socket growth
+    /// when the backend can't keep up with the target rate: issuances that
+    /// would exceed this cap are skipped and counted instead of queued.
+    #[arg(long, default_value_t = 100_000)]
+    pub max_in_flight: usize,
+
     /// Run duration in seconds. Omit for indefinite.
     #[arg(long)]
     pub duration: Option<u64>,
@@ -97,6 +130,12 @@ pub struct Config {
     #[arg(long, default_value_t = 300)]
     pub report_interval: u64,
 
+    /// Log detailed diagnostics for successful operations slower than this threshold in milliseconds.
+    ///
+    /// Omit to disable diagnostics logging for successful operations.
+    #[arg(long)]
+    pub diagnostics_threshold_ms: Option<u64>,
+
     /// Throughput (RU/s) to provision when creating the container.
     #[arg(long, default_value_t = 100000)]
     pub throughput: usize,

sdk/cosmos/azure_data_cosmos_perf/src/main.rs

@@ -130,10 +130,24 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         eprintln!("Error: --concurrency cannot exceed {}.", u32::MAX);
         std::process::exit(1);
     }
+    if let Some(rate) = config.target_rate {
+        if rate == 0 {
+            eprintln!("Error: --target-rate must be at least 1.");
+            std::process::exit(1);
+        }
+        if config.max_in_flight == 0 {
+            eprintln!("Error: --max-in-flight must be at least 1 when --target-rate is set.");
+            std::process::exit(1);
+        }
+    }
     if config.seed_count == 0 {
         eprintln!("Error: --seed-count must be at least 1.");
         std::process::exit(1);
     }
+    if config.feed_range_query_max_pages == 0 {
+        eprintln!("Error: --feed-range-query-max-pages must be at least 1.");
+        std::process::exit(1);
+    }
 
     // Resolve the (optional) user-agent suffix. Composes
     // `{configured}-{first-8-chars-of-workload_id}` so a single perf
@@ -324,6 +338,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
             .ok()
             .and_then(|v| v.parse::<bool>().ok())
             .unwrap_or(true),
+        diagnostics_threshold_ms: config.diagnostics_threshold_ms,
         pyroscope_enabled: std::env::var("PYROSCOPE_SERVER_URL")
             .map(|v| !v.is_empty())
             .unwrap_or(false),
@@ -346,6 +361,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         feed_range_refresher,
         stats,
         concurrency: config.concurrency,
+        target_rate: config.target_rate,
+        max_in_flight: config.max_in_flight,
+        diagnostics_threshold: config.diagnostics_threshold_ms.map(Duration::from_millis),
         duration,
         report_interval: Duration::from_secs(config.report_interval),
         results_container,

sdk/cosmos/azure_data_cosmos_perf/src/operations/create_item.rs

@@ -11,7 +11,7 @@ use azure_data_cosmos::options::ItemWriteOptions;
 use rand::RngExt;
 use uuid::Uuid;
 
-use super::{extract_backend_duration, Operation, PerfItem};
+use super::{extract_backend_duration, Operation, OperationResult, PerfItem};
 use crate::seed::{SeededItem, SharedItems};
 
 /// Creates a new item with a unique ID and partition key.
@@ -39,7 +39,8 @@ impl Operation for CreateItemOperation {
     async fn execute(
         &self,
         container: &ContainerClient,
-    ) -> azure_data_cosmos::Result<Option<std::time::Duration>> {
+        capture_diagnostics: bool,
+    ) -> azure_data_cosmos::Result<OperationResult> {
         let id = Uuid::new_v4().to_string();
         let partition_key = Uuid::new_v4().to_string();
         let value = rand::rng().random_range(0..u64::MAX);
@@ -55,8 +56,9 @@ impl Operation for CreateItemOperation {
             .create_item(&item.partition_key, &id, &item, self.options.clone())
             .await?;
         let backend = extract_backend_duration(response.headers());
+        let diagnostics = capture_diagnostics.then(|| response.diagnostics());
 
         self.items.push(SeededItem { id, partition_key });
-        Ok(backend)
+        Ok(OperationResult::single(backend, diagnostics))
     }
 }

sdk/cosmos/azure_data_cosmos_perf/src/operations/feed_range_query.rs

@@ -6,8 +6,8 @@
 //! Each `execute()` runs `SELECT * FROM c` against ONE `FeedRange`,
 //! round-robin'd from a cache shared with
 //! [`FeedRangeRefresher`](super::feed_range_refresher::FeedRangeRefresher),
-//! and drains every page so the harness exercises continuation-token
-//! handling across multi-page responses per feed range.
+//! and drains a bounded number of pages so the harness exercises
+//! continuation-token handling without query work growing with the container.
 //! The harness's existing worker pool provides concurrency — N workers
 //! each issue one query at a time, naturally matching the concurrency
 //! of every other operation.
@@ -22,7 +22,7 @@ use azure_data_cosmos::{feed::FeedRange, CosmosStatus, FeedScope, Query};
 use azure_data_cosmos_driver::error::CosmosError as DriverCosmosError;
 use futures::StreamExt;
 
-use super::{extract_backend_duration, Operation};
+use super::{extract_backend_duration, Operation, OperationResult};
 
 /// Shared, swappable feed-range cache.
 ///
@@ -34,13 +34,15 @@ pub type FeedRangeCache = Arc<RwLock<Arc<Vec<FeedRange>>>>;
 pub struct FeedRangeQueryOperation {
     cache: FeedRangeCache,
     cursor: AtomicUsize,
+    max_pages: usize,
 }
 
 impl FeedRangeQueryOperation {
-    pub fn new(cache: FeedRangeCache) -> Self {
+    pub fn new(cache: FeedRangeCache, max_pages: usize) -> Self {
         Self {
             cache,
             cursor: AtomicUsize::new(0),
+            max_pages,
         }
     }
 }
@@ -54,7 +56,8 @@ impl Operation for FeedRangeQueryOperation {
     async fn execute(
         &self,
         container: &ContainerClient,
-    ) -> azure_data_cosmos::Result<Option<Duration>> {
+        capture_diagnostics: bool,
+    ) -> azure_data_cosmos::Result<OperationResult> {
         // Snapshot the current ranges; release the lock immediately.
         let snapshot = {
             let guard = self.cache.read().expect("feed-range cache lock poisoned");
@@ -80,20 +83,25 @@ impl Operation for FeedRangeQueryOperation {
         let mut stream =
             Box::pin(container.query_items::<super::PerfItem>(query, FeedScope::range(fr), None))
                 .await?
-                .into_pages();
+                .into_pages()
+                .take(self.max_pages);
 
         // Sum backend durations across pages so a multi-page response
         // reports the total server processing time, matching the
         // client-observed elapsed which wraps the entire stream drain.
         let mut backend_total: Option<Duration> = None;
+        let mut diagnostics = Vec::new();
         while let Some(result) = stream.next().await {
             let page = result?;
             if let Some(d) = extract_backend_duration(page.headers()) {
                 backend_total = Some(backend_total.unwrap_or_default() + d);
             }
+            if capture_diagnostics {
+                diagnostics.push(page.diagnostics());
+            }
         }
 
-        Ok(backend_total)
+        Ok(OperationResult::paged(backend_total, diagnostics))
     }
 }
 
@@ -108,7 +116,7 @@ mod tests {
 
     #[test]
     fn cursor_round_robins_and_wraps() {
-        let op = FeedRangeQueryOperation::new(cache_with_len(3));
+        let op = FeedRangeQueryOperation::new(cache_with_len(3), 4);
         // 6 fetches over a 3-element cache should visit each index twice
         // (0,1,2,0,1,2 modulo 3).
         let picks: Vec<usize> = (0..6)
@@ -122,7 +130,7 @@ mod tests {
         // Construct against a zero-length cache and confirm snapshot
         // is empty (execute() would surface this as an error instead
         // of dividing by zero).
-        let op = FeedRangeQueryOperation::new(cache_with_len(0));
+        let op = FeedRangeQueryOperation::new(cache_with_len(0), 4);
         let guard = op.cache.read().unwrap();
         assert!(guard.is_empty());
     }