[Cosmos] Share PartitionKeyRangeCache across CosmosClients targeting the same account

sdk/cosmos/azure-cosmos-benchmark/src/test/java/com/azure/cosmos/benchmark/WorkflowTest.java

@@ -122,7 +122,7 @@ public void readThroughput(boolean useNameLink) throws Exception {
             "SESSION",
             2,
             numberOfOperations,
-            0);
+            100);
         try {
             Main.main(new String[]{"-workloadConfig", configFile.getAbsolutePath()});
         } finally {

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/AsyncCacheNonBlockingIntegrationTest.java

@@ -112,8 +112,12 @@ public void createItem_withCacheRefresh() throws InterruptedException {
         assertThat(processedDoc.get()).isEqualTo(totalRequest);
         RxDocumentClientImpl rxDocumentClient =
             (RxDocumentClientImpl) this.bulkClient.getDocClientWrapper();
+        // The partition-key-range (routing map) cache is shared across clients on the same service
+        // endpoint, so it may hold multiple containers' entries. Assert on this container's own entry
+        // rather than an arbitrary key from the shared map.
+        String collectionRid = container.read().block().getProperties().getResourceId();
         ConcurrentHashMap<String, ?> routingMap = getRoutingMap(rxDocumentClient);
-        String cacheKeyBeforePartition = routingMap.keys().nextElement();
+        assertThat(routingMap).containsKey(collectionRid);
 
         // introduce a split and continue bulk operations after split. The partition key range cache has to be
         // refreshed and bulk processing should complete without errors
@@ -157,10 +161,10 @@ public void createItem_withCacheRefresh() throws InterruptedException {
             .as("Partition ranges should increase after split");
         logger.info("After split num partitions = {}", partitionKeyRangesAfterSplit.size());
 
+        // After the split the routing map entry for this container is refreshed in place under the
+        // same collection rid, so the entry must still be present.
         routingMap = getRoutingMap(rxDocumentClient);
-        String cacheKeyAfterPartition = routingMap.keys().nextElement();
-
-        assertThat(cacheKeyBeforePartition).isEqualTo(cacheKeyAfterPartition);
+        assertThat(routingMap).containsKey(collectionRid);
 
         responseFlux = container.executeBulkOperations(cosmosItemOperationFlux2, cosmosBulkExecutionOptions);
 

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosContainerChangeFeedTest.java

@@ -7,6 +7,7 @@
 package com.azure.cosmos;
 
 import com.azure.cosmos.implementation.DocumentCollection;
+import com.azure.cosmos.implementation.HttpConstants;
 import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
 import com.azure.cosmos.implementation.RetryAnalyzer;
 import com.azure.cosmos.implementation.Utils;
@@ -26,7 +27,6 @@
 import com.azure.cosmos.models.CosmosContainerResponse;
 import com.azure.cosmos.models.CosmosItemResponse;
 import com.azure.cosmos.models.FeedRange;
-import com.azure.cosmos.models.ModelBridgeInternal;
 import com.azure.cosmos.models.PartitionKey;
 import com.azure.cosmos.models.PartitionKeyDefinition;
 import com.azure.cosmos.models.PartitionKeyDefinitionVersion;
@@ -52,6 +52,7 @@
 import org.testng.annotations.Test;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.util.retry.Retry;
 
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
@@ -1141,7 +1142,15 @@ List<String> insertDocumentsCore(
 
         ArrayList<Mono<CosmosItemResponse<ObjectNode>>> result = new ArrayList<>();
         for (int i = 0; i < docs.size(); i++) {
-            result.add(container.createItem(docs.get(i)));
+            result.add(
+                container.createItem(docs.get(i)).retryWhen(
+                    Retry
+                        .max(40)
+                        .filter(throwable ->
+                            throwable instanceof CosmosException
+                                && ((CosmosException) throwable).getStatusCode() == HttpConstants.StatusCodes.TOO_MANY_REQUESTS)
+                )
+            );
         }
 
         List<ObjectNode> insertedDocs = Flux.merge(

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosContainerOpenConnectionsAndInitCachesTest.java

@@ -145,7 +145,9 @@ public void openConnectionsAndInitCachesForDirectMode(boolean useAsync) throws I
 
         assertThat(provider.count()).isEqualTo(0);
         assertThat(collectionInfoByNameMap.size()).isEqualTo(0);
-        assertThat(routingMap.size()).isEqualTo(0);
+        // The partition-key-range (routing map) cache is shared across clients targeting the same
+        // service endpoint, so its size reflects every container routed to in this JVM. Assert on this
+        // container's own entry (populated by openConnectionsAndInitCaches) instead of the total size.
         assertThat(ReflectionUtils.isInitialized(asyncContainer).get()).isFalse();
 
         // Calling it twice to make sure no side effect of second time no-op call
@@ -158,7 +160,8 @@ public void openConnectionsAndInitCachesForDirectMode(boolean useAsync) throws I
         }
 
         assertThat(collectionInfoByNameMap.size()).isEqualTo(1);
-        assertThat(routingMap.size()).isEqualTo(1);
+        String collectionRid = asyncContainer.read().block().getProperties().getResourceId();
+        assertThat(routingMap).containsKey(collectionRid);
         assertThat(ReflectionUtils.isInitialized(asyncContainer).get()).isTrue();
 
         GlobalAddressResolver globalAddressResolver = ReflectionUtils.getGlobalAddressResolver(rxDocumentClient);
@@ -229,11 +232,12 @@ public void openConnectionsAndInitCachesForGatewayMode(boolean useAsync) {
 
         RxDocumentClientImpl rxDocumentClient = (RxDocumentClientImpl) asyncClient.getDocClientWrapper();
 
-        ConcurrentHashMap<String, ?> routingMap = getRoutingMap(rxDocumentClient);
         ConcurrentHashMap<String, ?> collectionInfoByNameMap = getCollectionInfoByNameMap(rxDocumentClient);
 
         assertThat(collectionInfoByNameMap.size()).isEqualTo(0);
-        assertThat(routingMap.size()).isEqualTo(0);
+        // The routing-map (partition-key-range) cache is shared per service endpoint; gateway-mode
+        // openConnectionsAndInitCaches does not populate it, but a sibling direct-mode test may have,
+        // so the shared routing map's size is not asserted here.
         assertThat(ReflectionUtils.isInitialized(asyncContainer).get()).isFalse();
 
         // Verifying no error when initializeContainer called on gateway mode
@@ -247,7 +251,6 @@ public void openConnectionsAndInitCachesForGatewayMode(boolean useAsync) {
         }
 
         assertThat(collectionInfoByNameMap.size()).isEqualTo(0);
-        assertThat(routingMap.size()).isEqualTo(0);
         assertThat(ReflectionUtils.isInitialized(asyncContainer).get()).isTrue();
     }
 

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosDiagnosticsTest.java

@@ -897,7 +897,11 @@ private void validateDirectModeDiagnosticsOnSuccess(
         assertThat(diagnostics).contains("gatewayStatisticsList");
         assertThat(diagnostics).contains("addressResolutionStatistics");
         assertThat(diagnostics).contains("\"metaDataName\":\"CONTAINER_LOOK_UP\"");
-        assertThat(diagnostics).contains("\"metaDataName\":\"PARTITION_KEY_RANGE_LOOK_UP\"");
+        // With the shared partition-key-range cache, a sibling client/test targeting the same service
+        // endpoint may have already populated this container's routing map. When that happens this client
+        // serves the partition-key-range lookup from the shared cache without issuing a /pkranges network
+        // request, so no PARTITION_KEY_RANGE_LOOK_UP metadata diagnostic is recorded for this operation.
+        // Its presence is therefore not asserted here.
         assertThat(diagnostics).contains("\"metaDataName\":\"SERVER_ADDRESS_LOOKUP\"");
         assertThat(diagnostics).contains("\"serializationType\":\"PARTITION_KEY_FETCH_SERIALIZATION\"");
         assertThat(diagnostics).contains("\"userAgent\":\"" + userAgent + "\"");

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosItemTest.java

@@ -176,7 +176,7 @@ public void readItemWithVeryLargePartitionKey() throws Exception {
         }
         docDefinition.set("mypk", sb.toString());
 
-        CosmosItemResponse<InternalObjectNode> itemResponse = container.createItem(docDefinition);
+        container.createItem(docDefinition);
 
         waitIfNeededForReplicasToCatchUp(getClientBuilder());
 
@@ -191,7 +191,9 @@ public void readItemWithVeryLargePartitionKey() throws Exception {
     @Test(groups = { "fast" }, timeOut = TIMEOUT)
     public void readItem() throws Exception {
         InternalObjectNode properties = getDocumentDefinition(UUID.randomUUID().toString());
-        CosmosItemResponse<InternalObjectNode> itemResponse = container.createItem(properties);
+        container.createItem(properties);
+
+        waitIfNeededForReplicasToCatchUp(getClientBuilder());
 
         CosmosItemResponse<InternalObjectNode> readResponse1 = container.readItem(properties.getId(),
                                                                                     new PartitionKey(properties.get("mypk")),
@@ -291,6 +293,8 @@ public void readMany() throws Exception {
             idSet.add(document.getId());
         }
 
+        waitIfNeededForReplicasToCatchUp(getClientBuilder());
+
         FeedResponse<InternalObjectNode> feedResponse = container.readMany(cosmosItemIdentities, InternalObjectNode.class);
 
         assertThat(feedResponse).isNotNull();
@@ -565,6 +569,7 @@ public void readManyWithPojoAndSingleTuple() throws Exception {
         CosmosItemIdentity cosmosItemIdentity = new CosmosItemIdentity(partitionKey, document.getId());
         cosmosItemIdentities.add(cosmosItemIdentity);
 
+        waitIfNeededForReplicasToCatchUp(getClientBuilder());
         FeedResponse<SampleType> feedResponse = container.readMany(cosmosItemIdentities, SampleType.class);
 
         assertThat(feedResponse.getResults()).isNotNull();
@@ -1078,6 +1083,8 @@ public void readItemWithSoftTimeoutAndFallback() throws Exception {
         ObjectNode fallBackProperties = getDocumentDefinition("justFallback", "justFallback");
         container.createItem(properties);
 
+        waitIfNeededForReplicasToCatchUp(getClientBuilder());
+
         String successfulResponse = wrapWithSoftTimeoutAndFallback(
             container
                 .asyncContainer
@@ -1334,6 +1341,8 @@ public void queryItemsWithEventualConsistency() throws Exception{
             ObjectNode properties = getDocumentDefinition(idAndPkValue, idAndPkValue);
             CosmosItemResponse<ObjectNode> itemResponse = container.createItem(properties);
 
+            waitIfNeededForReplicasToCatchUp(this.getClientBuilder());
+
             String query = String.format("SELECT * from c where c.id = '%s'", idAndPkValue);
             CosmosQueryRequestOptions cosmosQueryRequestOptions =
                 new CosmosQueryRequestOptions()
@@ -1449,10 +1458,12 @@ public void queryItemsWithContinuationTokenAndPageSize() throws Exception{
     public void readAllItemsOfLogicalPartition() throws Exception{
         String pkValue = UUID.randomUUID().toString();
         ObjectNode properties = getDocumentDefinition(UUID.randomUUID().toString(), pkValue);
-        CosmosItemResponse<ObjectNode> itemResponse = container.createItem(properties);
+        container.createItem(properties);
 
         CosmosQueryRequestOptions cosmosQueryRequestOptions = new CosmosQueryRequestOptions();
 
+        waitIfNeededForReplicasToCatchUp(getClientBuilder());
+
         CosmosPagedIterable<ObjectNode> feedResponseIterator1 =
             container.readAllItems(
                 new PartitionKey(pkValue),
@@ -1482,6 +1493,8 @@ public void readAllItemsOfLogicalPartitionWithContinuationTokenAndPageSize() thr
         properties = getDocumentDefinition(UUID.randomUUID().toString(), pkValue);
         container.createItem(properties);
 
+        waitIfNeededForReplicasToCatchUp(getClientBuilder());
+
         CosmosQueryRequestOptions cosmosQueryRequestOptions = new CosmosQueryRequestOptions();
         String continuationToken = null;
         int pageSize = 1;

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ProactiveConnectionManagementTest.java

@@ -243,7 +243,11 @@ public void openConnectionsAndInitCachesWithContainer(ProactiveConnectionManagem
 
             assertThat(provider.count()).isEqualTo(endpoints.size());
             assertThat(collectionInfoByNameMap.size()).isEqualTo(cosmosContainerIdentities.size());
-            assertThat(routingMap.size()).isEqualTo(cosmosContainerIdentities.size());
+            // The partition-key-range (routing map) cache is shared across clients targeting the same
+            // service endpoint, so its size reflects every container routed to in this JVM. Assert on
+            // this container's own entry instead of the total size.
+            String collectionRid = cosmosAsyncContainer.read().block().getProperties().getResourceId();
+            assertThat(routingMap).containsKey(collectionRid);
         } finally {
             if (cosmosAsyncContainer != null) {
                 safeDeleteCollection(cosmosAsyncContainer);
@@ -408,7 +412,14 @@ public void openConnectionsAndInitCachesWithCosmosClient_And_PerContainerConnect
 
             assertThat(provider.count()).isEqualTo(endpoints.size());
             assertThat(collectionInfoByNameMap.size()).isEqualTo(cosmosContainerIdentities.size());
-            assertThat(routingMap.size()).isEqualTo(cosmosContainerIdentities.size());
+            // The partition-key-range (routing map) cache is shared across clients targeting the same
+            // service endpoint, so its size reflects every container routed to in this JVM. Assert that
+            // each proactively-initialized container's own entry is present instead of the total size.
+            for (CosmosAsyncContainer proactivelyInitializedContainer : asyncContainers) {
+                String collectionRid =
+                    proactivelyInitializedContainer.read().block().getProperties().getResourceId();
+                assertThat(routingMap).containsKey(collectionRid);
+            }
 
             int totalConnectionCountForAllEndpoints = 0;
 
@@ -560,7 +571,14 @@ public void openConnectionsAndInitCachesWithCosmosClient_And_PerContainerConnect
 
             assertThat(provider.count()).isEqualTo(endpoints.size());
             assertThat(collectionInfoByNameMap.size()).isEqualTo(cosmosContainerIdentities.size());
-            assertThat(routingMap.size()).isEqualTo(cosmosContainerIdentities.size());
+            // The partition-key-range (routing map) cache is shared across clients targeting the same
+            // service endpoint, so its size reflects every container routed to in this JVM. Assert that
+            // each proactively-initialized container's own entry is present instead of the total size.
+            for (CosmosAsyncContainer proactivelyInitializedContainer : asyncContainers) {
+                String collectionRid =
+                    proactivelyInitializedContainer.read().block().getProperties().getResourceId();
+                assertThat(routingMap).containsKey(collectionRid);
+            }
 
             int totalConnectionCountForAllEndpoints = 0;
 
@@ -737,7 +755,14 @@ public void openConnectionsAndInitCachesWithCosmosClient_And_PerContainerConnect
 
             assertThat(provider.count()).isEqualTo(endpoints.size());
             assertThat(collectionInfoByNameMap.size()).isEqualTo(cosmosContainerIdentities.size());
-            assertThat(routingMap.size()).isEqualTo(cosmosContainerIdentities.size());
+            // The partition-key-range (routing map) cache is shared across clients targeting the same
+            // service endpoint, so its size reflects every container routed to in this JVM. Assert that
+            // each proactively-initialized container's own entry is present instead of the total size.
+            for (CosmosAsyncContainer proactivelyInitializedContainer : asyncContainers) {
+                String collectionRid =
+                    proactivelyInitializedContainer.read().block().getProperties().getResourceId();
+                assertThat(routingMap).containsKey(collectionRid);
+            }
 
             int totalConnectionCountForAllEndpoints = 0;