AddRetryCountConfigForOpenAsync (Azure#24209)

* add retry count config for open async Co-authored-by: annie-mac <[email protected]>
AlonsoMondal · Sep 22, 2021 · 75e9d94 · 75e9d94
1 parent 35e6766
commit 75e9d94
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 8 deletions.
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java
@@ -4,6 +4,7 @@
 
 import com.azure.core.util.Context;
 import com.azure.cosmos.implementation.AsyncDocumentClient;
+import com.azure.cosmos.implementation.Configs;
 import com.azure.cosmos.implementation.Constants;
 import com.azure.cosmos.implementation.CosmosPagedFluxOptions;
 import com.azure.cosmos.implementation.Document;
@@ -452,6 +453,7 @@ public <T> CosmosPagedFlux<T> queryItems(String query, Class<T> classType) {
      */
     @Beta(value = Beta.SinceVersion.V4_14_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
     public Mono<Void> openConnectionsAndInitCaches() {
+        int retryCount = Configs.getOpenConnectionsRetriesCount();
 
         if(isInitialized.compareAndSet(false, true)) {
             return this.getFeedRanges().flatMap(feedRanges -> {
@@ -460,13 +462,17 @@ public Mono<Void> openConnectionsAndInitCaches() {
                 querySpec.setQueryText("select * from c where c.id = @id");
                 querySpec.setParameters(Collections.singletonList(new SqlParameter("@id",
                     UUID.randomUUID().toString())));
-                for (FeedRange feedRange : feedRanges) {
-                    CosmosQueryRequestOptions options = new CosmosQueryRequestOptions();
-                    options.setFeedRange(feedRange);
-                    CosmosPagedFlux<ObjectNode> cosmosPagedFlux = this.queryItems(querySpec, options,
-                        ObjectNode.class);
-                    fluxList.add(cosmosPagedFlux.byPage());
+
+                for (int i = 0; i < retryCount; i++) {
+                    for (FeedRange feedRange : feedRanges) {
+                        CosmosQueryRequestOptions options = new CosmosQueryRequestOptions();
+                        options.setFeedRange(feedRange);
+                        CosmosPagedFlux<ObjectNode> cosmosPagedFlux = this.queryItems(querySpec, options,
+                            ObjectNode.class);
+                        fluxList.add(cosmosPagedFlux.byPage());
+                    }
                 }
+
                 Mono<List<FeedResponse<ObjectNode>>> listMono = Flux.merge(fluxList).collectList();
                 return listMono.flatMap(objects -> Mono.empty());
             });

diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java
@@ -92,9 +92,13 @@ public class Configs {
     private static final int DEFAULT_SESSION_TOKEN_MISMATCH_MAXIMUM_BACKOFF_TIME_IN_MILLISECONDS = 50;
 
     // Whether to process the response on a different thread
-    private static final String DEFAULT_SWITCH_OFF_IO_THREAD_FOR_RESPONSE_NAME = "COSMOS.SWITCH_OFF_IO_THREAD_FOR_RESPONSE";
+    private static final String SWITCH_OFF_IO_THREAD_FOR_RESPONSE_NAME = "COSMOS.SWITCH_OFF_IO_THREAD_FOR_RESPONSE";
     private static final boolean DEFAULT_SWITCH_OFF_IO_THREAD_FOR_RESPONSE = false;
 
+    // OpenConnectionsAndInitCaches Constants
+    private static final String OPEN_CONNECTIONS_RETRIES_COUNT_NAME = "COSMOS.OPEN_CONNECTIONS_RETRIES_COUNT";
+    private static final int DEFAULT_OPEN_CONNECTIONS_RETRIES_COUNT = 1;
+
     public Configs() {
         this.sslContext = sslContextInit();
     }
@@ -258,10 +262,16 @@ public static int getSessionTokenMismatchMaximumBackoffTimeInMs() {
 
     public static boolean shouldSwitchOffIOThreadForResponse() {
         return getJVMConfigAsBoolean(
-            DEFAULT_SWITCH_OFF_IO_THREAD_FOR_RESPONSE_NAME,
+            SWITCH_OFF_IO_THREAD_FOR_RESPONSE_NAME,
             DEFAULT_SWITCH_OFF_IO_THREAD_FOR_RESPONSE);
     }
 
+    public static int getOpenConnectionsRetriesCount() {
+        return getJVMConfigAsInt(
+            OPEN_CONNECTIONS_RETRIES_COUNT_NAME,
+            DEFAULT_OPEN_CONNECTIONS_RETRIES_COUNT);
+    }
+
     private static int getJVMConfigAsInt(String propName, int defaultValue) {
         String propValue = System.getProperty(propName);
         return getIntValue(propValue, defaultValue);