[NO REVIEW][Cosmos] Share PartitionKeyRangeCache across CosmosClients targeting the same account

[xinlian12]

Description

Today every CosmosClient / CosmosAsyncClient owns its own RxPartitionKeyRangeCache, even when many clients in the same JVM target the same Cosmos account (a common pattern for multi-tenant / multi-credential apps and frameworks that recreate clients). The routing-map data is duplicated N times and /pkranges calls fan out N times for the same containers.

This PR moves the routing-map storage to a process-wide, refcounted registry keyed by service endpoint. The fetching path (which depends on the per-client network stack, auth, collection cache, diagnostics) stays per-client.

Design

Split RxPartitionKeyRangeCache into two layers:

1. Storage — AsyncCacheNonBlocking<String, CollectionRoutingMap>. Account-level data, naturally shareable. Now obtained from SharedRoutingMapCacheRegistry (process-wide singleton) keyed by the service endpoint URL string.
2. Fetcher — issues /pkranges, depends on per-client RxDocumentClientImpl, RxCollectionCache, diagnostics. Unchanged.

Concurrency model

All registry state transitions go through ConcurrentHashMap.compute(...), which provides atomic per-key check-and-update.

Lifecycle

• RxPartitionKeyRangeCache ctor acquires from the registry (bumps refcount).
• RxPartitionKeyRangeCache now implements Closeable; close() releases the refcount and is idempotent (guarded by AtomicBoolean).
• RxDocumentClientImpl.close() calls LifeCycleUtils.closeQuietly(partitionKeyRangeCache).
• When the last reference is released, the registry entry is evicted so idle endpoints don't pin memory.

Why this is safe

• Data sensitivity: PK-range data is account-level metadata; not per-credential, not per-consistency-level, not per-preferred-region.
• Single-flight: AsyncCacheNonBlocking already enforces a single in-flight fetch per key (AsyncLazyWithRefresh). Sharing the instance strengthens that to "one in-flight /pkranges per (account, container) across all clients" — a stronger guarantee than today.
• Endpoint identity: Two clients only collide when serviceEndpoint.toString() is equal, which means they target the same account.
• Back-compat: The two-arg constructor still exists; it resolves the endpoint from the client. Existing tests that mock RxDocumentClientImpl (mock returns null endpoint) silently fall back to an isolated cache — identical behaviour to today.

Opt-out

System property COSMOS.SHARED_PARTITION_KEY_RANGE_CACHE_ENABLED=false restores per-client private caches.

Reference

The Python SDK uses essentially this pattern (sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_map_provider.py) — module-level endpoint-keyed dicts with refcounted cleanup. This PR adapts that to Java idioms: ConcurrentHashMap.compute instead of explicit RLock, Closeable instead of __del__.

Files

File	Change
caches/SharedRoutingMapCacheRegistry.java	NEW — process-wide registry singleton
caches/RxPartitionKeyRangeCache.java	Endpoint ctor, registry-backed storage, idempotent close()
Configs.java	New system property
RxDocumentClientImpl.java	Pass endpoint to ctor, release cache on close
caches/SharedRoutingMapCacheRegistryTest.java	NEW — 8 unit tests inc. 32-thread concurrency stress
caches/RxPartitionKeyRangeCacheTest.java	+3 tests: cross-client sharing, cross-endpoint isolation, close idempotency
CHANGELOG.md	Entry under 4.82.0-beta.1

Test plan

• ✅ mvn compile (azure-cosmos)
• ✅ mvn checkstyle:check spotbugs:check (azure-cosmos)
• ✅ All cache unit tests pass (18 cache tests, 0 failures)
• ✅ ChangeFeedFetcherTest passes (14 tests)

Key behavioural tests

• twoCachesForSameEndpointShareRoutingMapStorage — populates routing map from client A, then asserts client B serves the same lookup with clientB.readPartitionKeyRanges invoked zero times.
• cachesForDifferentEndpointsDoNotShareStorage — both clients invoke their own readPartitionKeyRanges exactly once.
• closeIsIdempotent — repeated close() calls do not drive refcount negative.
• concurrentAcquireAndReleaseProducesConsistentRefcount — 32 threads × 200 ops, refcount ends at 0.
• disabledFlagReturnsIsolatedCachesAndPreservesRegistryEmpty — opt-out preserves pre-sharing behaviour.

Breaking changes

None. The public API of RxPartitionKeyRangeCache (an implementation package class) gains a new ctor overload and now implements Closeable. The two-arg ctor signature is preserved. No customer-visible APIs change.

[Copilot]

Pull request overview

This PR reduces duplicated routing-map cache memory and redundant /pkranges requests by sharing the storage layer of RxPartitionKeyRangeCache across CosmosClient / CosmosAsyncClient instances that target the same Cosmos account (keyed by service endpoint), while keeping the per-client fetch path unchanged. The shared cache is managed by a process-wide, refcounted registry and can be disabled via a new system property for opt-out.

Changes:

• Introduces SharedRoutingMapCacheRegistry (endpoint-keyed, refcounted) to share AsyncCacheNonBlocking<String, CollectionRoutingMap> across clients.
• Updates RxPartitionKeyRangeCache to acquire shared storage by endpoint and to implement Closeable for refcount release on client shutdown.
• Wires RxDocumentClientImpl.close() to release the cache reference, adds config flag plumbing, and adds targeted unit tests + changelog entry.

Show a summary per file

File	Description
sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java	Passes endpoint into the cache ctor and releases the cache reference during client close.
sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java	Adds COSMOS.SHARED_PARTITION_KEY_RANGE_CACHE_ENABLED flag (default true).
sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/caches/SharedRoutingMapCacheRegistry.java	New process-wide singleton registry for shared routing-map cache storage with refcounted eviction.
sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/caches/RxPartitionKeyRangeCache.java	Splits “storage” vs “fetcher” by sourcing storage from the shared registry and adding close() ref-release.
sdk/cosmos/azure-cosmos/CHANGELOG.md	Documents the new sharing behavior and opt-out property.
sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/caches/SharedRoutingMapCacheRegistryTest.java	New unit tests validating sharing, eviction, disabled behavior, and concurrency refcount correctness.
sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/caches/RxPartitionKeyRangeCacheTest.java	Adds tests validating cross-client sharing, cross-endpoint isolation, and idempotent close behavior.

Copilot's findings

• Files reviewed: 7/7 changed files
• Comments generated: 1