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 index 96d3b5203c674..b18307f32bfb1 100644 --- 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 CosmosPagedFlux queryItems(String query, Class classType) { */ @Beta(value = Beta.SinceVersion.V4_14_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) public Mono openConnectionsAndInitCaches() { + int retryCount = Configs.getOpenConnectionsRetriesCount(); if(isInitialized.compareAndSet(false, true)) { return this.getFeedRanges().flatMap(feedRanges -> { @@ -460,13 +462,17 @@ public Mono 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 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 cosmosPagedFlux = this.queryItems(querySpec, options, + ObjectNode.class); + fluxList.add(cosmosPagedFlux.byPage()); + } } + Mono>> 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 index 41078d6b54b5b..eb7f152eb17b0 100644 --- 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);