Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
xinlian12 wants to merge 6 commits into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
5a4b74d
Cosmos Java: share PartitionKeyRangeCache across CosmosClients target…
xinlian12 Jun 18, 2026
f3fa638
Cosmos Java: use URI (not String) as registry key for case-insensitiv…
xinlian12 Jun 19, 2026
75f93d5
Cosmos Java: keep registry key as serviceEndpoint URI (not _rid) for …
xinlian12 Jun 19, 2026
05f6780
Cosmos Java: add PhantomReference-based leak safety net for unclosed …
xinlian12 Jun 19, 2026
cbd47a8
Cosmos Java: use azure-core ReferenceManager for leaked-client safety…
xinlian12 Jun 19, 2026
9b43616
remove kafka test output
xinlian12 Jun 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
162 changes: 161 additions & 1 deletion ...ts/src/test/java/com/azure/cosmos/implementation/caches/RxPartitionKeyRangeCacheTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,10 +20,13 @@
import reactor.core.publisher.Mono;
import reactor.test.StepVerifier;

import java.net.URI;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.concurrent.atomic.AtomicInteger;

import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Fail.fail;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
Expand All @@ -35,7 +38,7 @@ public class RxPartitionKeyRangeCacheTest {
private RxDocumentClientImpl client;
private RxCollectionCache collectionCache;
private RxPartitionKeyRangeCache cache;

@BeforeMethod(groups = "unit")
public void before_test() {
client = Mockito.mock(RxDocumentClientImpl.class);
Expand Down Expand Up @@ -245,4 +248,161 @@ public void tryLookupAsync_RetriesOnceAndConvertsToNotFoundException() {
.expectNextMatches(s -> s.v == null)
.verifyComplete();
}

@Test(groups = "unit")
public void twoCachesForSameEndpointShareRoutingMapStorage() throws Exception {
// Two RxPartitionKeyRangeCache instances pointing at the same endpoint should
// share their underlying AsyncCacheNonBlocking, so a routing map populated by
// one is immediately visible to the other without a second /pkranges call.
URI endpoint = new URI("https://test-shared-pkr-1.documents.azure.com:443/");

RxDocumentClientImpl clientA = Mockito.mock(RxDocumentClientImpl.class);
RxDocumentClientImpl clientB = Mockito.mock(RxDocumentClientImpl.class);
RxCollectionCache collA = Mockito.mock(RxCollectionCache.class);
RxCollectionCache collB = Mockito.mock(RxCollectionCache.class);

String collectionRid = "shared-coll-1";
DocumentCollection collection = new DocumentCollection();
collection.setResourceId(collectionRid);
collection.setSelfLink("dbs/db1/colls/coll1");

PartitionKeyRange range = new PartitionKeyRange();
range.setId("0");
range.setMinInclusive(PartitionKeyRange.MINIMUM_INCLUSIVE_EFFECTIVE_PARTITION_KEY);
range.setMaxExclusive(PartitionKeyRange.MAXIMUM_EXCLUSIVE_EFFECTIVE_PARTITION_KEY);

FeedResponse<PartitionKeyRange> response = Mockito.mock(FeedResponse.class);
when(response.getResults()).thenReturn(Arrays.asList(range));
when(response.getContinuationToken()).thenReturn("etag-1");

AtomicInteger clientACalls = new AtomicInteger();
AtomicInteger clientBCalls = new AtomicInteger();

when(collA.resolveCollectionAsync(any(), any()))
.thenReturn(Mono.just(new Utils.ValueHolder<>(collection)));
when(clientA.readPartitionKeyRanges(eq(collection.getSelfLink()), any(CosmosQueryRequestOptions.class)))
.thenAnswer(invocation -> {
clientACalls.incrementAndGet();
return Flux.just(response);
});

when(collB.resolveCollectionAsync(any(), any()))
.thenReturn(Mono.just(new Utils.ValueHolder<>(collection)));
when(clientB.readPartitionKeyRanges(eq(collection.getSelfLink()), any(CosmosQueryRequestOptions.class)))
.thenAnswer(invocation -> {
clientBCalls.incrementAndGet();
return Flux.just(response);
});

RxPartitionKeyRangeCache cacheA = new RxPartitionKeyRangeCache(clientA, collA, endpoint);
RxPartitionKeyRangeCache cacheB = new RxPartitionKeyRangeCache(clientB, collB, endpoint);

try {
// First client populates the cache.
StepVerifier.create(cacheA.tryLookupAsync(null, collectionRid, null, new HashMap<>()))
.expectNextMatches(v -> v != null && v.v != null)
.verifyComplete();

// Second client must hit the shared cache without issuing its own /pkranges read.
StepVerifier.create(cacheB.tryLookupAsync(null, collectionRid, null, new HashMap<>()))
.expectNextMatches(v -> v != null && v.v != null)
.verifyComplete();

assertThat(clientACalls.get())
.as("client A populates the shared routing map")
.isEqualTo(1);
assertThat(clientBCalls.get())
.as("client B must hit the shared cache without issuing its own /pkranges call")

@xinlian12 xinlian12 Jun 19, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Recommendation: Test Coverage — Cross-client invalidation/refresh visibility is not tested

assertThat(clientBCalls.get())
    .as("client B must hit the shared cache without issuing its own /pkranges call")
    .isZero();

This test proves positive sharing (B sees A's populated value), but the inverse — when A invalidates or force-refreshes a routing map, B must observe the new value — is the most safety-critical property of the shared model and is currently uncovered. Concretely, the following scenarios are not exercised:

  1. After A's lookup populates the cache, simulate a PartitionKeyRangeGoneException (split scenario) on A causing AsyncCacheNonBlocking.remove(rid); assert B's next tryLookupAsync triggers a fresh fetch and observes the new routing map.
  2. After A's lookup populates the cache, A's tryLookupAsync(previousValue=…) force-refreshes (since shouldForceRefresh(prev, current) returns true); assert B's subsequent lookup sees the refreshed value, not the stale one.
  3. A 404 NotFound on a background refresh evicts the entry — assert that both clients then re-populate consistently rather than B silently keeping a stale reference to the removed AsyncLazyWithRefresh.

These are the exact scenarios the .NET SDK exercises (see EmulatorTests/PartitionKeyRangeCacheTests.cs::TestRidRefreshOnNotFoundAsync and VerifyPkRangeCacheRefreshOnSplitWithErrorsAsync for the patterns). Without them, a regression where invalidation doesn't propagate across sharing clients would slip past CI.

Additionally, twoCachesForSameEndpointShareRoutingMapStorage could be tightened by asserting object identity of the resolved CollectionRoutingMap (B sees the exact same instance A populated), which is the strongest evidence that storage truly shares.

⚠️ AI-generated review — may be incorrect. Agree? → resolve the conversation. Disagree? → reply with your reasoning.

.isZero();
} finally {
cacheA.close();
cacheB.close();
}

assertThat(SharedRoutingMapCacheRegistry.getInstance().referenceCount(endpoint))
.as("close() releases the shared cache reference")
.isZero();
}

@Test(groups = "unit")
public void cachesForDifferentEndpointsDoNotShareStorage() throws Exception {
URI endpointA = new URI("https://test-shared-pkr-2a.documents.azure.com:443/");
URI endpointB = new URI("https://test-shared-pkr-2b.documents.azure.com:443/");

RxDocumentClientImpl clientA = Mockito.mock(RxDocumentClientImpl.class);
RxDocumentClientImpl clientB = Mockito.mock(RxDocumentClientImpl.class);
RxCollectionCache collA = Mockito.mock(RxCollectionCache.class);
RxCollectionCache collB = Mockito.mock(RxCollectionCache.class);

String collectionRid = "shared-coll-2";
DocumentCollection collection = new DocumentCollection();
collection.setResourceId(collectionRid);
collection.setSelfLink("dbs/db1/colls/coll2");

PartitionKeyRange range = new PartitionKeyRange();
range.setId("0");
range.setMinInclusive(PartitionKeyRange.MINIMUM_INCLUSIVE_EFFECTIVE_PARTITION_KEY);
range.setMaxExclusive(PartitionKeyRange.MAXIMUM_EXCLUSIVE_EFFECTIVE_PARTITION_KEY);

FeedResponse<PartitionKeyRange> response = Mockito.mock(FeedResponse.class);
when(response.getResults()).thenReturn(Arrays.asList(range));
when(response.getContinuationToken()).thenReturn("etag-2");

AtomicInteger clientACalls = new AtomicInteger();
AtomicInteger clientBCalls = new AtomicInteger();

when(collA.resolveCollectionAsync(any(), any()))
.thenReturn(Mono.just(new Utils.ValueHolder<>(collection)));
when(clientA.readPartitionKeyRanges(eq(collection.getSelfLink()), any(CosmosQueryRequestOptions.class)))
.thenAnswer(invocation -> {
clientACalls.incrementAndGet();
return Flux.just(response);
});
when(collB.resolveCollectionAsync(any(), any()))
.thenReturn(Mono.just(new Utils.ValueHolder<>(collection)));
when(clientB.readPartitionKeyRanges(eq(collection.getSelfLink()), any(CosmosQueryRequestOptions.class)))
.thenAnswer(invocation -> {
clientBCalls.incrementAndGet();
return Flux.just(response);
});

RxPartitionKeyRangeCache cacheA = new RxPartitionKeyRangeCache(clientA, collA, endpointA);
RxPartitionKeyRangeCache cacheB = new RxPartitionKeyRangeCache(clientB, collB, endpointB);

try {
StepVerifier.create(cacheA.tryLookupAsync(null, collectionRid, null, new HashMap<>()))
.expectNextMatches(v -> v != null && v.v != null)
.verifyComplete();

StepVerifier.create(cacheB.tryLookupAsync(null, collectionRid, null, new HashMap<>()))
.expectNextMatches(v -> v != null && v.v != null)
.verifyComplete();

// Two different endpoints → no sharing, each client issues its own read.
assertThat(clientACalls.get()).isEqualTo(1);
assertThat(clientBCalls.get()).isEqualTo(1);
} finally {
cacheA.close();
cacheB.close();
}
}

@Test(groups = "unit")
public void closeIsIdempotent() throws Exception {

@xinlian12 xinlian12 Jun 19, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟢 Suggestion: Test Coverage — Cross-client single-flight invariant is asserted in the PR description but not proven by any test

public void closeIsIdempotent() throws Exception {

The CHANGELOG and PR description both call out a stronger guarantee — "only one in-flight /pkranges fetch per (account, container) at any time, even across clients". But twoCachesForSameEndpointShareRoutingMapStorage blocks on client A's fetch (verifyComplete()) before issuing client B's lookup, so B is guaranteed to hit a fully-populated cache. That tests cross-client visibility of completed fetches, not cross-client coalescing of concurrent fetches.

Suggest adding a test that:

  1. Makes clientA.readPartitionKeyRanges and clientB.readPartitionKeyRanges both block on a CountDownLatch before returning the Flux.just(response).
  2. Concurrently subscribes both cacheA.tryLookupAsync(...) and cacheB.tryLookupAsync(...) (e.g. subscribeOn(Schedulers.parallel())).
  3. Counts down the latch.
  4. Asserts clientACalls.get() + clientBCalls.get() == 1 — proving exactly one client's network call was actually invoked.

Without this, a future regression where tryLookupAsync accidentally creates a fresh AsyncLazyWithRefresh per call (instead of going through the shared AsyncCacheNonBlocking lookup) would not break any existing test, but would silently invalidate one of the PR's central claims.

⚠️ AI-generated review — may be incorrect. Agree? → resolve the conversation. Disagree? → reply with your reasoning.

URI endpoint = new URI("https://test-shared-pkr-3.documents.azure.com:443/");
RxDocumentClientImpl mockClient = Mockito.mock(RxDocumentClientImpl.class);
RxCollectionCache mockColl = Mockito.mock(RxCollectionCache.class);

RxPartitionKeyRangeCache c = new RxPartitionKeyRangeCache(mockClient, mockColl, endpoint);
assertThat(SharedRoutingMapCacheRegistry.getInstance().referenceCount(endpoint))
.isEqualTo(1);

c.close();
c.close(); // second call must be a no-op
c.close();

assertThat(SharedRoutingMapCacheRegistry.getInstance().referenceCount(endpoint))
.as("repeated close() must not drive refcount negative")
.isZero();
}
}
Loading
Loading