Azure / Azure/azure-sdk-for-rust

[Cosmos] - Address Known Gaps in Fetching PkRange From Service

Open
#4,301 0 comments 0 reactions 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

# Cosmos DB Rust Driver — Partition Key Range Cache: Behavioral Gaps vs. .NET / Java

The Rust driver's `PartitionKeyRangeCache` (in
[`sdk/cosmos/azure_data_cosmos_driver/src/driver/cache/partition_key_range_cache.rs`](../sdk/cosmos/azure_data_cosmos_driver/src/driver/cache/partition_key_range_cache.rs))
and the per-page fetch in `cosmos_driver::fetch_pk_ranges_from_service` share most of
their behavior with the .NET (`Microsoft.Azure.Cosmos.Routing.PartitionKeyRangeCache`)
and Java (`com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache`) SDKs:

- Cross-region failover on 503 / 410 / 408 / connection failure (via the operation pipeline).
- `A-IM: Incremental feed` + `x-ms-max-item-count: -1` request headers.
- `If-None-Match` continuation token threading; ETag captured from response.
- Multi-page change-feed loop until HTTP 304 Not Modified.
- Incremental merge into the previous routing map via `try_combine`.
- 404 → cache miss; 401 / 403 → terminal.

The gaps below are the remaining areas where the Rust driver does **not** match the
behavior of the .NET and/or Java SDKs. None of them are blocking for steady-state
operation, but they affect tail behavior under throttling, partition splits, and
post-split races.

---

## Gap 1 — 429 throttle backoff is not honored (rotates region instead)

**Severity:** Medium
**Affects:** Tail latency and cross-region budget under global throttle events.

### Current behavior (Rust)

When `/pkranges` returns HTTP 429 the driver's pipeline treats it like any other
transient failure and rotates to the next preferred read endpoint via the
in-flight failover loop, ignoring the server-supplied `Retry-After` header.

### .NET behavior

`PartitionKeyRangeCache.GetRoutingMapForCollectionAsync` wraps each fetch in
`MetadataRequestThrottleRetryPolicy`, which composes a
`ResourceThrottleRetryPolicy`:

- Default `MaxRetryAttemptsOnThrottledRequests = 9`
- Default `MaxRetryWaitTimeInSeconds = 30`
- Honors `x-ms-retry-after-ms` from the response and waits that long
**on the same region** before retrying.

Cross-region failover only happens after the throttle budget is exhausted.

### Java behavior

Same shape: `client.readPartitionKeyRanges(...)` flows through
`RxGatewayStoreModel`, which applies `ResourceThrottleRetryPolicy`
(default 9 attempts / 30s) on 429 with `retry-after` honored on the same region
before any cross-region action.

### Why it matters

Under a global throttle (e.g., regional capacity event), rotating regions on the
first 429 amplifies load on alternate regions and burns the cross-region failover
budget for nothing — the alternate is likely to be throttled too.

### Suggested approach

1. Plumb `x-ms-retry-after-ms` (or HTTP `Retry-After`) out of the response
headers in `fetch_pk_ranges_from_service`.
2. On 429 returned from a single page fetch, sleep for the indicated duration
(capped at a max wait, e.g., 30s total budget) and retry the same page
on the same endpoint, up to N attempts (default 9).
3. Only fall through to the existing cross-region failover when the throttle
budget is exhausted.

---

## Gap 2 — Incomplete routing map is not retried

**Severity:** Medium-High
**Affects:** Correctness of routing immediately after a partition split.

### Current behavior (Rust)

`fetch_and_build_routing_map` ([partition_key_range_cache.rs](../sdk/cosmos/azure_data_cosmos_driver/src/driver/cache/partition_key_range_cache.rs))
calls `ContainerRoutingMap::try_create` (or `try_combine` for incremental).
Both can return `Ok(None)` when the assembled set of ranges does not cover
`[00, FF...)` end-to-end (e.g., the change-feed loop captured a snapshot in
the middle of a split). Today the cache logs a warning and returns an
**empty `ContainerRoutingMap`**, which is then cached.

The next routing lookup will get an empty map → no range → fall back to
default routing. The driver only re-fetches when an explicit `force_refresh`
is requested (typically driven by 410/1002 PartitionKeyRangeGone on a data
operation). Until then, PPAF/PPCB pre-resolution silently no-ops.

### .NET behavior

`CollectionRoutingMap.TryCreateCompleteRoutingMap` returns `null` when the set
is incomplete, and `GetRoutingMapForCollectionAsync` raises
`NotFoundException("... Range information either doesn't exist or is not
complete.")`. Callers (`TryLookupAsync`) treat this as a transient cache miss;
the next lookup retries and is expected to see the post-split state.

### Java behavior

`InCompleteRoutingMapRetryPolicy` is wired around
`getRoutingMapForCollectionAsync` via `ObservableHelper.inlineIfPossible`.
When `tryCreateCompleteRoutingMap` / `tryCombine` returns `null` the policy
re-runs the entire fetch (a small bounded number of times) until the map is
complete or the budget is exhausted, at which point it surfaces a
`NotFoundException`.

### Why it matters

During a partition split, the gateway briefly publishes a transient state
where the parent range is gone but a child range hasn't appeared yet (or
vice versa). The Rust behavior caches an empty map and stays incorrect until
something else triggers `force_refresh`; .NET/Java retry the fetch within the
same call and converge.

### Suggested approach

Add a small bounded retry inside `fetch_and_build_routing_map` (e.g., 3
attempts with short backoff) for the "incomplete map" case (`Ok(None)` from
`try_create` / `try_combine` when there are non-zero ranges in the page set).
Do not cache an empty map on incomplete; cache only a complete map or
return `None` from `try_lookup` so the next call re-fetches.

---

## Gap 3 — `404:1002` (PartitionKeyRangeGone) is not propagated for downstream retries

**Severity:** Low
**Affects:** Cooperation between the routing-map cache and document-level
retry policies after a partition split.

### Current behavior (Rust)

`fetch_pk_ranges_from_service` lumps `404` into a single "permanent error"
branch and returns `None` for any 404, regardless of sub-status code.

### Java behavior

`RxPartitionKeyRangeCache.tryLookupAsync` and `tryGetRangeByPartitionKeyRangeId`
explicitly check `Exceptions.isNotFound(dce) && !Exceptions.isSubStatusCode(
dce, HttpConstants.SubStatusCodes.READ_SESSION_NOT_AVAILABLE)`:

- Plain 404 (no sub-status, or sub-status other than `1002`) → return
`ValueHolder(null)` (cache miss).
- 404:1002 (`READ_SESSION_NOT_AVAILABLE`) → re-throw so document-level
retry policies (`SessionTokenMismatchRetryPolicy`, etc.) can react.

### .NET behavior

.NET does not have an exact equivalent at the cache layer (it has a slightly
different sub-status taxonomy on metadata reads), so this gap is Java-only.

### Why it matters

If the gateway returns `404:1002` for a `/pkranges` lookup, surfacing it as a
plain cache miss prevents the document-level retry policy from doing the right
thing (forcing a session-token refresh / read retry).

### Suggested approach

In `fetch_pk_ranges_from_service`, distinguish:

- `404` with no sub-status, or sub-status that indicates "not found" →
return `None` (current behavior).
- `404:1002` → propagate the error rather than swallowing it.

This requires plumbing the response sub-status from the operation pipeline
into the error returned to `fetch_pk_ranges_from_service`.

---

## Gap 4 — No negative cache TTL on permanent 401 / 403 / 404

**Severity:** Low
**Affects:** Wasted gateway round-trips and noisy error logs when a container
is misconfigured, deleted, or the credential lacks access.

### Current behavior (Rust)

`fetch_pk_ranges_from_service` logs an error and returns `None` on 401 / 403 /
404. The cache layer (`fetch_and_build_routing_map`) then caches an **empty
`ContainerRoutingMap`**, but on the next routing lookup any `force_refresh`
(e.g., triggered by a 410 / 1002 elsewhere, or simply because the empty map
yields no usable range) re-issues the same `/pkranges` request, which fails
identically. There is a `// TODO: Consider adding a negative-cache TTL ...`
comment in `fetch_pk_ranges_from_service` acknowledging this.

### .NET behavior

.NET does not implement a dedicated negative-cache TTL for permanent metadata
errors at the `PartitionKeyRangeCache` layer either. It relies on the
`AsyncCacheNonBlocking` entry created on the failed fetch to keep the failure
"sticky" until something explicitly invalidates it, plus
`enableAsyncCacheExceptionNoSharing` to avoid sharing the exception across
concurrent waiters.

### Java behavior

Same as .NET — no explicit negative TTL at the routing-map cache layer.

### Why it matters

Even though no SDK fully solves this, the Rust driver caches an *empty* routing
map (rather than the failure itself) on 401 / 403 / 404. Subsequent
`force_refresh` calls (which are common in PPAF/PPCB code paths after any
410/1002 anywhere in the system) keep hammering the gateway for a container
that will never resolve. A small negative TTL — even a few seconds — would
suppress the spam without changing semantics.

### Suggested approach

Cache a sentinel "negative entry" with a short TTL (configurable; default
on the order of 5–30 seconds) on 401 / 403 / 404 in `fetch_pk_ranges_from_service`.
While the entry is fresh, `try_lookup` short-circuits and returns the same
error (or `None`) without issuing another HTTP request. After the TTL elapses
the next lookup re-fetches normally. Implement this in the cache layer rather
than in `fetch_pk_ranges_from_service` so it composes cleanly with the
existing `force_refresh` path (force refresh should still bypass the negative
cache).

---

## Out-of-scope (not a real gap, captured for completeness)

- **`MAX_FETCH_ITERATIONS = 10` cap** in the Rust pagination loop. .NET/Java
have no such cap. The cap is a defensive safety net that emits a
`tracing::warn!` if hit; in practice the loop completes in 1–2 iterations.
This is intentional and not a behavioral gap.

---

## Gap 5 — PPCB Needs to Track All 5XX Errors.

## Gap 6 — Create `PartitionFailoverOptions` type of their own as a new type.

**Severity:** Medium

## References

- Rust: [`sdk/cosmos/azure_data_cosmos_driver/src/driver/cache/partition_key_range_cache.rs`](../sdk/cosmos/azure_data_cosmos_driver/src/driver/cache/partition_key_range_cache.rs)
- Rust: [`sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs`](../sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs) — `fetch_pk_ranges_from_service`
- .NET: [`Microsoft.Azure.Cosmos/src/Routing/PartitionKeyRangeCache.cs`](https://github.com/Azure/azure-cosmos-dotnet-v3/blob/master/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyRangeCache.cs)
- .NET: `Microsoft.Azure.Cosmos/src/Routing/MetadataRequestThrottleRetryPolicy.cs`
- Java: [`sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/caches/RxPartitionKeyRangeCache.java`](https://github.com/Azure/azure-sdk-for-java/blob/main/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/caches/RxPartitionKeyRangeCache.java)
- Java: `com.azure.cosmos.implementation.caches.InCompleteRoutingMapRetryPolicy`

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.