Cosmos PPCB: EPK-range feed-range queries never engage the circuit breaker

## Summary Per-Partition Circuit Breaker (PPCB) never engages for **feed-range queries** that target a physical partition by an **EPK `Range`** (e.g. `SELECT * FROM c` scoped to one partition via `FeedRange::try_from(partition_key_range)`). Because the breaker never trips, the faulted hub region is contacted on **every** query, even within a fast (sub-5-minute) window — failover happens per-operation (via the failover-retry budget) but the partition is never latched away from the hub. ## Root cause PPCB partition-level marking is keyed on `partition_key_range_id`. The chain that drops it for EPK-range feed ranges: 1. **Seeding only handles a logical partition key.** `CosmosDriver::pre_resolve_partition_key_range_id` (`sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs`) returned `None` unless the operation target exposed a logical partition key: ```rust let Some(partition_key) = operation.target().and_then(|t| t.partition_key()) else { return None; }; ``` 2. **EPK-range feed ranges have no logical partition key.** `FeedRange::partition_key()` (`sdk/cosmos/azure_data_cosmos_driver/src/models/feed_range.rs:128`) returns `None` for the `FeedRangeRepr::Range { min_inclusive, max_exclusive }` variant — exactly what a per-physical-partition `SELECT * FROM c` uses. So `pre_resolved_pk_range_id = None`. 3. **The retry state therefore starts with no pk range id.** `operation_pipeline.rs:231`: `retry_state.partition_key_range_id = pre_resolved_pk_range_id;` (= `None`). 4. **Every `MarkPartitionUnavailable` effect is then dropped.** `LocationStateStore` (`sdk/cosmos/azure_data_cosmos_driver/src/driver/routing/location_state_store.rs:265`) skips partition marking when the id is absent: ```rust LocationEffect::MarkPartitionUnavailable(partition) => { if partition.partition_key_range_id.is_none() { // No partition key range ID available (first attempt); skip partition-level marking. continue; } ... } ``` The header-capture fallback (`operation_pipeline.rs:503`) only repopulates `retry_state.partition_key_range_id` for **subsequent attempts within the same operation**. The hub-region failure on attempt 1 is already dropped (id was `None`), and the operation then succeeds via failover (no failure recorded). Each new query starts over with `None`, so the hub-region failure counter never increments → `read_failure_count` never reaches `read_failure_threshold` (default 10) → the breaker never trips → the hub is contacted on every query. ### Why a prior successful request does *not* mask the bug The per-operation `partition_key_range_id` learned from a healthy response (via the header-capture fallback) is **discarded** when the operation completes — it is only ever written to `retry_state`, never to the shared routing-map cache. Only the routing-map cache (`pk_range_cache`, populated by `fetch_pk_ranges_from_service`) persists across operations. So whether or not earlier queries succeeded, every faulting query starts again from `None`, and the breaker still never accumulates failures. This is captured by the second repro test below. ## Repro tests Both run against a multi-region (East US 2 hub + West US 3) account, with 503s injected on the query path in the hub region. Both **fail** on the `!regions.contains(&HUB_REGION)` assert before the fix. ```bash export AZURE_COSMOS_CONNECTION_STRING=$(az cosmosdb keys list --type connection-strings \ -g <rg> -n <multi-region-account> --query "connectionStrings[0].connectionString" -o tsv) RUSTFLAGS='--cfg test_category="multi_region"' cargo test -p azure_data_cosmos_driver \ --test multi_region --features fault_injection \ ppcb_feed_range -- --nocapture --test-threads=1 ``` ### Test 1 — breaker should stop contacting the hub after it trips ```rust #[tokio::test] #[cfg_attr(not(test_category = "multi_region"), ignore = "requires test_category 'multi_region'")] pub async fn ppcb_enabled_503_on_feed_range_query_fails_over_after_threshold( ) -> Result<(), Box<dyn Error>> { let condition = FaultInjectionConditionBuilder::new() .with_operation_type(FaultOperationType::QueryItem) .with_region(HUB_REGION) .build(); let result = FaultInjectionResultBuilder::new() .with_error(FaultInjectionErrorType::ServiceUnavailable) .with_probability(1.0) .build(); let rule = Arc::new( FaultInjectionRuleBuilder::new("ppcb-503-feed-range-query", result) .with_condition(condition) .build(), ); let operation_options = OperationOptionsBuilder::new() .with_per_partition_circuit_breaker_enabled(true) .build(); Box::pin(DriverTestClient::run_with_unique_db_and_fault_injection_options( vec![Arc::clone(&rule)], operation_options, async |context, database| { let container_name = context.unique_container_name(); let container = context.create_container(&database, &container_name, "/pk").await?; for i in 0..5 { let pk = format!("pk{i}"); let item_json = format!(r#"{{"id": "ppcb-query-item-{i}", "pk": "{pk}", "value": "test"}}"#); context.create_item_with_pk(&container, pk, item_json.as_bytes()).await?; } // Scope the query to one physical partition's EPK range. let ranges = context .resolve_all_partition_key_ranges(&container, true) .await? .ok_or("expected partition key ranges for the container")?; let first_range = ranges.first().ok_or("container has no partition key ranges")?; let feed_range = FeedRange::try_from(first_range)?; // Issue enough queries to exceed read_failure_threshold (default 10). let mut query_success_count = 0; let mut last_regions_contacted = Vec::new(); for _ in 0..15 { if let Ok(response) = context .query_feed_range(&container, feed_range.clone(), "SELECT * FROM c") .await { if response.status().is_success() { query_success_count += 1; } last_regions_contacted = response.diagnostics_ref().regions_contacted(); } } assert!(rule.hit_count() > 0, "fault rule should have been hit in the hub region"); assert!(query_success_count > 0, "queries should recover via failover"); // EXPECTED: after the breaker trips the partition is latched to the // healthy region and the hub is no longer contacted. // ACTUAL (without fix): regions_contacted = [eastus2, westus3] on every // query — the breaker never trips because pk_range_id is never seeded. assert!( !last_regions_contacted.contains(&HUB_REGION), "after the circuit breaker trips, feed-range queries should NOT contact the hub \ region, but regions_contacted={last_regions_contacted:?}" ); Ok(()) }, )) .await } ``` ### Test 2 — prior success must not mask the breaker This guards the "what if a successful request happens first?" concern: the fault is injected **only after** a burst of genuinely successful warm-up queries. The breaker must still trip and exclude the hub. ```rust #[tokio::test] #[cfg_attr(not(test_category = "multi_region"), ignore = "requires test_category 'multi_region'")] pub async fn ppcb_feed_range_query_trips_even_after_initial_success() -> Result<(), Box<dyn Error>> { let condition = FaultInjectionConditionBuilder::new() .with_operation_type(FaultOperationType::QueryItem) .with_region(HUB_REGION) .build(); let result = FaultInjectionResultBuilder::new() .with_error(FaultInjectionErrorType::ServiceUnavailable) .with_probability(1.0) .build(); let rule = Arc::new( FaultInjectionRuleBuilder::new("ppcb-503-feed-range-query-after-success", result) .with_condition(condition) .build(), ); // Start with the fault DISABLED so warm-up queries genuinely succeed. rule.disable(); let operation_options = OperationOptionsBuilder::new() .with_per_partition_circuit_breaker_enabled(true) .build(); Box::pin(DriverTestClient::run_with_unique_db_and_fault_injection_options( vec![Arc::clone(&rule)], operation_options, async |context, database| { let container_name = context.unique_container_name(); let container = context.create_container(&database, &container_name, "/pk").await?; for i in 0..5 { let pk = format!("pk{i}"); let item_json = format!(r#"{{"id": "ppcb-after-success-item-{i}", "pk": "{pk}", "value": "test"}}"#); context.create_item_with_pk(&container, pk, item_json.as_bytes()).await?; } let ranges = context .resolve_all_partition_key_ranges(&container, true) .await? .ok_or("expected partition key ranges for the container")?; let first_range = ranges.first().ok_or("container has no partition key ranges")?; let feed_range = FeedRange::try_from(first_range)?; // Warm-up: 3 successful feed-range queries against the healthy hub. let mut warmup_success = 0; for _ in 0..3 { let response = context .query_feed_range(&container, feed_range.clone(), "SELECT * FROM c") .await?; if response.status().is_success() { warmup_success += 1; } } assert!(warmup_success == 3, "warm-up queries should all succeed while fault disabled"); assert!(rule.hit_count() == 0, "fault rule should not be hit while disabled"); // Now enable the fault in the hub region. rule.enable(); let mut query_success_count = 0; let mut last_regions_contacted = Vec::new(); for _ in 0..15 { if let Ok(response) = context .query_feed_range(&container, feed_range.clone(), "SELECT * FROM c") .await { if response.status().is_success() { query_success_count += 1; } last_regions_contacted = response.diagnostics_ref().regions_contacted(); } } assert!(rule.hit_count() > 0, "fault rule should have been hit after enable"); assert!(query_success_count > 0, "queries should recover via failover"); // Even though the partition just served successful traffic, once the // breaker trips the hub must no longer be contacted. // ACTUAL (without fix): regions_contacted = [eastus2, westus3]. assert!( !last_regions_contacted.contains(&HUB_REGION), "after the circuit breaker trips, feed-range queries should NOT contact the hub \ region even following initial success, but regions_contacted={last_regions_contacted:?}" ); Ok(()) }, )) .await } ``` **Expected (both tests):** after ~10 failures the partition latches to West US 3; `regions_contacted` no longer includes East US 2. **Actual (both tests, without fix):** `regions_contacted = [eastus2, westus3]` on every query (verified across consecutive queries); the breaker never trips. ## Proposed fix Seed `partition_key_range_id` for EPK-`Range` feed ranges, not just logical-partition-key targets. `pre_resolve_partition_key_range_id` should resolve the EPK range to its owning physical partition's `partitionKeyRangeId` via the PK-range cache when the target is a `FeedRange::Range`, seeding the id only when the range maps to exactly one physical partition (a cross-partition range has no single owning pk range id): ```rust // Logical-partition-key targets resolve directly from the partition key. if let Some(partition_key) = target.partition_key() { return self .pk_range_cache .resolve_partition_key_range_id(container, partition_key, false, |c, cont| { Box::pin(self.fetch_pk_ranges_from_service(c, cont)) }) .await .map(PartitionKeyRangeId::from); } // EPK-range feed ranges (e.g. `SELECT * FROM c` scoped to a single physical // partition) carry no logical partition key. Resolve the owning physical // partition by EPK range so PPCB/PPAF can attribute failures from the first // attempt. Only seed when the range maps to exactly one physical partition. let ranges = self .pk_range_cache .resolve_overlapping_ranges( container, target.min_inclusive()..target.max_exclusive(), false, |c, cont| Box::pin(self.fetch_pk_ranges_from_service(c, cont)), ) .await?; match ranges.as_slice() { [single] => Some(PartitionKeyRangeId::from(single.id.clone())), _ => None, } ``` With the id seeded from the first attempt, `MarkPartitionUnavailable` is applied and the breaker accumulates failures and trips as designed. ### Verified | Scenario | Without fix | With fix | |---|---|---| | Test 1 (breaker stops contacting hub) | ❌ `[eastus2, westus3]` | ✅ passes | | Test 2 (success-first does not mask) | ❌ `[eastus2, westus3]` | ✅ passes | Both verified live against a multi-master East US 2 + West US 3 account. `cargo fmt` / `cargo clippy --tests` clean.