perf(cosmos): improve pkrange cache memory usage

[tvaron3]

Summary

Four optimizations to drastically reduce CollectionRoutingMap memory footprint, motivated by a customer regression where PPCB (Per Partition Circuit Breaker) being silently enabled by default in 4.15.0 caused a ~730 MB heap. Each CosmosClient
previously kept its own routing-map cache of every partition key range; for accounts with many partitions × many clients this dominates memory.

Changes

1. Shared routing-map cache across clients to the same endpoint

_routing/routing_map_provider.py, _routing/aio/routing_map_provider.py

Module-level _shared_routing_map_cache: dict keyed by client.url_connection. All PartitionKeyRangeCache instances pointing at the same endpoint share a single dict. clear_cache() uses dict.clear() (not reassignment) so every sharing client
continues to see the same dict.

Dominant win — independent of per-range size, the cache is now O(endpoints) rather than O(clients × endpoints).

2. Strip unused fields → compact PKRange namedtuple

_routing/routing_range.py, _routing/routing_map_provider.py, _routing/aio/routing_map_provider.py

The service returns 13 fields per partition key range; CollectionRoutingMap only uses 4 (id, minInclusive, maxExclusive, parents). After fetching, ranges are converted into a PKRange namedtuple that supports dict-style [key], .get(),
in, .items(), and equality with plain dicts for full backward compatibility.

Dropped fields: _rid, _etag, ridPrefix, _self, ownedArchivalPKRangeIds, _ts, lsn.

3. __slots__ on Range and _PartitionHealthInfo

_routing/routing_range.py, _global_partition_endpoint_manager_circuit_breaker_core.py

Eliminates the per-instance __dict__ (~100 bytes per object). Material in PPCB scenarios with many partitions × many clients.

4. Skip redundant .upper() on hex strings

_routing/routing_range.py

Range.__init__ previously called .upper() unconditionally on min/max. Cosmos returns uppercase hex (e.g. 10F0F0F0...); we skip the copy when already uppercase.

Memory Profiling Results

Test setup (real account, measured with tracemalloc, retained memory):

• ~100 physical partitions, 2 regions (East US 2 + West US 3), multi-write
• VM: Standard_D16s_v5
• Operations per client: 1 read_item + 1 upsert_item
• PPCB: AZURE_COSMOS_ENABLE_CIRCUIT_BREAKER=True

Total cache memory (MB) — patches 2, 3, 4 measured

Clients	4.15.0 Original	Strip only (#2)	Strip + slots + upper (#2 + #3 + #4)	PPCB=false baseline
1	14.3	14.3	14.3	14.0
25	23.0	20.5	20.0	17.9
50	31.9	27.4	25.8	21.7
100	44.9	39.9	36.6	29.4
150	63.8	52.9	43.3	36.4

PPCB overhead reduction (vs PPCB=false), patches 2-4

Clients	Original PPCB overhead	After patches #2 + #3 + #4	Reduction
25	5.1 MB	2.1 MB	-58%
50	10.3 MB	4.1 MB	-60%
100	15.4 MB	7.2 MB	-53%
150	27.4 MB	6.9 MB	-74%

Shared-cache impact (#1) — separately measured

The shared cache makes routing-map memory independent of client count for clients sharing an endpoint. Measured with tracemalloc (real allocated bytes incl. unique string contents) on Python 3.13, 100K PK ranges per cache:

Per-cache cost — isolates the per-record container choice (#2-#4 only, no #1):

Representation	MiB / 100K ranges	Bytes / entry
Pre-PR: dict PK-range record + Range without __slots__ + always .upper()	81.5 MiB	854 B
After #2-#4: PKRange namedtuple + Range.__slots__ + skip-redundant-.upper()	43.1 MiB	452 B
Reduction	−38.3 MiB / cache (−47%)	−402 B / entry

Cluster-wide cost — adds the shared-cache effect (#1):

Scenario (150 clients on the same endpoint, 100K PK ranges)	Total routing-map memory
Pre-PR (no patches, each client owns its own cache)	~12.2 GB (81.5 MiB × 150)
With #2-#4 only (smaller per-cache, still per-client)	~6.3 GB (43.1 MiB × 150)
With all 4 (single shared cache for the endpoint)	~43 MiB total

The shared cache flattens the cost from O(clients × ranges) to O(ranges), which dominates at high client counts. Combined with patches #2-#4, PPCB-enabled customers with many clients per endpoint should see PPCB cache overhead drop close to zero.

Note on a previously-cited number: an earlier draft reported "44 MB → 7 MB / −84%" per 100K entries. That figure used sys.getsizeof() on bare containers (which excludes string contents and counts the dict's empty-list overhead but not PKRange's tuple overhead) and is not representative of the real allocated footprint. The numbers above use tracemalloc peak-allocation deltas, which is the honest measurement.

Test Coverage

New tests added on top of existing routing/cache coverage:

• tests/routing/test_shared_pk_range_cache.py (sync) + test_shared_pk_range_cache_async.py (async) — cache sharing, isolation across endpoints, clear_cache() semantics, dict-identity preservation.
• tests/test_shared_cache_integration.py (sync) + test_shared_cache_integration_async.py (async) — emulator integration covering multi-client cache sharing, clear_cache() repopulation, PKRange compatibility through full CRUD lifecycles, change
  feed, and endpoint isolation.
• tests/test_shared_cache_fault_injection.py (sync) + test_shared_cache_fault_injection_async.py (async) — fault-injection coverage for the shared cache.
• tests/test_routing_map.py — PKRange dict-equality, __getitem__, .get, __contains__, .items().
• tests/routing/test_collection_routing_map.py, test_routing_map_provider.py (+ async), test_routing_map_provider_unit.py (+ async) — substantial expansion of unit coverage for refresh / merge / overlap behavior.
• Sync and async paths have full parity coverage.

All routing, session, query, and circuit-breaker tests pass on the latest pipeline run.

Reproduction Script

import asyncio, os, tracemalloc
tracemalloc.start()
os.environ["AZURE_COSMOS_ENABLE_CIRCUIT_BREAKER"] = "True"

from azure.cosmos.aio import CosmosClient

N = int(os.environ.get("NUM_CLIENTS", "150"))

async def main():
    clients = []
    for i in range(N):
        c = CosmosClient(os.environ["COSMOS_URI"], os.environ["COSMOS_KEY"],
                         preferred_locations=["East US 2"])
        db = c.get_database_client("mydb")
        cont = db.get_container_client("mycont")
        try:
            await cont.read_item("x", partition_key="x")
        except Exception:
            pass
        clients.append(c)
    curr, peak = tracemalloc.get_traced_memory()
    print(f"{N} clients: {curr/1024/1024:.1f} MB current, {peak/1024/1024:.1f} MB peak")

asyncio.run(main())

[kushagraThapar]

Thanks @tvaron3
I am curious, are there other places where we build the collection routing map? Shall we fix those as well?

[tvaron3]

Superseded by shared cache approach.

[tvaron3]

/azp run python - cosmos - tests

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[xinlian12]

✅ Review complete (36:43)

Posted 5 inline comment(s).

_{Steps: ✓ context, correctness, cross-sdk, design, history, past-prs, synthesis, test-coverage}

[tvaron3]

/azp run python - cosmos - tests

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[tvaron3]

/azp run python - cosmos - tests

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[xinlian12]

✅ Review complete (44:14)

Posted 2 inline comment(s).

_{Steps: ✓ context, correctness, cross-sdk, design, history, past-prs, synthesis, test-coverage}

[tvaron3]

/azp run python - cosmos - tests

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[kushagraThapar]

LGTM, thanks @tvaron3

[simorenoh]

_resolve_endpoint() is identical in sync/async so can be moved to routing_common and the PkRange construction is also duplicated in collection_routing_map and in routing_common - won't block on that though, thanks Tomas!

[tvaron3]

/azp run python - cosmos - tests

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).