Azure / Azure/azure-sdk-for-python
[Cosmos] HTTP 431 on cross-partition feed_range queries against large session-consistency containers — compound session token grows unbounded
- Dominant language
- Python
- Stars
- 5.6k
- Forks
- 3.4k
- Avg merge
- 1d 21h
- Merged PRs (30d)
- 193
Description
## Summary
`Session.get_session_token_async` in `azure/cosmos/_session.py` builds an unbounded compound session token (`",".join("{pkrange_id}:{vector_clock}" for every cached pk range)`) for cross-partition streamable queries that have neither an explicit `partition_key_range_id` nor a `partitionKey` value. On large containers (thousands of physical partitions), this header grows past the Cosmos gateway frontend's HTTP header-size cap and the gateway responds with **HTTP 431 (Request Header Fields Too Large)** with an empty body, killing the request.
The behavior is acknowledged in a literal in-code TODO:
https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/cosmos/azure-cosmos/azure/cosmos/_session.py#L226-L234
```python
else:
# we're executing a cross partition streamable query that can be resolved by the gateway
# send the entire compound session token for the container to target all partitions
# TODO: this logic breaks large containers, needs to be addressed along with requesting
# a query plan for every query
session_token_list = []
for key in token_dict.keys():
session_token_list.append("{0}:{1}".format(key, token_dict[key].convert_to_string()))
session_token = ",".join(session_token_list)
```
## Why we only see this now
Three conditions are required for the symptom to surface:
1. **Session consistency account** (so `self.session` is populated and a SessionToken header is attached at all).
2. **Large container** — many physical partitions whose vector-clock tokens accumulate in `rid_to_session_token`.
3. **Cross-partition streamable query path** — caller passes neither `partition_key` nor `partition_key_range_id`, so `get_session_token_async` lands in the compound branch.
PPCB (`AZURE_COSMOS_ENABLE_CIRCUIT_BREAKER=true`) is an **amplifier**, not a cause: it drives more cross-region retries and PK-range health probes, populating the session cache across more partitions faster, so the compound token crosses the gateway limit within minutes.
Point operations (`read_item`, `upsert_item`, single-PK `query_items`) are immune — they pass an explicit `partition_key`, so `get_session_token_async` takes the `pk_value is not None` branch (line 212) and returns a single-entry `"{rid}:{vc}"` token (~50 bytes).
## Reproduction
Reproduced against a Cosmos DB account in production with a 16,384-physical-partition session-consistency container (`/id` v2 hash). Reproduced both at scale (production workload) and in isolation (instrumented probe below).
### A. Production workload (reproducible failure)
Per-process Cosmos client config:
- `AZURE_COSMOS_ENABLE_CIRCUIT_BREAKER=true`
- session consistency
- `query_items(query="SELECT * FROM c", feed_range=fr)` looped over all 16,384 feed_ranges with `asyncio.Semaphore(25)`
- 20 processes × 5 concurrent feed-range queries per process
**Result:** ~3,000 HTTP 431 errors per process per 5-minute window. Captured `error_message = "Status code: 431\nb''"` — empty body confirms gateway-level reject.
### B. Three-run experiment (toggle-isolation)
Same workload, account, container, VM, and SDK build. Only the two env vars toggled:
| Run | PPCB | Client `__aenter__()` called | 431 errors |
|---|---|---|---|
| 1 | on | yes (default) | **~3000/proc/5min** |
| 2 | off | yes | 0 |
| 3 | on | no (`WORKLOAD_SKIP_CLOSE=true` skips it) | 0 |
Run 3 is the diagnostic clincher: skipping `__aenter__()` skips `_setup()`, which means `self.session` is never instantiated and no SessionToken header is ever attached → the bug cannot fire. (This is a diagnostic confirmation only — disabling session consistency is not a real fix.)
### C. Isolated probe — compound token growth measurement
Monkey-patch `Session.get_session_token_async` to record token size; run 25 concurrent `query_items(..., feed_range=fr)` calls over a 16,384-partition session container with PPCB enabled.
```python
import asyncio, os, time
os.environ["AZURE_COSMOS_ENABLE_CIRCUIT_BREAKER"] = "true"
from azure.cosmos.aio import CosmosClient
from azure.cosmos import _session
from azure.identity.aio import DefaultAzureCredential
STATS = {"compound_calls": 0, "max_token_bytes": 0, "max_entries": 0,
">4KB": 0, ">8KB": 0}
_orig = _session.Session.get_session_token_async
async def patched(self, resource_path, pk_value, container_properties_cache,
routing_map_provider, partition_key_range_id, options):
token = await _orig(self, resource_path, pk_value, container_properties_cache,
routing_map_provider, partition_key_range_id, options)
if token and "," in token:
n = len(token)
STATS["compound_calls"] += 1
if n > STATS["max_token_bytes"]:
STATS["max_token_bytes"] = n
STATS["max_entries"] = token.count(",") + 1
if n > 4096: STATS[">4KB"] += 1
if n > 8192: STATS[">8KB"] += 1
return token
_session.Session.get_session_token_async = patched
async def worker(cont, frs, start):
idx = start
while True:
try:
async for _ in cont.query_items("SELECT TOP 1 c.id FROM c",
feed_range=frs[idx % len(frs)],
max_item_count=1):
pass
except Exception:
pass
idx += 1
async def main():
cred = DefaultAzureCredential()
async with CosmosClient(URI, credential=cred) as client:
cont = client.get_database_client(DB).get_container_client(CONTAINER)
frs = [fr async for fr in cont.read_feed_ranges()]
tasks = [asyncio.create_task(worker(cont, frs, i * (len(frs) // 25))) for i in range(25)]
await asyncio.sleep(240)
for t in tasks: t.cancel()
await cred.close()
print(STATS)
asyncio.run(main())
```
**Observed growth over 4 minutes (against a 16,384-partition session container):**
| t (s) | compound_calls | max_token_bytes | max_entries | >4 KB | >8 KB |
|---|---|---|---|---|---|
| 31 | 520 | 8,856 | 521 | 280 | 40 |
| 61 | 1,217 | 20,705 | 1,218 | 977 | 737 |
| 91 | 1,819 | 30,939 | 1,820 | 1,579 | 1,339|
| 121 | 2,441 | 41,513 | 2,442 | 2,201 | 1,961|
| 152 | 3,061 | 52,053 | 3,062 | 2,821 | 2,581|
| 182 | 3,581 | 60,893 | 3,582 | 3,341 | 3,101|
| 213 | 4,123 | 70,107 | 4,124 | 3,883 | 3,643|
| 243 | 4,556 | **77,468** | **4,557** | 4,316 | 4,076 |
Compound token exceeded **8 KB on 4,076 / 4,556 calls** (89%) in this short run, and was still growing linearly when stopped. A typical edge / gateway / proxy header cap is 8–16 KB, so the per-request payload is well above any reasonable limit.
## Why the per-iteration mitigation at `_cosmos_client_connection_async.py:3162` isn't enough
The feed_range query path already tries to mitigate this: in the loop at https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py#L3140-L3187 it calls `set_session_token_header_async` per `over_lapping_range["id"]` and overwrites `req_headers[SessionToken]` with the single-partition token before each POST. The comment on L3161 says:
> `# set the session token for this specific partition to avoid sending compound token for all partitions`
But the **initial** call at L3123 (with `partition_key_range_id=None`) still triggers the compound branch and stores the bulk token on `req_headers`. Whenever a sibling code path or a retry re-enters with the bulk header already set (without the per-iteration overwrite), the giant token reaches the wire and the gateway 431s.
## Suggested fixes
**Option 1 (surgical, fixes feed_range path):** drop the initial bulk `set_session_token_header_async` call at https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py#L3122-L3124. The per-iteration call at L3162 with the explicit `over_lapping_range["id"]` already produces the correct single-partition token before each POST.
**Option 2 (defensive, fixes all cross-partition paths):** in `_session.py:226-234`, cap the compound token — either truncate above a configurable max entries / max bytes, or return `""` past a threshold (the gateway will resolve session consistency itself). This is what the existing TODO comment hints at.
I'd recommend doing both: option 1 closes the specific call site; option 2 prevents recurrence on any other path that hits the compound branch.
## Environment
- `azure-cosmos` (Python) from `main` (commit at time of repro: post-#41588 PPCB)
- Python 3.10, async client
- Ubuntu 22.04 on `Standard_D16s_v5`
- Direct gateway connection (no Envoy / sidecar proxy)
- Cosmos DB session-consistency multi-region account, 16,384 physical partitions in target container
## Related code
- TODO: https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/cosmos/azure-cosmos/azure/cosmos/_session.py#L226-L234
- Sync sibling: https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/cosmos/azure-cosmos/azure/cosmos/_session.py#L130-L141
- Feed_range query call site (async): https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py#L3122-L3187
- Per-iteration mitigation comment: https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py#L3161
Contributor guide
Assessment
This issue has not been assessed yet.