Skip to content

Expose async state of replace throughput API#4096

Draft
tvaron3 wants to merge 1 commit intoAzure:mainfrom
tvaron3:tvaron3/replace-throughput-lro
Draft

Expose async state of replace throughput API#4096
tvaron3 wants to merge 1 commit intoAzure:mainfrom
tvaron3:tvaron3/replace-throughput-lro

Conversation

@tvaron3
Copy link
Copy Markdown
Member

@tvaron3 tvaron3 commented Apr 5, 2026

What

Resolves #3792 — makes replace_throughput clearly expose the async nature of the Cosmos DB throughput API.

The Cosmos DB service may process throughput changes asynchronously (HTTP 202 with x-ms-offer-replace-pending: true), but the previous API returned a simple CosmosResponse<ThroughputProperties> with no indication of whether the change was still in-flight.

Changes

  • Renamed replace_throughputbegin_replace_throughput on ContainerClient and DatabaseClient per the Azure SDK LRO naming guideline (begin_ prefix for LROs).

  • Added ThroughputPoller — a custom type that implements both IntoFuture (simple .await) and Stream (manual polling) to handle sync vs async throughput operations:

    • HTTP 200: Completes immediately — single-element stream, .await resolves instantly.
    • HTTP 202 / x-ms-offer-replace-pending: true: Polls the offer resource at 5-second intervals until the operation completes.
    • Caches the offer RID from the initial response for efficient single-GET polling (no query per tick).
    • Propagates caller Context through polling for distributed tracing.
  • Why not azure_core::Poller? Cosmos throughput LROs use a header-based detection pattern (x-ms-offer-replace-pending), not the Operation-Location URL pattern that azure_core::Poller and StatusMonitor expect. A custom type follows the Cosmos SDK convention of wrapping azure_core internals (like CosmosResponse wraps Response, FeedItemIterator wraps BoxStream).

Migration

// Before
let throughput = client
    .replace_throughput(ThroughputProperties::manual(500), None)
    .await?
    .into_model()?;

// After (simple — just await the final result)
let throughput = client
    .begin_replace_throughput(ThroughputProperties::manual(500), None)
    .await?   // sends the replace request
    .await?   // awaits the poller for final result
    .into_model()?;

// After (advanced — poll for progress via Stream)
use futures::TryStreamExt;
let mut poller = client
    .begin_replace_throughput(ThroughputProperties::manual(500), None)
    .await?;
while let Some(status) = poller.try_next().await? {
    println!("Request charge: {:?}", status.request_charge());
}

Testing

  • Updated all existing throughput test callsites.
  • Added container_throughput_high_ru — replaces throughput to 11000 RU/s (may trigger async processing).
  • Added container_throughput_stream_polling — exercises the Stream interface directly.
  • Added unit tests for ThroughputPoller completed path and is_offer_replace_pending header detection.

Rename replace_throughput to begin_replace_throughput on ContainerClient
and DatabaseClient, returning a new ThroughputPoller type that detects
whether the operation completed synchronously (HTTP 200) or
asynchronously (HTTP 202 / x-ms-offer-replace-pending header).

ThroughputPoller implements IntoFuture for simple .await usage and
Stream for manual polling. It caches the offer RID from the initial
response for efficient single-GET polling and propagates the caller
Context for distributed tracing.

Fixes Azure#3792

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Cosmos The azure_cosmos crate

Projects

Status: Todo

Development

Successfully merging this pull request may close these issues.

Cosmos: make replaceThroughput clearly expose the async state/nature of the API

1 participant