[NO REVIEW][Cosmos] Share PartitionKeyRangeCache across CosmosClients targeting the same account by xinlian12 · Pull Request #49560 · Azure/azure-sdk-for-java

xinlian12 · 2026-06-18T18:42:17Z

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:

Storage — AsyncCacheNonBlocking<String, CollectionRoutingMap>. Account-level data, naturally shareable. Now obtained from SharedRoutingMapCacheRegistry (process-wide singleton) keyed by the service endpoint URL string.
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.

…ing the same account Move the partition-key-range routing-map cache from per-CosmosClient to a process-wide, refcounted registry keyed by service endpoint. Multiple CosmosClient / CosmosAsyncClient instances in the same JVM targeting the same Cosmos account now share a single AsyncCacheNonBlocking instance for collection -> CollectionRoutingMap, eliminating duplicate routing-map memory and redundant /pkranges fetches. Design - New SharedRoutingMapCacheRegistry (process-wide singleton) holds an AsyncCacheNonBlocking per endpoint URL plus an AtomicInteger refcount. All state transitions go through ConcurrentHashMap.compute, giving atomic per-key check-and-update without a global lock. - RxPartitionKeyRangeCache: new ctor accepts the service endpoint; underlying routingMapCache is obtained from the registry. Implements Closeable; close() releases this client's reference and is idempotent. - RxDocumentClientImpl: passes serviceEndpoint to the cache ctor and releases the cache reference in its close() path. - Opt-out: COSMOS.SHARED_PARTITION_KEY_RANGE_CACHE_ENABLED=false restores the pre-sharing behaviour (each client owns a private cache). Why this is safe - PK-range data is account-level metadata, not credential-bound. - AsyncCacheNonBlocking already enforces single-flight per key; sharing the instance strengthens that to "single in-flight /pkranges per (account, container) across all clients". - The two-arg back-compat ctor resolves the endpoint from the client, so existing mocked tests continue to work (mock returns null endpoint -> isolated cache, matching today's behaviour). Tests - New SharedRoutingMapCacheRegistryTest: acquire/release sharing, refcount eviction, idempotent release, null-endpoint isolation, opt-out flag, 32-thread concurrent acquire/release stress. - New RxPartitionKeyRangeCacheTest cases: two caches at same endpoint share storage (verified by mock /pkranges call count = 1, not 2), caches at different endpoints stay independent, close() is idempotent. - Existing 7 RxPartitionKeyRangeCacheTest cases unchanged and passing. Reference Pattern matches Python (sdk/cosmos/azure-cosmos/azure/cosmos/_routing/ routing_map_provider.py) which uses module-level endpoint-keyed dicts with refcounted cleanup. Adapted to Java idioms (ConcurrentHashMap.compute instead of explicit RLock, Closeable instead of __del__). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

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

xinlian12 · 2026-06-19T15:55:55Z

@sdkReviewAgent

…e host matching Switch SharedRoutingMapCacheRegistry's key type from String to URI so URI.equals() — which is case-insensitive on the host component per RFC 3986 — is used for sharing identity. Previously, two clients built with 'https://Acct.documents.azure.com/' and 'https://acct.documents.azure.com/' would fragment into two registry entries even though they target the same account. With URI as the key the two collapse into a single shared entry. This matches the spirit of the Rust SDK, which uses Url-based equality on its AccountReference identity. Python uses raw string comparison; Java's URI gives us strictly better behaviour for free. Added a new test (acquireTreatsHostCaseInsensitivelyMatchingUriEquals) that asserts URI.equals() considers the two casings equal AND that the registry produces a single shared entry for them. Ran 34 cache unit tests, 0 failures. No public API change. RxPartitionKeyRangeCache's three-arg ctor still takes URI; only the internal field type changed (String -> URI). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

…cross-SDK consistency Confirmed via cross-SDK review that both peer Cosmos SDKs key sharing on the user-supplied account endpoint URL, not on the account _rid: - Python (sdk/cosmos/azure-cosmos/azure/cosmos/_routing/_routing_map_provider_common.py): _resolve_endpoint() returns client.url_connection (the input endpoint string) with no normalisation and no _rid lookup. - Rust (sdk/cosmos/azure_data_cosmos_driver/src/models/account_reference.rs): AccountReference identity is endpoint-only via AccountEndpoint(Url) which Hash/Eq on the Url; PartialEq deliberately excludes credentials and backup endpoints. No _rid involvement. This SDK should match. The "regional vs global endpoint to the same account" case stays a known fragmentation case across all three SDKs rather than something Java solves alone via _rid. Why _rid keying was rejected after exploration: 1. Diverges from Python and Rust — increases mental-model and maintenance cost for cross-SDK contributors. 2. DatabaseAccount.getResourceId() returns the empty string in emulator and some service paths where the account JSON has no _rid (Resource.java:130 delegates to JsonSerializable.getString(R_ID)). Would silently fall back and fragment differently than peers. 3. Brittle to init reorders: today GlobalEndpointManager.init() runs before cache construction, but any future refactor (lazy account fetch, offline-mode init) would silently break sharing. Endpoint URI is constructor-immutable; _rid depends on a successful prior network call. Final shape: - Registry keyed by URI (case-insensitive host via URI.equals). - RxPartitionKeyRangeCache 3-arg ctor takes (client, collectionCache, serviceEndpoint URI). Two-arg ctor delegates with client.getServiceEndpoint(). - JavaDoc on SharedRoutingMapCacheRegistry now explicitly documents the cross-SDK alignment and the regional-endpoint fragmentation tradeoff. All 34 cache unit tests still pass. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

xinlian12 · 2026-06-19T16:32:03Z

✅ Review complete (35:07)

Posted 7 inline comment(s).

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

…clients Without this safety net, a customer that forgets to call CosmosClient.close() would pin the shared partition-key-range cache entry for the lifetime of the JVM. The owning RxPartitionKeyRangeCache holds a strong reference to the shared AsyncCacheNonBlocking and the registry's refcount stays > 0 forever. Peer SDKs handle this: - Python: __del__ in PartitionKeyRangeCache calls release() as a GC fallback (sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_map_provider.py L192). - Rust: no Drop impl needed — the cache lives as a field on the driver and Rust ownership guarantees cleanup on driver drop. Java cannot use java.lang.ref.Cleaner because azure-cosmos targets Java 8 (verified: sdk/parents/azure-client-sdk-parent/pom.xml <source>1.8</source>). Solution uses the pre-Cleaner pattern: PhantomReference + ReferenceQueue + daemon reaper thread. All Java 1.2+ APIs. Design - SharedRoutingMapCacheRegistry holds: * ReferenceQueue<Object> reaperQueue * Set<OwnerPhantom> livePhantoms (concurrent) — critical for correctness: the JVM only enqueues phantoms that are themselves still strongly reachable, so the registry must hold them alive until processed. * One daemon thread (cosmos-shared-pkr-cache-reaper) blocking on reaperQueue.remove(). - acquire(URI endpoint, Object owner): registers an OwnerPhantom on the owner, adds it to livePhantoms, returns AcquireResult { cache, phantom }. - release(URI, cache, PhantomReference) — new 3-arg overload — clears the phantom and removes it from livePhantoms in addition to decrementing the refcount. This is the path RxPartitionKeyRangeCache.close() uses. - When the owner becomes phantom-reachable, the reaper drains the queue, logs a WARN ("Leaked (unclosed) RxPartitionKeyRangeCache detected..."), calls release(endpoint, cache) to decrement refcount, then removes the phantom from livePhantoms. - close() is still the right primary path; the reaper is a safety net that prevents permanent JVM-lifetime cache pinning, not a substitute. Tests - reaperReleasesSharedCacheWhenOwnerIsGarbageCollected: acquires in a helper method (so the test frame cannot keep owner alive), polls referenceCount while forcing System.gc() in a 15s window. Reaper warning is observable in test output. - promptCloseClearsPhantomSoReaperDoesNotDoubleRelease: validates the prompt-close path clears the phantom and a subsequent GC produces no extra release. 36 cache unit tests pass (was 34, +2 new leak tests). Key correctness note in code The first attempt at this had a subtle bug: acquire() returned the phantom in AcquireResult but the registry didn't hold it. Once the test discarded the AcquireResult, the phantom became unreachable and the JVM never enqueued it — the reaper sat idle forever. The livePhantoms set fixes this. The fields/JavaDoc explicitly document the why. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

… net Replace the bespoke PhantomReference + ReferenceQueue + daemon-thread reaper with com.azure.core.util.ReferenceManager.INSTANCE, the SDK-wide singleton that already encapsulates this pattern. ReferenceManagerImpl: - On Java 9+ delegates reflectively to java.lang.ref.Cleaner. - On Java 8 (our baseline) uses an internal PhantomReference + daemon thread named "azure-sdk-referencemanager" — exactly the same mechanism this PR was reimplementing. Confirmed in test output: the leak WARN is logged on the "azure-sdk-referencemanager" thread, proving the azure-core path is wired. Why this is better: - Reuses supported, well-tested azure-core machinery instead of rolling our own. One thread per JVM regardless of how many SDK components opt into the pattern, instead of cosmos adding its own competing thread. - Java 9+ automatically gets the Cleaner-based implementation (better shutdown semantics, less thread-stack overhead). - Drops ~100 lines of bespoke phantom plumbing from SharedRoutingMapCacheRegistry (OwnerPhantom inner class, livePhantoms set, reaper loop). Net negative on code we maintain. Design notes preserved: - The lambda registered with ReferenceManager.INSTANCE.register MUST NOT capture `owner`, otherwise the owner never becomes phantom-reachable. We capture only the endpoint URI and the cache reference (both independent of the owner) and document this constraint in code. - ReleaseHandle is a one-shot AtomicBoolean fulfilment flag shared between the prompt close() path and the deferred ReferenceManager cleanup, so whichever runs first wins via compareAndSet and the refcount is decremented exactly once. 36 cache unit tests still pass; the leak test was renamed to referenceManagerReleasesSharedCacheWhenOwnerIsGarbageCollected to reflect the new mechanism. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Per PR feedback, comments in the shared-cache implementation were too verbose and contained cross-SDK comparisons that don't add value to maintainers reading the Java code. Trimmed everywhere: - SharedRoutingMapCacheRegistry: removed Python/Rust comparison paragraphs, the "Cross-SDK consistency" and "Leaked-client safety net" walls of text, and condensed JavaDoc on individual methods. Kept only the critical "lambda must not capture owner" comment because it's a correctness invariant that's easy to break in a refactor. - RxPartitionKeyRangeCache: removed the long ownerPhantom-style field comments; consolidated the class JavaDoc into two sentences. - Configs: condensed the system-property comment to two lines. - RxDocumentClientImpl: shortened the close-path log message. - CHANGELOG entry: condensed to a single sentence describing the change and the opt-out flag. - Tests: stripped the "First client / Second client" narration, the "must hit the shared cache" explanations, and the multi-paragraph preambles on the leak tests. Kept enough to explain the GC-related test setup since that's not obvious from the code. Behavior unchanged; 36 cache unit tests still pass. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

- Renamed SharedRoutingMapCacheRegistry → SharedPartitionKeyRangeCacheRegistry for consistency with the class it serves (RxPartitionKeyRangeCache). - Removed the test-only acquire(URI) overload that bypassed ReferenceManager registration; tests now use acquire(URI, owner) so the cleanup-action path is exercised end-to-end. - Added clientWithServiceEndpointAcquiresAndReleasesRegistryRefcount: regression test guarding the RxDocumentClientImpl.close() → partitionKeyRangeCache.close() → refcount-- wiring. Constructs the cache via the 2-arg ctor (matching production) and asserts the refcount delta on construct and close. - Added forceRefreshOnSharedCacheIsVisibleToSiblingClient: cross-client invalidation propagation. Client A populates → A force-refreshes after a simulated split → B's lookup sees A's refreshed value (same routing-map instance) without issuing its own /pkranges call. Asserts object identity on the shared CollectionRoutingMap. 38 cache unit tests pass (was 36). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

xinlian12 · 2026-06-19T22:32:03Z

@sdkReviewAgent

Previous run failed in azure-cosmos-spark_3-3_2-12 with a scala-maven-plugin classpath flake (xsbt/ZincCompiler$sbtAnalyzer$ ClassNotFoundException) unrelated to this PR's changes (PR touches azure-cosmos core; Spark connector is unaffected). Empty commit to re-run the pipeline. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Previous run hit a JVM crash on the macOS-latest JDK 8 runner (SIGBUS in libzip.dylib's newEntry while reading a jar/zip, JRE 8.0_492-b09). Native-code crash unrelated to this PR. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

xinlian12 · 2026-06-19T23:11:55Z

✅ Review complete (38:45)

Posted 3 inline comment(s).

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

Adds Javadoc to the 2-arg constructor explaining that it routes through SharedPartitionKeyRangeCacheRegistry for any real client (since client.getServiceEndpoint() returns non-null). Production call sites should prefer the explicit 3-arg ctor; the 2-arg form is kept for back- compat and the tests-with-Mockito-mocks path where getServiceEndpoint() returns null and the isolated cache is wanted. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Previous run hit the recurring scala-maven-plugin classpath flake in azure-cosmos-spark_3-3_2-12 (xsbt.JavaUtils$JavaForEach$ CNFE). Third occurrence of this flake on this PR; pattern matches resource starvation in the ADO agents' classloader. Unrelated to this PR (touches azure-cosmos core only). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

4th scala-maven-plugin Zinc classloader flake on this PR (xsbt.CompilerInterface CNFE in azure-cosmos-spark_3-4_2-12 this time; previously hit different xsbt classes in different Spark modules). Known issue with scala-maven-plugin 4.8.1 under parallel mvn -T 2. Unrelated to this PR. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

5th scala-maven-plugin Zinc classloader flake on this PR (xsbt.ExtractAPI$ReusableTreePrinter$ in spark_3-3_2-12 Windows). Different xsbt.* class hits each rerun; same plugin under mvn -T 2. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

xinlian12 · 2026-06-20T00:13:49Z

CI infra flake — 5th occurrence on this PR, recurring scala-maven-plugin Zinc classloader CNFE

Documenting the pattern so reviewers don't conflate this with PR breakage. This PR touches only azure-cosmos core (no Spark connector changes), but every retriggered build keeps failing at the Spark connector compile step with the same root cause: scala-maven-plugin:4.8.1 losing an xsbt.* class from its classloader under mvn -T 2 parallel execution.

#	Build	OS / JDK	Missing class	Module
1	6460399	Ubuntu / JDK 17	`xsbt/ZincCompiler$sbtAnalyzer$`	`spark_3-3_2-12`
2	6460548	macOS / JDK 8	(different: native JVM SIGBUS in libzip.dylib)	(tests)
3	6460575	Windows / JDK 8	`xsbt.JavaUtils$JavaForEach$`	`spark_3-3_2-12`
4	6460624	Ubuntu / JDK 17	`xsbt.CompilerInterface`	`spark_3-4_2-12`
5	6460694	Windows / JDK 8	`xsbt.ExtractAPI$ReusableTreePrinter$`	`spark_3-3_2-12`

Continuing to retrigger but this may need an infra-side look at scala-maven-plugin parallel-build behavior on ADO agents.

Copilot AI review requested due to automatic review settings June 18, 2026 18:42

xinlian12 requested review from a team and kirankumarkolli as code owners June 18, 2026 18:42

github-actions Bot added the Cosmos label Jun 18, 2026

Copilot started reviewing on behalf of xinlian12 June 18, 2026 18:42 View session

Copilot AI reviewed Jun 18, 2026

View reviewed changes

Comment thread sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java Outdated

xinlian12 changed the title ~~[Cosmos] Share PartitionKeyRangeCache across CosmosClients targeting the same account~~ [NO REVIEW][Cosmos] Share PartitionKeyRangeCache across CosmosClients targeting the same account Jun 18, 2026

xinlian12 and others added 2 commits June 19, 2026 09:22