Azure / Azure/azure-sdk-for-rust

[Cosmos] Extend hedge concurrency budget to data-plane hedging

Open
#4,916 0 comments 1 reaction 0 assignees View on GitHub
Dominant language
Rust
Stars
884
Forks
365
Avg merge
2d 19h
Merged PRs (30d)
109

Description

## Context

Split out of #4914 during review of #4896 ([thread](https://github.com/Azure/azure-sdk-for-rust/pull/4896#discussion_r3678684849)).

#4896 added `HedgeBudget` (`src/driver/pipeline/hedge_budget.rs`), a per-client cap on how many cross-region hedge races may be open at once. It is deliberately scoped to **metadata** hedging only: `HedgeBudget::try_admit(PipelineType::DataPlane)` returns an unbudgeted no-op permit. This issue tracks doing the same for the data plane.

## Why it wasn't done in #4896

The metadata budget counts hedge **races**, not spawned **legs**, because admission happens at the two points where the pipeline *decides* to hedge:

- **STAGE 2b** — the first-attempt pre-emptive race.
- **STAGE 5b** — `maybe_upgrade_to_hedge`, on `FailoverRetry` / `SessionRetry`.

Both already have a natural refusal path: decline the upgrade and fall through to the ordinary sequential failover path. That shape existed for the failover budget, so admission slotted in without touching `execute_hedged`.

For metadata, races-are-legs is a tight bound in practice: races are gated by distinct container / PK-range cache misses, and the caches coalesce the rest. For the data plane it is the wrong unit. Point-read concurrency scales with application throughput, so a race-scoped cap would refuse hedges to clients that were never going to spawn an alternate leg at all — a hedge only spawns its secondary if the primary is still outstanding when the threshold fires, which under normal conditions is almost never.

The obvious fix — acquire the permit inside `execute_hedged` where the secondary is actually launched — accounts correctly but has nowhere to put a refusal. `HedgedRaceResult` has exactly two variants, `Terminal` and `BothTransient`, and `BothTransient` contractually means "the race consumed two regions": the caller increments `failover_retry_count` by 2 and advances the `LocationIndex` past both. Returning it when only the primary was ever tried would skip a healthy region and burn double the failover budget. So a leg-scoped budget needs a third outcome ("the primary lost the race but the secondary was never launched, retry from the next region") threaded through `execute_hedged`'s result handling — a real change to the hedge state machine rather than an admission check bolted onto the edges.

## Suggested direction

1. Add a `HedgedRaceResult` variant for *"secondary declined, one region consumed"*, and have the STAGE 2b/5b fallback advance the `LocationIndex` by 1 rather than 2 for it.
2. Move data-plane admission into `execute_hedged`, at the point the secondary leg is spawned, so the counter tracks in-flight alternate-region requests.
3. Give the data plane its own limit and env override, separate from `AZURE_COSMOS_MAX_CONCURRENT_METADATA_HEDGES` — the fan-out characteristics are different enough that one shared number would be wrong for both. `HedgeSlots` is already a standalone type, so this is a second field on `HedgeBudget`, not a new mechanism.
4. Keep admission non-blocking, matching the metadata behavior: a hedge that queues for a slot has already lost the latency argument it exists to win.

## Measuring it

The fault-injection harness added in #4896 is the natural place to prove it, though it currently only exercises metadata:

```text
RUN_HEDGE_STRESS=1 cargo test -p azure_data_cosmos_driver \
--features fault_injection,__internal_in_memory_emulator \
--test in_memory_emulator metadata_hedging_stress -- --ignored --nocapture
```

`scenario_budgeted_storm` is the template — it fixes the hedge ceiling below the fan-out and asserts the resulting amplification drop (2.00× → 1.25× at 32 concurrent reads with a ceiling of 8). A data-plane equivalent would inject a `Document` read delay instead of a metadata one and assert on `HedgeDiagnostics` from the point reads directly, which is simpler than the metadata case since data-plane operations surface their own diagnostics.

Contributor guide

Open the contributing guide

Research direction

Start with src/driver/pipeline/hedge_budget.rs and the execute_hedged entry point, then inspect the STAGE 2b/5b result handling and scenario_budgeted_storm. Run the documented metadata stress command as a baseline. Done means data-plane admission is non-blocking, has its own limit and environment override, preserves one-region fallback semantics, and a Document-delay stress test verifies HedgeDiagnostics and reduced amplification.

Written by the indexing model from the issue text.

Assessment

Tech stack
azure, rust
Domain
backend-api-design, databases
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.