[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 `ThroughputControlGroupOptions`, `PriorityLevel`, `ThroughputControlGroupRegistry`, and related types from the driver. Users can register throughput control groups on `CosmosClientBuilder` via `with_throughput_control_group()` to configure priority-based execution and throughput bucket server features. ([#4078](https://github.com/Azure/azure-sdk-for-rust/pull/4078))
 
 ### Breaking Changes
 

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
@@ -575,6 +567,7 @@ impl ContainerClient {
 
         // Execute through the driver.
         let driver_response = self
+            .context
             .driver
             .execute_operation(operation, options.operation)
             .await?;
@@ -706,7 +699,7 @@ impl ContainerClient {
         options.apply_headers(&mut headers);
 
         crate::query::executor::QueryExecutor::new(
-            self.pipeline.clone(),
+            self.context.pipeline.clone(),
             self.items_link.clone(),
             Context::default(),
             query,

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

@@ -2,22 +2,16 @@
 // Licensed under the MIT License.
 
 use crate::{
-    clients::DatabaseClient,
+    clients::{ClientContext, DatabaseClient},
     cosmos_request::CosmosRequest,
     models::{DatabaseProperties, ResourceResponse},
     operation_context::OperationType,
-    pipeline::GatewayPipeline,
+    options::ThroughputControlGroupRegistry,
     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 +64,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 +96,20 @@ 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
+    }
+
+    /// Returns the throughput control group registry.
+    ///
+    /// Groups are registered at client creation time via
+    /// [`CosmosClientBuilder::with_throughput_control_group()`].
+    pub fn throughput_control_groups(&self) -> &ThroughputControlGroupRegistry {
+        self.context.driver.runtime().throughput_control_groups()
     }
 
     /// Executes a query against databases in the account.
@@ -149,7 +142,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 +177,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,8 @@
 //! Builder for creating [`CosmosClient`] instances.
 
 use crate::{
+    clients::ClientContext,
+    options::ThroughputControlGroupOptions,
     pipeline::{AuthorizationPolicy, CosmosHeadersPolicy, GatewayPipeline},
     resource_context::{ResourceLink, ResourceType},
     CosmosAccountReference, CosmosClient, CosmosClientOptions, CosmosCredential, RoutingStrategy,
@@ -79,6 +81,8 @@ use azure_core::http::{ClientOptions, LoggingOptions, RetryOptions};
 #[derive(Default)]
 pub struct CosmosClientBuilder {
     options: CosmosClientOptions,
+    /// Throughput control groups to register on the driver runtime.
+    throughput_control_groups: Vec<ThroughputControlGroupOptions>,
     /// Whether to accept invalid TLS certificates when connecting to the emulator.
     #[cfg(feature = "allow_invalid_certificates")]
     allow_emulator_invalid_certificates: bool,
@@ -139,6 +143,16 @@ impl CosmosClientBuilder {
         self
     }
 
+    /// Registers a throughput control group on the driver runtime.
+    ///
+    /// Groups define throughput policies (priority level, throughput bucket) that
+    /// are applied to requests referencing the group name via
+    /// [`OperationOptions::throughput_control_group_names`](crate::OperationOptions::throughput_control_group_names).
+    pub fn with_throughput_control_group(mut self, group: ThroughputControlGroupOptions) -> Self {
+        self.throughput_control_groups.push(group);
+        self
+    }
+
     /// Builds the [`CosmosClient`] with the specified account reference and region selection strategy.
     ///
     /// The account reference bundles an endpoint and credential. You can create one using
@@ -348,17 +362,24 @@ impl CosmosClientBuilder {
             driver_runtime_builder =
                 driver_runtime_builder.with_fault_injection_rules(driver_fi_rules);
         }
+        for group in self.throughput_control_groups {
+            driver_runtime_builder = driver_runtime_builder
+                .register_throughput_control_group(group)
+                .map_err(|e| azure_core::Error::new(azure_core::error::ErrorKind::Other, e))?;
+        }
         let driver_runtime = driver_runtime_builder.build().await?;
         let driver = driver_runtime
             .get_or_create_driver(driver_account, None)
             .await?;
 
         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,
+            },
         })
     }
 }