[Cosmos] ThroughputControl API

Add SDK-level throughput control API for server-side priority-based
execution and throughput buckets (issue #3904).

New SDK types:
- ThroughputControlGroup enum with PriorityBased and ThroughputBucket
  variants, hiding the driver's unimplemented ClientSide variant
- Re-exports PriorityLevel, ThroughputControlGroupName, ContainerReference

Registration and resolution:
- CosmosClientOptions::with_throughput_control_group() for registration
- ThroughputControlGroupRegistry threaded through CosmosClient,
  DatabaseClient, and ContainerClient
- apply_throughput_control() on ItemWriteOptions, QueryOptions, and
  BatchOptions resolves group name to x-ms-cosmos-priority-level and
  x-ms-cosmos-throughput-bucket headers

Driver:
- ContainerReference::stub() doc-hidden helper for unit testing

Tests cover registration errors, header application for priority and
bucket groups, default group resolution, and empty registry behavior.

Fixes #3904

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

Author: tvaron3

sdk/cosmos/azure_data_cosmos/CHANGELOG.md

@@ -5,6 +5,7 @@
 ### Features Added
 
 - Added `CustomResponseBuilder` and `FaultInjectionRule::hit_count()` APIs for fault injection, enabling ergonomic construction of synthetic HTTP responses and test verification of rule activation counts. ([#3888](https://github.com/Azure/azure-sdk-for-rust/pull/3888))
+- Added throughput control API: re-exported `PriorityLevel`, `ThroughputControlGroupOptions`, `ThroughputControlGroupRegistry`, and related types from the driver. Users can register throughput control groups on `CosmosClientOptions` via `with_throughput_control_group()` to configure priority-based execution and throughput bucket server features. ([#3904](https://github.com/Azure/azure-sdk-for-rust/issues/3904))
 
 ### Breaking Changes
 

sdk/cosmos/azure_data_cosmos/src/clients/client_context.rs

@@ -0,0 +1,31 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+//! Shared infrastructure passed from [`CosmosClient`](super::CosmosClient) down to child clients.
+
+use std::sync::Arc;
+
+use crate::{
+    options::ThroughputControlGroupRegistry,
+    pipeline::GatewayPipeline,
+    routing::{
+        global_endpoint_manager::GlobalEndpointManager,
+        global_partition_endpoint_manager::GlobalPartitionEndpointManager,
+    },
+};
+use azure_data_cosmos_driver::CosmosDriver;
+
+/// Shared infrastructure threaded from [`CosmosClient`](super::CosmosClient)
+/// through [`DatabaseClient`](super::DatabaseClient) to
+/// [`ContainerClient`](super::ContainerClient).
+///
+/// Bundling these fields avoids passing them individually through every
+/// constructor in the client hierarchy.
+#[derive(Clone, Debug)]
+pub(crate) struct ClientContext {
+    pub pipeline: Arc<GatewayPipeline>,
+    pub driver: Arc<CosmosDriver>,
+    pub global_endpoint_manager: Arc<GlobalEndpointManager>,
+    pub global_partition_endpoint_manager: Arc<GlobalPartitionEndpointManager>,
+    pub throughput_control_groups: ThroughputControlGroupRegistry,
+}

sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs

@@ -2,13 +2,12 @@
 // Licensed under the MIT License.
 
 use crate::{
-    clients::OffersClient,
+    clients::{ClientContext, OffersClient},
     models::{
         BatchResponse, ContainerProperties, CosmosResponse, ItemResponse, ResourceResponse,
         ThroughputProperties,
     },
     options::{BatchOptions, QueryOptions, ReadContainerOptions},
-    pipeline::GatewayPipeline,
     resource_context::{ResourceLink, ResourceType},
     transactional_batch::TransactionalBatch,
     DeleteContainerOptions, FeedItemIterator, ItemReadOptions, ItemWriteOptions, PartitionKey,
@@ -19,13 +18,10 @@ use std::sync::Arc;
 use crate::cosmos_request::CosmosRequest;
 use crate::handler::container_connection::ContainerConnection;
 use crate::operation_context::OperationType;
-use crate::routing::global_endpoint_manager::GlobalEndpointManager;
-use crate::routing::global_partition_endpoint_manager::GlobalPartitionEndpointManager;
 use crate::routing::partition_key_range_cache::PartitionKeyRangeCache;
 use azure_core::http::headers::AsHeaders;
 use azure_core::http::Context;
 use azure_data_cosmos_driver::models::{ContainerReference, CosmosOperation, ItemReference};
-use azure_data_cosmos_driver::CosmosDriver;
 use serde::{de::DeserializeOwned, Serialize};
 
 /// A client for working with a specific container in a Cosmos DB account.
@@ -35,30 +31,27 @@ use serde::{de::DeserializeOwned, Serialize};
 pub struct ContainerClient {
     link: ResourceLink,
     items_link: ResourceLink,
-    pipeline: Arc<GatewayPipeline>,
     container_connection: Arc<ContainerConnection>,
     container_id: String,
-    driver: Arc<CosmosDriver>,
     container_ref: ContainerReference,
+    context: ClientContext,
 }
 
 impl ContainerClient {
     pub(crate) async fn new(
-        pipeline: Arc<GatewayPipeline>,
+        context: ClientContext,
         database_link: &ResourceLink,
         container_id: &str,
         database_id: &str,
-        driver: Arc<CosmosDriver>,
-        global_endpoint_manager: Arc<GlobalEndpointManager>,
-        global_partition_endpoint_manager: Arc<GlobalPartitionEndpointManager>,
     ) -> azure_core::Result<Self> {
         let link = database_link
             .feed(ResourceType::Containers)
             .item(container_id);
         let items_link = link.feed(ResourceType::Documents);
 
         // Eagerly resolve immutable container metadata from the driver.
-        let container_ref = driver
+        let container_ref = context
+            .driver
             .resolve_container(database_id, container_id)
             .await
             .map_err(|e| {
@@ -68,25 +61,24 @@ impl ContainerClient {
             })?;
 
         let partition_key_range_cache = Arc::from(PartitionKeyRangeCache::new(
-            pipeline.clone(),
+            context.pipeline.clone(),
             database_link.clone(),
-            global_endpoint_manager.clone(),
+            context.global_endpoint_manager.clone(),
         ));
         let container_connection = Arc::from(ContainerConnection::new(
-            pipeline.clone(),
+            context.pipeline.clone(),
             partition_key_range_cache,
-            global_partition_endpoint_manager.clone(),
+            context.global_partition_endpoint_manager.clone(),
             container_ref.clone(),
         ));
 
         Ok(Self {
             link,
             items_link,
-            pipeline,
             container_connection,
             container_id: container_id.to_string(),
-            driver,
             container_ref,
+            context,
         })
     }
 
@@ -190,7 +182,7 @@ impl ContainerClient {
             .resource_id
             .expect("service should always return a '_rid' for a container");
 
-        let offers_client = OffersClient::new(self.pipeline.clone(), resource_id);
+        let offers_client = OffersClient::new(self.context.pipeline.clone(), resource_id);
         offers_client.read(Context::default()).await
     }
 
@@ -222,7 +214,7 @@ impl ContainerClient {
             .resource_id
             .expect("service should always return a '_rid' for a container");
 
-        let offers_client = OffersClient::new(self.pipeline.clone(), resource_id);
+        let offers_client = OffersClient::new(self.context.pipeline.clone(), resource_id);
         offers_client
             .replace(Context::default(), throughput)
             .await
@@ -332,6 +324,11 @@ impl ContainerClient {
                 .excluded_regions(excluded_regions)
                 .build()?;
         options.apply_headers(&mut cosmos_request.headers);
+        options.apply_throughput_control(
+            &mut cosmos_request.headers,
+            &self.context.throughput_control_groups,
+            &self.container_ref,
+        );
 
         self.container_connection
             .send(cosmos_request, Context::default())
@@ -420,6 +417,11 @@ impl ContainerClient {
             .excluded_regions(excluded_regions)
             .build()?;
         options.apply_headers(&mut cosmos_request.headers);
+        options.apply_throughput_control(
+            &mut cosmos_request.headers,
+            &self.context.throughput_control_groups,
+            &self.container_ref,
+        );
 
         self.container_connection
             .send(cosmos_request, Context::default())
@@ -510,6 +512,11 @@ impl ContainerClient {
                 .excluded_regions(excluded_regions)
                 .build()?;
         options.apply_headers(&mut cosmos_request.headers);
+        options.apply_throughput_control(
+            &mut cosmos_request.headers,
+            &self.context.throughput_control_groups,
+            &self.container_ref,
+        );
 
         return self
             .container_connection
@@ -575,6 +582,7 @@ impl ContainerClient {
 
         // Execute through the driver.
         let driver_response = self
+            .context
             .driver
             .execute_operation(operation, options.operation)
             .await?;
@@ -620,6 +628,11 @@ impl ContainerClient {
             .excluded_regions(excluded_regions)
             .build()?;
         options.apply_headers(&mut cosmos_request.headers);
+        options.apply_throughput_control(
+            &mut cosmos_request.headers,
+            &self.context.throughput_control_groups,
+            &self.container_ref,
+        );
 
         self.container_connection
             .send(cosmos_request, Context::default())
@@ -704,9 +717,14 @@ impl ContainerClient {
             headers.insert(name, value);
         }
         options.apply_headers(&mut headers);
+        options.apply_throughput_control(
+            &mut headers,
+            &self.context.throughput_control_groups,
+            &self.container_ref,
+        );
 
         crate::query::executor::QueryExecutor::new(
-            self.pipeline.clone(),
+            self.context.pipeline.clone(),
             self.items_link.clone(),
             Context::default(),
             query,
@@ -771,6 +789,11 @@ impl ContainerClient {
                 .json(batch.operations())
                 .build()?;
         options.apply_headers(&mut cosmos_request.headers);
+        options.apply_throughput_control(
+            &mut cosmos_request.headers,
+            &self.context.throughput_control_groups,
+            &self.container_ref,
+        );
 
         self.container_connection
             .send(cosmos_request, Context::default())

sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client.rs

@@ -2,22 +2,15 @@
 // Licensed under the MIT License.
 
 use crate::{
-    clients::DatabaseClient,
+    clients::{ClientContext, DatabaseClient},
     cosmos_request::CosmosRequest,
     models::{DatabaseProperties, ResourceResponse},
     operation_context::OperationType,
-    pipeline::GatewayPipeline,
     resource_context::ResourceLink,
-    routing::{
-        global_endpoint_manager::GlobalEndpointManager,
-        global_partition_endpoint_manager::GlobalPartitionEndpointManager,
-    },
     CreateDatabaseOptions, FeedItemIterator, Query, QueryDatabasesOptions,
 };
 use azure_core::http::{Context, Url};
-use azure_data_cosmos_driver::CosmosDriver;
 use serde::Serialize;
-use std::sync::Arc;
 
 pub use super::cosmos_client_builder::CosmosClientBuilder;
 
@@ -70,10 +63,7 @@ pub use super::cosmos_client_builder::CosmosClientBuilder;
 #[derive(Debug, Clone)]
 pub struct CosmosClient {
     pub(crate) databases_link: ResourceLink,
-    pub(crate) pipeline: Arc<GatewayPipeline>,
-    pub(crate) driver: Arc<CosmosDriver>,
-    pub(crate) global_endpoint_manager: Arc<GlobalEndpointManager>,
-    pub(crate) global_partition_endpoint_manager: Arc<GlobalPartitionEndpointManager>,
+    pub(crate) context: ClientContext,
 }
 
 impl CosmosClient {
@@ -105,18 +95,12 @@ impl CosmosClient {
     /// # Arguments
     /// * `id` - The ID of the database.
     pub fn database_client(&self, id: &str) -> DatabaseClient {
-        DatabaseClient::new(
-            self.pipeline.clone(),
-            id,
-            self.driver.clone(),
-            self.global_endpoint_manager.clone(),
-            self.global_partition_endpoint_manager.clone(),
-        )
+        DatabaseClient::new(self.context.clone(), id)
     }
 
     /// Gets the endpoint of the database account this client is connected to.
     pub fn endpoint(&self) -> &Url {
-        &self.pipeline.endpoint
+        &self.context.pipeline.endpoint
     }
 
     /// Executes a query against databases in the account.
@@ -149,7 +133,7 @@ impl CosmosClient {
         _options: Option<QueryDatabasesOptions>,
     ) -> azure_core::Result<FeedItemIterator<DatabaseProperties>> {
         crate::query::executor::QueryExecutor::new(
-            self.pipeline.clone(),
+            self.context.pipeline.clone(),
             self.databases_link.clone(),
             Context::default(),
             query.into(),
@@ -184,7 +168,8 @@ impl CosmosClient {
                 .json(&RequestBody { id })
                 .build()?;
 
-        self.pipeline
+        self.context
+            .pipeline
             .send(cosmos_request, Context::default())
             .await
             .map(ResourceResponse::new)

sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client_builder.rs

@@ -4,6 +4,7 @@
 //! Builder for creating [`CosmosClient`] instances.
 
 use crate::{
+    clients::ClientContext,
     pipeline::{AuthorizationPolicy, CosmosHeadersPolicy, GatewayPipeline},
     resource_context::{ResourceLink, ResourceType},
     CosmosAccountReference, CosmosClient, CosmosClientOptions, CosmosCredential, RoutingStrategy,
@@ -318,6 +319,8 @@ impl CosmosClientBuilder {
             },
         ));
 
+        let throughput_control_groups = self.options.throughput_control_groups.clone();
+
         let pipeline = Arc::new(GatewayPipeline::new(
             endpoint.clone(),
             pipeline_core,
@@ -355,10 +358,13 @@ impl CosmosClientBuilder {
 
         Ok(CosmosClient {
             databases_link: ResourceLink::root(ResourceType::Databases),
-            pipeline,
-            driver,
-            global_endpoint_manager,
-            global_partition_endpoint_manager,
+            context: ClientContext {
+                pipeline,
+                driver,
+                global_endpoint_manager,
+                global_partition_endpoint_manager,
+                throughput_control_groups,
+            },
         })
     }
 }