Azure / Azure/azure-sdk-for-rust

Cosmos: Percentage-based PPCB partition failback

Open
#4,676 0 comments 1 reaction 1 assignee Claimed by @kundadebdatta View on GitHub
Cosmos
Dominant language
Rust
Stars
884
Forks
365
Avg merge
2d 19h
Merged PRs (30d)
109

Description

## 1. Summary

Replace (or augment) the current **time + jitter + single-probe** PPCB failback model with a **percentage-based (success-rate) failback** model: once a tripped partition's original region demonstrates a configurable success rate — i.e. at least *X%* of a representative sample of requests succeed — the partition is failed back to its original region.

The trip side (how a partition becomes `Unhealthy`) is unchanged. Only the **recovery / failback** decision changes.

## 2. Background — how PPCB failback works today

Per-`(partition_key_range_id, region)` state lives in `PartitionEndpointState` (`azure_data_cosmos_driver/src/driver/routing/partition_endpoint_state.rs`), with each tripped partition represented by a `PartitionFailoverEntry` carrying a `HealthStatus` of either `Unhealthy` or `ProbeCandidate`.

Today's recovery is a three-step state machine:

1. **Trip → `Unhealthy`.** When a partition exceeds `read_failure_threshold` (default 10) / `write_failure_threshold` (default 5) within `counter_reset_window` (default 5 min) — or after `consecutive_hedge_win_threshold` (default 5) alternate-region hedge wins — an `Unhealthy` entry is installed and traffic routes to an alternate region.

2. **`Unhealthy` → `ProbeCandidate` (timer + jitter).** A background sweep (`failback_loop` in `azure_data_cosmos_driver/src/driver/routing/location_state_store.rs`, every `failback_sweep_interval`, default 300 s) calls `expire_partition_overrides` (`azure_data_cosmos_driver/src/driver/routing/routing_systems.rs`). An entry flips to `ProbeCandidate` once:

```text
now - first_failure_time >= partition_unavailability_duration (5s)
+ failback_jitter
```

where `failback_jitter` is sampled uniformly from `[0, partition_unavailability_duration / 2]` (thundering-herd mitigation). **This jitter timer is the "failback trigger" the task is changing.**

3. **`ProbeCandidate` → recovered (single probe).** `resolve_endpoint` (`operation_pipeline.rs`) routes the **next single request** for that partition back to `first_failed_endpoint`. If it **succeeds**, `remove_probe_succeeded_entry` deletes the override (full failback). If it **fails**, the entry returns to `Unhealthy` with re-sampled jitter.

### Limitations of the current model

- Failback hinges on **one** probe request. A single lucky success fully fails back even if the original region is still flaky; a single unlucky failure resets the whole window even if the region is mostly healthy.
- The trigger is **purely time-based** (`unavailability_duration + jitter`), carrying no signal about whether the original region has actually recovered.
- No notion of "partially healthy" — recovery is all-or-nothing on a coin-flip sample of size 1.

## 3. Goal

Make failback **evidence-based**: fail a partition back to its original region only when a statistically meaningful sample of requests to that region succeeds at or above a configured percentage.

Concretely, when a tripped partition becomes eligible for recovery:

- Route a **sample** of its requests to the original region (a canary), instead of a single probe.
- Track a **rolling success ratio** of those canary requests.
- **Fail back** (remove the override) once the sample size is large enough **and** the success ratio is `>= failback_success_threshold` (e.g. 90%).
- **Re-trip** (back to `Unhealthy`, canary abandoned) if the ratio falls below a floor, so a still-degraded region does not flap back and forth.

## 4. Non-goals

- Changing the **trip** thresholds or the failure-counting logic.
- Changing PPAF (`failover_overrides`) — PPAF does not use probe-based failback.
- Changing account-level endpoint failback (the `endpoint_probe_loop` connectivity-probe path is separate and stays as-is).

## 5. Proposed design

### 5.1 New health state: a half-open / canary phase

Extend the `HealthStatus` state machine (`partition_endpoint_state.rs`) so recovery has an explicit measurement phase rather than a one-shot probe. Proposed states:

```text
Unhealthy ──(dwell time elapsed)──▶ Recovering ──(success% >= threshold over N)──▶ removed (failed back)
▲ │
└────────────(success% < floor)──────┘
```

- `Unhealthy` — all traffic to the alternate region (unchanged).
- `Recovering` (replaces / generalizes `ProbeCandidate`) — a configurable **percentage of requests** for this partition is routed to the original region as canaries; the rest continue to the alternate. The entry accumulates a rolling success/failure tally of the canary requests.
- Removal (full failback) happens when the canary success ratio clears the threshold over a minimum sample.

> Keeping `ProbeCandidate` as an alias of "Recovering with canary%=one-shot" is an option to minimize churn; see Open Question 9.4.

### 5.2 Per-entry measurement state

Add to `PartitionFailoverEntry` (`partition_endpoint_state.rs`):

```rust
/// Canary requests routed to the original region during `Recovering`.
recovery_canary_total: u32,
/// Of those, how many succeeded.
recovery_canary_success: u32,
/// When the recovery (canary) phase began — bounds the measurement window.
recovery_started_at: Instant,
```

Use **saturating** arithmetic on the counters. Because `PartitionEndpointState` is immutable-and-CAS-swapped, each canary result produces a new entry via the existing `apply_partition` path (same pattern as `record_partition_success` / `record_partition_failure`).

### 5.3 New configuration knobs

Add to `PartitionFailoverOptions` (`azure_data_cosmos_driver/src/options/partition_failover.rs`), each with an `AZURE_COSMOS_PPCB_*` env binding to match the existing namespace:

| Field | Purpose | Proposed default | Env |
|---|---|---|---|
| `failback_success_threshold` | Min success ratio (0.0–1.0) of canary requests to fail back | `0.90` | `AZURE_COSMOS_PPCB_FAILBACK_SUCCESS_THRESHOLD` |
| `failback_min_sample` | Min number of canary requests before the ratio is evaluated | `10` | `AZURE_COSMOS_PPCB_FAILBACK_MIN_SAMPLE` |
| `failback_canary_percentage` | % of partition requests routed to the original region while `Recovering` | `10` | `AZURE_COSMOS_PPCB_FAILBACK_CANARY_PERCENTAGE` |
| `failback_eval_window` | Time window over which the canary sample is counted (reset on expiry) | `60s` | `AZURE_COSMOS_PPCB_FAILBACK_EVAL_WINDOW_MS` |

Validate ranges (ratio in `[0,1]`, percentage in `1..=100`, min sample `>= 1`), following the existing builder's min-value validation style.

Retain `partition_unavailability_duration` as the **dwell floor** before the `Recovering` phase begins (so we don't canary a region that just failed milliseconds ago), and keep `failback_sweep_interval` for the background sweep that promotes `Unhealthy → Recovering`.

### 5.4 Routing change (`resolve_endpoint`)

In the PPCB branch of `resolve_endpoint` (`operation_pipeline.rs`, the block that currently special-cases `HealthStatus::ProbeCandidate`):

- For a `Recovering` entry, route this request to `first_failed_endpoint` with probability `failback_canary_percentage`% (deterministic sampling — e.g. a per-entry counter or hashing the request, not a global RNG, to keep behavior testable); otherwise route to `current_endpoint` (the alternate).
- Preserve the existing `ppcb_should_skip` / excluded-region guard so canaries never target an excluded or in-flight-failed region.

### 5.5 Result feedback

Where probe results are consumed today (the `remove_probe_succeeded_entry` call site and the failure path that re-trips a probe), record the canary outcome into the new counters instead:

- **Canary success:** increment `recovery_canary_success` and `recovery_canary_total`. If `recovery_canary_total >= failback_min_sample` **and** `success/total >= failback_success_threshold`, remove the override (full failback) via the existing `remove_probe_succeeded_entry` swap.
- **Canary failure:** increment `recovery_canary_total` only. If the running ratio has fallen below a floor once the min sample is met, transition back to `Unhealthy` and reset the counters (re-sample `failback_jitter`).
- **Window expiry:** if `now - recovery_started_at > failback_eval_window` without a decision, reset the counters (and optionally re-enter the dwell), so a low-traffic partition does not fail back on a stale sample.

Only **canary** requests (those routed to the original region) feed the ratio; requests served by the alternate region are not counted.

### 5.6 Background sweep change (`expire_partition_overrides`)

`expire_partition_overrides` changes its transition target from `ProbeCandidate` to `Recovering` (initializing the canary counters and `recovery_started_at`). The dwell-time + jitter gate for *entering* recovery is unchanged; the jitter no longer **decides** failback — it only staggers when canary routing begins.

## 6. Safety / correctness considerations

1. **Low-traffic partitions.** With few requests, the min-sample bar may never be met. Mitigation: the `failback_eval_window` reset, plus an optional fallback to a single-probe decision after a max number of windows (Open Question 9.2).
2. **Flapping.** A region oscillating around the threshold could trip/recover repeatedly. Mitigation: hysteresis — separate `failback_success_threshold` (to recover) from a lower re-trip floor, and require the full min sample before either decision.
3. **Canary blast radius.** Canary traffic hits a possibly-still-unhealthy region. `failback_canary_percentage` bounds the customer-visible impact; keep the default low (10%) and ensure canary failures still get normal failover retries so the *caller* is not penalized by a canary loss.
4. **Determinism for tests.** Use a deterministic sampler (counter/hash), not a process RNG, so integration tests can assert exact canary counts.
5. **Concurrency.** Counter updates ride the existing CAS (`apply_partition`) path; ensure read-modify-write of the ratio is done inside the CAS closure to avoid lost updates under concurrent canary completions.
6. **Memory.** Three small fields per entry; no unbounded growth (entries are already bounded by the partition/region cardinality PPCB tracks).
7. **Backward compatibility.** Default thresholds should approximate today's behavior closely enough that enabling this does not surprise existing deployments — or gate the new model behind a flag (Open Question 9.1).

## 7. Implementation checklist

- [ ] `partition_failover.rs`: add the four new knobs + env bindings + builder methods + range validation + getters; update `Default`.
- [ ] `partition_endpoint_state.rs`: add `Recovering` to `HealthStatus` (or repurpose `ProbeCandidate`); add the three canary counter fields to `PartitionFailoverEntry`; saturating arithmetic.
- [ ] `routing_systems.rs`: update `expire_partition_overrides` to transition to `Recovering` and initialize counters; add the canary-result reducer (success/failure/window-expiry → new state); keep `remove_probe_succeeded_entry` as the "failed back" terminal.
- [ ] `operation_pipeline.rs`: update `resolve_endpoint` PPCB branch to do percentage-based canary routing; update the result-feedback site to record canary outcomes instead of one-shot probe removal.
- [ ] `location_state_store.rs`: `failback_loop` unchanged structurally (still drives the sweep); confirm it threads the new config.
- [ ] Diagnostics/tracing: emit the canary tally + ratio on transitions so SREs can observe failback progress.
- [ ] CHANGELOG entry for `azure_data_cosmos_driver` (and `azure_data_cosmos` if any surface changes).

## 8. Testing plan

**Unit (`routing_systems.rs` tests):**

- Canary success reducer fails back only at `>= min_sample` **and** `>= success_threshold` (boundary cases: exactly at threshold, one below).
- Canary failure reducer re-trips below the floor; does not re-trip before min sample.
- Window expiry resets counters without failing back.
- `expire_partition_overrides` transitions `Unhealthy → Recovering` after dwell + jitter and initializes counters; does not re-transition an already-`Recovering` entry.
- Deterministic sampler routes exactly `canary_percentage`% over a known request sequence.

**Env / options (`partition_failover.rs` tests):**

- Each new env var parses, validates ranges, and layers correctly; `Default` values match the table in §5.3.

**Integration (multi-region / multi-write fault injection):**

- Trip a partition; inject a region that recovers to ~95% success; assert the partition fails back after the canary sample clears 90%, and that it did **not** fail back on the first success alone.
- Inject a region stuck at ~50% success; assert the partition does **not** fail back and re-trips.
- Low-traffic partition: assert window-expiry behavior.
- Extend the existing `ppcb_failback_to_hub_region_after_*` tests (`tests/multi_region_tests/driver_partition_failover.rs`, `tests/multi_write_tests/driver_partition_failover.rs`) for the new model.

## 9. Open questions

1. **Replace vs. flag.** Replace the probe model outright, or gate percentage-based failback behind a new opt-in (e.g. `failback_strategy = Probe | SuccessRate`) so rollout is staged?
2. **Low-traffic fallback.** After K eval windows without enough samples, should the driver fall back to the legacy single-probe decision, or keep waiting?
3. **What counts in "total"?** Confirm the denominator is **canary requests routed to the original region** (proposed), not "all requests for the partition." The latter would conflate alternate-region health with original-region recovery.
4. **State reuse.** Add a distinct `Recovering` variant, or overload `ProbeCandidate` with canary counters (smaller diff, fuzzier semantics)?
5. **Re-trip threshold.** Single threshold with hysteresis margin, or two explicit knobs (`failback_success_threshold` to recover, `re_trip_failure_threshold` to abandon)?
6. **.NET / Java parity.** Do the sibling SDKs use a success-rate failback we should match (thresholds, window, canary %), or is this Rust-specific?
7. **Interaction with hedging-driven trips.** A partition tripped purely by `consecutive_hedge_win_threshold` has no hard-failure history — does canary success on the original region mean the same thing there?

## 10. Acceptance criteria

- A tripped partition fails back to its original region only after a configurable percentage of a minimum-size canary sample succeeds — never on a single lucky request.
- A still-degraded region does **not** fail back and re-trips cleanly without flapping.
- New knobs resolve through the standard `AZURE_COSMOS_PPCB_*` env + builder layering with range validation.
- Canary routing percentage bounds the customer-visible blast radius; canary losses still get normal failover so callers are not penalized.
- Unit + integration coverage for success, sustained-failure, boundary, and low-traffic cases; CHANGELOG updated.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.