[Cosmos] ThroughputControl API

[tvaron3]

Add server-side throughput control API enabling priority-based execution and throughput buckets.

Flow

Registration (build time)

graph LR
    A[CosmosClientBuilder] -->|with_throughput_control_group| B[stores Vec of groups]
    B -->|build| C[CosmosDriverRuntimeBuilder]
    C -->|register_throughput_control_group| D[ThroughputControlGroupRegistry]
    D --> E[CosmosDriverRuntime]
    E --> F[CosmosDriver]
    F --> G[CosmosClient]

Loading

Per-request resolution

graph TD
    A[OperationOptions::throughput_control_group_names] -->|resolve| B{Names provided?}
    B -->|Yes| C[Look up each name in registry]
    C --> D{Any found?}
    D -->|Yes| E[Merge into ThroughputControlGroupSnapshot]
    D -->|No| F[Return None - no headers set]
    B -->|No| G[Look up default group for container]
    G -->|Found| E
    G -->|Not found| F
    E --> H[build_transport_request]
    H -->|priority_level| I[x-ms-cosmos-priority-level header]
    H -->|throughput_bucket| J[x-ms-cosmos-throughput-bucket header]

Loading

Combining groups

A single request can reference multiple groups to combine features:

let mut operation = OperationOptions::default();
operation.throughput_control_group_names = Some(vec![
    ThroughputControlGroupName::new("my-priority-group"),
    ThroughputControlGroupName::new("my-bucket-group"),
]);

The driver merges the first priority_level and first throughput_bucket found across all named groups into a single snapshot.

Driver changes (azure_data_cosmos_driver)

Throughput control types

• Removed ClientSide variant from ThroughputControlGroupOptions enum (deferred for GA; enum is #[non_exhaustive], can be re-added later — tracked in [Cosmos] Client-side throughput control #4089)
• Removed dead ThroughputTarget enum (will return with ClientSide)
• Removed target_throughput and is_client_side from ThroughputControlGroupSnapshot

Multi-group per request

• Changed OperationOptions::throughput_control_group_name (singular) to throughput_control_group_names (Option<Vec<...>>) so one request can reference both a priority group and a bucket group
• Resolution merges groups into a single snapshot (first-wins for each field)
• Explicit names that fail to resolve return None with tracing::warn! (no silent fallback to default)

Pipeline header application

• effective_throughput_control_group() passes resolved ThroughputControlGroupSnapshot through execute_operation_pipeline() to build_transport_request()
• Inserts x-ms-cosmos-priority-level and x-ms-cosmos-throughput-bucket headers for all operations

Other

• Added PRIORITY_LEVEL and THROUGHPUT_BUCKET constants to request_header_names
• Extracted TransportRequestContext struct to reduce build_transport_request parameter count
• Fixed ServerSidePriorityBasedThrottling doc URL (was linking to throughput-buckets)
• Added tracing::debug! in set_* no-op paths (calling setter on wrong variant)

SDK changes (azure_data_cosmos)

Re-exports

• ThroughputControlGroupOptions, ThroughputControlGroupRegistry, ThroughputControlGroupKey, ThroughputControlGroupSnapshot, ThroughputControlGroupRegistrationError
• PriorityLevel, ThroughputControlGroupName, ContainerReference

Registration

• CosmosClientBuilder::with_throughput_control_group() registers groups, passes to CosmosDriverRuntimeBuilder at build time
• CosmosClient::throughput_control_groups() read-only access to the driver's registry

ClientContext refactor

• Extracted ClientContext struct bundling shared infrastructure (pipeline, driver, global_endpoint_manager, global_partition_endpoint_manager) threaded through CosmosClient to DatabaseClient to ContainerClient
• Reduces ContainerClient::new from 8 arguments to 4

Depends on #4059

Fixes #3904

[tvaron3]

/azp run rust - cosmos - weekly

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[Copilot]

Pull request overview

This PR introduces an SDK-level Throughput Control API for Cosmos DB by allowing users to register named throughput control groups (priority-based execution or throughput buckets) at client creation time, then applying the appropriate service headers on supported operations.

Changes:

• Added ThroughputControlGroup (SDK-facing) and re-exported PriorityLevel, ThroughputControlGroupName, and ContainerReference for public API usage.
• Threaded a ThroughputControlGroupRegistry through the client hierarchy and applied throughput control headers for write/query/batch requests.
• Added a ContainerReference::stub() helper in the driver for unit tests and documented the feature in the SDK changelog.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
sdk/cosmos/azure_data_cosmos/src/options/mod.rs	Adds SDK throughput control group type + client registration API and implements header application helpers + tests.
sdk/cosmos/azure_data_cosmos/src/clients/mod.rs	Introduces ClientContext to carry shared client infrastructure including throughput control registry.
sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client.rs	Switches CosmosClient to store shared dependencies via ClientContext.
sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client_builder.rs	Captures throughput control registry during build and populates it into ClientContext.
sdk/cosmos/azure_data_cosmos/src/clients/database_client.rs	Refactors to use ClientContext and passes it through to container construction.
sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs	Applies throughput control headers for gateway-pipeline operations and uses ClientContext throughout.
sdk/cosmos/azure_data_cosmos/CHANGELOG.md	Adds an unreleased changelog entry for the new public Throughput Control API.
sdk/cosmos/azure_data_cosmos_driver/src/models/resource_reference.rs	Adds ContainerReference::stub() for unit testing without live service metadata resolution.

[analogrelay]

Rather than adding this to the SDK, should we adjust the driver's options to match this and then re-export it? I'd like to avoid adding new conversion layers if we can.

[analogrelay]

Java's ThroughputControlGroupBuilder seems to allow setting both a throughput bucket AND a priority level. I kinda viewed TCGs as a struct holding optional Priority Level, and Throughput Bucket values (for server-side throughput control) as well as (in the future) Target Throughput (for client-side throughput control).

[tvaron3]

Addressed — instead of merging into a single struct, we kept the enum (each variant represents a distinct server-side feature) and changed OperationOptions::throughput_control_group_name to throughput_control_group_names: Option<Vec<ThroughputControlGroupName>>. This lets a single request reference multiple groups (e.g., one priority group + one bucket group) and the driver merges their snapshots, applying both x-ms-cosmos-priority-level and x-ms-cosmos-throughput-bucket headers. The ClientSide variant was removed for GA (#[non_exhaustive] allows re-adding it later with target_throughput).