[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 are configured with the same service endpoint (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 the service endpoint URI configured on CosmosClientBuilder. 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 SharedPartitionKeyRangeCacheRegistry (process-wide singleton) keyed by the service endpoint URI.
Fetcher — issues /pkranges, depends on per-client RxDocumentClientImpl, RxCollectionCache, diagnostics. Unchanged.

Scope of sharing

Two clients share the cache only when their service endpoint URIs compare equal via URI.equals (case-insensitive on host per RFC 3986). Clients configured with different endpoint URIs — including the global endpoint vs a regional endpoint of the same logical account — do not share.

The natural-looking alternative of keying by DatabaseAccount.getId() (so global + regional clients of the same account would share) was tried and rejected: the id returned from a regional endpoint is <globalId>-<service-normalised-region>, and recovering the global form requires brittle suffix-stripping against the readable/writable locations list. DatabaseAccount.getResourceId() (the _rid field) is not a documented canonical id at the protocol level. Rather than ship a fragile canonicalisation, the registry honestly keys on the builder-supplied URI.

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 implements Closeable; close() releases the refcount and is idempotent (guarded by AtomicBoolean).
RxDocumentClientImpl.close() calls LifeCycleUtils.closeQuietly(partitionKeyRangeCache).
A leak-safety net registers a one-shot cleanup with com.azure.core.util.ReferenceManager: if a client is GC'd without calling close(), the cleanup decrements the refcount once. A WARN log identifies the leaking endpoint.
When the last reference is released, the registry entry is evicted so idle endpoints don't pin memory.

Diagnostics

Diagnostics behaviour is unchanged: the PARTITION_KEY_RANGE_LOOK_UP metadata diagnostic is recorded only when a real /pkranges network fetch happens (inside getRoutingMapForCollectionAsync), exactly as before this change. A consequence of sharing the routing-map storage is that a client can serve a PK-range lookup from a cache already populated by a sibling client on the same endpoint without issuing any /pkranges fetch — in which case no PARTITION_KEY_RANGE_LOOK_UP diagnostic is recorded for that operation. Tests that previously assumed the diagnostic is always present were updated accordingly: CosmosDiagnosticsTest.validateDirectModeDiagnosticsOnSuccess no longer asserts its presence, while FaultInjectionMetadataRequestRuleTests keeps its original single-entry assertion because it forces a routing-map refresh (so a network fetch — and the delayed diagnostic — is guaranteed).

Opt-out

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

Files

File	Change
`caches/SharedPartitionKeyRangeCacheRegistry.java`	NEW — process-wide registry singleton, keyed by service-endpoint `URI`
`caches/RxPartitionKeyRangeCache.java`	3-arg ctor `(client, collectionCache, URI)`; registry-backed storage; idempotent `close()` (diagnostics emission unchanged — still recorded only on the `/pkranges` network fetch path)
`Configs.java`	New system property `COSMOS.SHARED_PARTITION_KEY_RANGE_CACHE_ENABLED` (default: enabled)
`RxDocumentClientImpl.java`	Pass `this.serviceEndpoint` to the cache ctor; release the cache in `close()`
`caches/SharedPartitionKeyRangeCacheRegistryTest.java`	NEW — 13 unit tests inc. 32-thread concurrency stress, GC-driven leak cleanup, URI host case-insensitivity, and a negative-case pin that regional vs global endpoints don't share
`caches/RxPartitionKeyRangeCacheTest.java`	+5 unit tests: cross-client sharing, cross-endpoint isolation, close idempotency, ctor-lifecycle, cross-client force-refresh visibility
`SharedPartitionKeyRangeCacheE2ETest.java`	NEW — e2e test against a real Cosmos endpoint: positive sharing on the same endpoint (cross-endpoint isolation is covered by unit tests, since `CosmosClientBuilder` normalises the endpoint URI so two distinct connectable endpoints can't be built in a single-endpoint test environment)
`CHANGELOG.md`	Entry under 4.82.0-beta.1

Test plan

✅ mvn install (azure-cosmos)
✅ mvn checkstyle:check spotbugs:check (azure-cosmos + azure-cosmos-tests)
✅ Unit tests pass: 24 tests, 0 failures (RxPartitionKeyRangeCacheTest + SharedPartitionKeyRangeCacheRegistryTest)
⏳ E2e tests (SharedPartitionKeyRangeCacheE2ETest) registered under the emulator and fast Maven profiles — executed in CI against the configured Cosmos endpoint.

Key behavioural tests (unit)

twoCachesForSameEndpointShareRoutingMapStorage — client A populates the routing map, client B serves the same lookup with clientB.readPartitionKeyRanges invoked zero times.
cachesForDifferentEndpointsDoNotShareStorage — clients with different endpoint URIs each invoke their own readPartitionKeyRanges exactly once.
forceRefreshOnSharedCacheIsVisibleToSiblingClient — client A's force-refresh propagates to client B without B issuing its own fetch.
closeIsIdempotent — repeated close() calls do not drive refcount negative.
clientWithServiceEndpointAcquiresAndReleasesRegistryRefcount — regression guard for the RxDocumentClientImpl.close() → partitionKeyRangeCache.close() wiring.
concurrentAcquireAndReleaseProducesConsistentRefcount — 32 threads × 200 ops, refcount ends at 0.
referenceManagerReleasesSharedCacheWhenOwnerIsGarbageCollected — leak-safety net: an unclosed client is reclaimed by ReferenceManager once GC'd.
acquireTreatsHostCaseInsensitivelyMatchingUriEquals — RFC 3986 host case-insensitivity flows through to the registry key.
regionalAndGlobalEndpointsDoNotShareStorage — pins the explicit scope: distinct endpoint URIs use distinct registry entries.
disabledFlagReturnsIsolatedCachesAndPreservesRegistryEmpty — opt-out preserves pre-sharing behaviour.

Key behavioural tests (e2e, real Cosmos endpoint)

twoClientsOnSameEndpointShareRoutingMapStorage — spins up two real CosmosAsyncClients configured with the same endpoint, performs PK-routed reads on both, and asserts they share the same AsyncCacheNonBlocking instance, the registry refcount accounts for both holders, and closing each client decrements the refcount by exactly one.
Cross-endpoint isolation (clients with distinct endpoint URIs use distinct registry entries) is pinned by unit tests — SharedPartitionKeyRangeCacheRegistryTest.acquireReturnsDifferentInstanceForDifferentEndpoints / regionalAndGlobalEndpointsDoNotShareStorage and RxPartitionKeyRangeCacheTest.cachesForDifferentEndpointsDoNotShareStorage — rather than e2e, because CosmosClientBuilder.validateConfig() strips path/query so two distinct connectable endpoint URIs can't be constructed against a single test endpoint.

Breaking changes

None. RxPartitionKeyRangeCache is in the implementation package; its ctor signature and its new Closeable supertype are not part of the public API surface. 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>

azure-pipelines · 2026-07-02T20:40:57Z

Azure Pipelines successfully started running 1 pipeline(s).

… coverage Resolves open review feedback on the shared PartitionKeyRangeCache PR: - CHANGELOG: add the PR 49560 link to the 4.82.0-beta.1 entry. - RxPartitionKeyRangeCacheTest: add 4 unit tests covering the sharing behavior that was previously untested: * concurrentLookupsOnSharedCacheIssueSingleFetch - single-flight across clients (concurrent lookups collapse to exactly one /pkranges fetch). * failedFetchDoesNotPoisonSharedCacheForSiblings - a failed fetch is not cached; a sibling can still populate and the failing client recovers. * diagnosticsRecordedOnlyByFetchingClientNotSibling - positive assertion that the fetching client records PARTITION_KEY_RANGE_LOOK_UP while a cache-served sibling records none. * sharingDisabledYieldsIsolatedCachesPerClient - opt-out exercised through the cache/client (not just the registry): each client fetches its own. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

xinlian12 · 2026-07-06T17:55:43Z

/azp run java - cosmos - tests

azure-pipelines · 2026-07-06T17:56:16Z

Azure Pipelines successfully started running 1 pipeline(s).

xinlian12 · 2026-07-06T22:11:54Z

/azp run java - cosmos - kafka

xinlian12 · 2026-07-06T22:12:00Z

/azp run java - cosmos - spark

azure-pipelines · 2026-07-06T22:12:10Z

Azure Pipelines successfully started running 1 pipeline(s).

azure-pipelines · 2026-07-06T22:12:20Z

Azure Pipelines successfully started running 1 pipeline(s).

The shared routing-map cache makes the partition-key-range cache JVM-wide (shared across all clients on the same service endpoint), so tests that assumed a per-client routing map now see sibling containers' entries. This is the same class of breakage the PR already fixed in CosmosContainerOpenConnectionsAndInitCachesTest, but three test files were missed: - ProactiveConnectionManagementTest (groups multi-master / flaky-multi-master): 4 assertions of routingMap.size() == cosmosContainerIdentities.size() now fail because the shared map also holds sibling containers. Replaced with per-container assertThat(routingMap).containsKey(rid) checks (matching the established fix pattern). This is the root cause of the failing *_Tcp_MultiMaster CI legs. - AsyncCacheNonBlockingIntegrationTest (group split) and ProactiveOpenConnectionsProcessorTest (group multi-region): routingMap.keys().nextElement() grabbed an arbitrary key from the now multi-entry shared map; the before/after-split key comparison was no longer tied to the tested container. Replaced with containsKey(collectionRid) on the specific container so the "entry refreshed in place across the split" intent is preserved and robust under sharing. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

xinlian12 · 2026-07-06T22:44:18Z

/azp run java - cosmos - tests

azure-pipelines · 2026-07-06T22:44:48Z

Azure Pipelines successfully started running 1 pipeline(s).

xinlian12 · 2026-07-06T22:57:46Z

/azp run java - cosmos - tests

azure-pipelines · 2026-07-06T22:58:11Z

Azure Pipelines successfully started running 1 pipeline(s).

kushagraThapar · 2026-07-07T04:29:59Z

+
+        this.sharedCacheEndpointKey = serviceEndpoint;
+        SharedPartitionKeyRangeCacheRegistry.AcquireResult acquired =
+            SharedPartitionKeyRangeCacheRegistry.getInstance().acquire(this.sharedCacheEndpointKey, this);


This sharing looks like a nice improvement for reducing duplicate /pkranges work across clients. One edge case I’d love to sanity-check is the cold-cache single-flight behavior when the winning initializer fails or stalls.

Since the shared AsyncCacheNonBlocking can make sibling clients attach to the same in-flight initializer, could client B end up waiting on or observing client A’s failed/stalled fetch instead of retrying through its own client-specific auth/transport path?

Would it be worth adding a deterministic test for this? For example: client A starts the shared initialization and fails or never completes, while same-endpoint client B would be able to fetch successfully. A Sinks.One, TestPublisher, or latch-based setup might make this easier to test without timing assumptions.

There's already a deterministic test for this: RxPartitionKeyRangeCacheTest.failedFetchDoesNotPoisonSharedCacheForSiblings.

It exercises exactly the failure case you describe — client A wins the shared initialization and its /pkranges fetch fails, while a same-endpoint client B must still succeed through its own path:

A's readPartitionKeyRanges always errors ⇒ A's lookup verifyError(), clientACalls == 1.

B then populates the shared entry via its own readPartitionKeyRanges (clientBCalls == 1) — proving the failed initializer is not cached/poisoned and B fetches through its own client-specific transport, not A's.

A's subsequent lookup then succeeds from the healed shared entry.

It uses a deterministic mock error (no timing assumptions), matching the AsyncCacheNonBlocking contract: a failed initializer is removed from the cache (shouldRemoveFromCache() → remove(key)), so siblings re-initialize rather than inherit the failure.

One honest distinction: this covers fail-fast. The "never completes / indefinitely stalled" case is genuinely different — with single-flight, a concurrent sibling attaches to the in-flight initializer rather than starting its own, so it does wait on A until that Mono errors/completes (there's no separate per-client timeout at this cache layer; request-level e2e-timeout/retry policies sit above it). So the "B doesn't inherit a failed fetch" guarantee is tested; "B races an unresolved fetch" is by-design single-flight, not a separate code path.

kushagraThapar · 2026-07-07T04:30:01Z

+        this.sharedCacheEndpointKey = serviceEndpoint;
+        SharedPartitionKeyRangeCacheRegistry.AcquireResult acquired =
+            SharedPartitionKeyRangeCacheRegistry.getInstance().acquire(this.sharedCacheEndpointKey, this);
+        this.routingMapCache = acquired.cache;


This should help reduce duplicate routing-map storage and /pkranges fan-out, which is great. Since this is on a pretty hot routing-cache path and is default-on, could we include a small targeted benchmark or some before/after numbers?

The most useful numbers would probably be:

/pkranges call count with N same-account clients

memory/allocation impact

lookup latency

contention under concurrent cold-cache access

Not necessarily asking for a huge benchmark suite — even a focused benchmark or reproducible measurement in the PR description would make the trade-off much easier to validate.

The correctness side of "N same-account clients ⇒ collapsed /pkranges" is already pinned deterministically (as assertions rather than a perf benchmark):

RxPartitionKeyRangeCacheTest.twoCachesForSameEndpointShareRoutingMapStorage — two clients on the same endpoint ⇒ exactly one /pkranges fetch (clientACalls == 1, clientBCalls == 0).

cachesForDifferentEndpointsDoNotShareStorage — different endpoints ⇒ each client fetches once (negative control).

concurrentLookupsOnSharedCacheIssueSingleFetch — two clients hitting the cold shared cache concurrently ⇒ exactly one fetch total (single-flight under contention).

SharedPartitionKeyRangeCacheRegistryTest.concurrentAcquireAndReleaseProducesConsistentRefcount — 32 threads × 200 acquire/release ⇒ consistent refcount (registry contention).

So the /pkranges call-count reduction and concurrent cold-cache contention behavior are covered. What these don't produce is the memory/allocation and lookup-latency numbers — those would need a separate micro-benchmark. Happy to add before/after figures to the PR description if that's the bar you'd like.

ack, will do a focused benchmark

…he initiating client Fixes a shared-cache owner-retention leak surfaced in PR review. On a cold key, AsyncLazyWithRefresh stored the initial value as `taskFactory.apply(null).cache()`. Reactor's `.cache()` keeps its upstream source chain reachable, and for the routing-map cache that chain captures the initiating client's object graph (RxDocumentClientImpl, collection cache, diagnostics) via the per-client fetch lambda. Now that the routing-map storage is shared across clients on the same endpoint (SharedPartitionKeyRangeCacheRegistry), the shared entry kept the first client to populate a key strongly reachable for the life of the entry: - after close(), a sibling keeping the entry alive prevented the closed client from being GC'd (memory leak), and - for an unclosed client, the owner never became phantom-reachable, so the ReferenceManager leak-safety net never fired. Fix: once the initial load succeeds, swap the cached value for a detached `Mono.just(response)` (mirroring createBackgroundRefreshTask), so the entry stops retaining the initializer's source chain. Single-flight and failure-isolation are preserved (the outer `.cache()` still runs the source once; on error the swap never runs and the key is removed). Adds RxPartitionKeyRangeCacheTest.initiatingClientReleasedAfterCloseWhileSibling- KeepsSharedEntryAlive, which fails without this fix (owner retained) and passes with it. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

xinlian12 · 2026-07-07T05:30:35Z

/azp run java - cosmos - tests

azure-pipelines · 2026-07-07T05:31:10Z

Azure Pipelines successfully started running 1 pipeline(s).

xinlian12 · 2026-07-07T05:36:35Z

/azp run java - cosmos - tests

xinlian12 · 2026-07-07T05:36:56Z

/azp run java - cosmos - spark

azure-pipelines · 2026-07-07T05:36:57Z

Azure Pipelines successfully started running 1 pipeline(s).

xinlian12 · 2026-07-07T05:37:02Z

/azp run java - cosmos - kafka

azure-pipelines · 2026-07-07T05:37:06Z

Azure Pipelines successfully started running 1 pipeline(s).

azure-pipelines · 2026-07-07T05:37:14Z

Azure Pipelines successfully started running 1 pipeline(s).

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

xinlian12 commented Jun 19, 2026

View reviewed changes

Comment thread ...tests/src/test/java/com/azure/cosmos/implementation/caches/RxPartitionKeyRangeCacheTest.java Outdated

xinlian12 commented Jun 19, 2026

View reviewed changes

Comment thread ...smos/src/main/java/com/azure/cosmos/implementation/caches/SharedRoutingMapCacheRegistry.java Outdated

xinlian12 commented Jun 19, 2026

View reviewed changes

Comment thread ...re-cosmos/src/main/java/com/azure/cosmos/implementation/caches/RxPartitionKeyRangeCache.java Outdated

xinlian12 commented Jun 19, 2026

View reviewed changes

Comment thread ...re-cosmos/src/main/java/com/azure/cosmos/implementation/caches/RxPartitionKeyRangeCache.java Outdated

xinlian12 commented Jun 19, 2026

View reviewed changes

Comment thread ...tests/src/test/java/com/azure/cosmos/implementation/caches/RxPartitionKeyRangeCacheTest.java

xinlian12 commented Jun 19, 2026

View reviewed changes

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

xinlian12 commented Jun 19, 2026

View reviewed changes

Comment thread ...c/main/java/com/azure/cosmos/implementation/caches/SharedPartitionKeyRangeCacheRegistry.java Outdated

xinlian12 and others added 3 commits June 19, 2026 09:54

remove kafka test output

9b43616

xinlian12 commented Jun 19, 2026

View reviewed changes

Comment thread ...c/main/java/com/azure/cosmos/implementation/caches/SharedPartitionKeyRangeCacheRegistry.java

xinlian12 commented Jun 19, 2026

View reviewed changes

Comment thread ...smos/src/main/java/com/azure/cosmos/implementation/caches/SharedRoutingMapCacheRegistry.java Outdated

xinlian12 force-pushed the feature/shared-partition-key-range-cache branch from 1a27dc2 to 9b43616 Compare June 19, 2026 20:39

xinlian12 and others added 2 commits June 19, 2026 14:56

xinlian12 and others added 3 commits July 6, 2026 09:39

merge from main and resolve conflicts

c93502a

revert unnecessary changes

b8b8379

kushagraThapar reviewed Jul 7, 2026

View reviewed changes

Comment thread ...s/azure-cosmos-tests/src/test/java/com/azure/cosmos/SharedPartitionKeyRangeCacheE2ETest.java

kushagraThapar reviewed Jul 7, 2026

View reviewed changes

Comment thread ...re-cosmos/src/main/java/com/azure/cosmos/implementation/caches/RxPartitionKeyRangeCache.java

revert unnecessary changes

a5c33f4

xinlian12 force-pushed the feature/shared-partition-key-range-cache branch from 4a92a85 to a5c33f4 Compare July 7, 2026 05:06

Uh oh!

Conversation

xinlian12 commented Jun 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Design

Scope of sharing

Concurrency model

Lifecycle

Diagnostics

Opt-out

Files

Test plan

Key behavioural tests (unit)

Key behavioural tests (e2e, real Cosmos endpoint)

Breaking changes

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Copilot's findings

Uh oh!

Uh oh!

xinlian12 commented Jun 19, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

xinlian12 commented Jun 19, 2026

Uh oh!

Uh oh!

Uh oh!

xinlian12 commented Jun 19, 2026

Uh oh!

azure-pipelines Bot commented Jul 2, 2026

Uh oh!

xinlian12 commented Jul 6, 2026

Uh oh!

azure-pipelines Bot commented Jul 6, 2026

Uh oh!

xinlian12 commented Jul 6, 2026

Uh oh!

xinlian12 commented Jul 6, 2026

Uh oh!

azure-pipelines Bot commented Jul 6, 2026

Uh oh!

azure-pipelines Bot commented Jul 6, 2026

Uh oh!

xinlian12 commented Jul 6, 2026

Uh oh!

azure-pipelines Bot commented Jul 6, 2026

Uh oh!

xinlian12 commented Jul 6, 2026

Uh oh!

azure-pipelines Bot commented Jul 6, 2026

Uh oh!

kushagraThapar Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

xinlian12 Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

kushagraThapar Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

xinlian12 Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

xinlian12 Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

xinlian12 commented Jul 7, 2026

Uh oh!

azure-pipelines Bot commented Jul 7, 2026

xinlian12 commented Jun 18, 2026 •

edited

Loading