Open connections through rntbd (Azure#28470)

* openAsync through rntbd

Co-authored-by: annie-mac <annie-mac@annie-macs-MacBook-Pro.local>
Co-authored-by: annie-mac <annie-mac@XBX-2505-B09E.redmond.corp.microsoft.com>
Co-authored-by: annie-mac <annie-mac@yindeng2019.fareast.corp.microsoft.com>

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java

@@ -4,7 +4,6 @@
 
 import com.azure.core.util.Context;
 import com.azure.cosmos.implementation.AsyncDocumentClient;
-import com.azure.cosmos.implementation.Configs;
 import com.azure.cosmos.implementation.CosmosPagedFluxOptions;
 import com.azure.cosmos.implementation.CosmosSchedulers;
 import com.azure.cosmos.implementation.Document;
@@ -14,6 +13,7 @@
 import com.azure.cosmos.implementation.InternalObjectNode;
 import com.azure.cosmos.implementation.ItemDeserializer;
 import com.azure.cosmos.implementation.Offer;
+import com.azure.cosmos.implementation.OpenConnectionResponse;
 import com.azure.cosmos.implementation.OperationType;
 import com.azure.cosmos.implementation.Paths;
 import com.azure.cosmos.implementation.RequestOptions;
@@ -25,7 +25,6 @@
 import com.azure.cosmos.implementation.batch.BulkExecutor;
 import com.azure.cosmos.implementation.feedranges.FeedRangeEpkImpl;
 import com.azure.cosmos.implementation.feedranges.FeedRangeInternal;
-import com.azure.cosmos.implementation.query.QueryInfo;
 import com.azure.cosmos.implementation.routing.Range;
 import com.azure.cosmos.implementation.throughputControl.config.GlobalThroughputControlGroup;
 import com.azure.cosmos.implementation.throughputControl.config.LocalThroughputControlGroup;
@@ -52,23 +51,19 @@
 import com.azure.cosmos.models.FeedResponse;
 import com.azure.cosmos.models.ModelBridgeInternal;
 import com.azure.cosmos.models.PartitionKey;
-import com.azure.cosmos.models.SqlParameter;
 import com.azure.cosmos.models.SqlQuerySpec;
 import com.azure.cosmos.models.ThroughputProperties;
 import com.azure.cosmos.models.ThroughputResponse;
 import com.azure.cosmos.util.Beta;
 import com.azure.cosmos.util.CosmosPagedFlux;
 import com.azure.cosmos.util.UtilBridgeInternal;
-import com.fasterxml.jackson.databind.node.ObjectNode;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
-import java.util.ArrayList;
-import java.util.Collections;
 import java.util.List;
-import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
 
@@ -440,53 +435,71 @@ public <T> CosmosPagedFlux<T> queryItems(String query, Class<T> classType) {
         return queryItemsInternal(new SqlQuerySpec(query), new CosmosQueryRequestOptions(), classType);
     }
 
-    /**
-     * Initializes the container by warming up the caches and connections for the current read region.
+    /***
+     *  Best effort to initializes the container by warming up the caches and connections for the current read region.
      *
-     * <p><br>The execution of this method is expected to result in some RU charges to your account.
-     * The number of RU consumed by this request varies, depending on data consistency, size of the overall data in the container,
-     * item indexing, number of projections. For more information regarding RU considerations please visit
-     * <a href="https://docs.microsoft.com/en-us/azure/cosmos-db/request-units#request-unit-considerations">https://docs.microsoft.com/en-us/azure/cosmos-db/request-units#request-unit-considerations</a>.
-     * </p>
+     *  Depending on how many partitions the container has, the total time needed will also change. But generally you can use the following formula
+     *  to get an estimated time:
+     *  If it took 200ms to establish a connection, and you have 100 partitions in your container
+     *  then it will take around (100 * 4 / CPUCores) * 200ms to open all connections after get the address list
      *
-     * <p>
-     * <br>NOTE: This API ideally should be called only once during application initialization before any workload.
-     * <br>In case of any transient error, caller should consume the error and continue the regular workload.
-     * </p>
+     *  <p>
+     *  <br>NOTE: This API ideally should be called only once during application initialization before any workload.
+     *  <br>In case of any transient error, caller should consume the error and continue the regular workload.
+     *  </p>
      *
-     * @return Mono of Void
+     *  @return Mono of Void.
      */
     @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 -> {
-                List<Flux<FeedResponse<ObjectNode>>> fluxList = new ArrayList<>();
-                SqlQuerySpec querySpec = new SqlQuerySpec();
-                querySpec.setQueryText("select * from c where c.id = @id");
-                querySpec.setParameters(Collections.singletonList(new SqlParameter("@id",
-                    UUID.randomUUID().toString())));
-
-                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());
-            });
+            return withContext(context -> openConnectionsAndInitCachesInternal()
+                                            .flatMap(openResult -> {
+                                                logger.info("OpenConnectionsAndInitCaches: {}", openResult);
+                                                return Mono.empty();
+                                            }));
         } else {
-            logger.warn("openConnectionsAndInitCaches is already called once on Container {}, no operation will take place in this call", this.getId());
+            logger.warn(
+                    String.format(
+                        "OpenConnectionsAndInitCaches is already called once on Container %s, no operation will take place in this call",
+                        this.getId()));
             return Mono.empty();
         }
     }
 
+    /***
+     * Internal implementation to try to initialize the container by warming up the caches and connections for the current read region.
+     *
+     * @return a string represents the open result.
+     */
+    private Mono<String> openConnectionsAndInitCachesInternal() {
+        return this.database.getDocClientWrapper().openConnectionsAndInitCaches(getLink())
+                .collectList()
+                .flatMap(openConnectionResponses -> {
+                    // Generate a simple statistics string for open connections
+                    int total = openConnectionResponses.size();
+
+                    ConcurrentHashMap<String, Boolean> endPointOpenConnectionsStatistics = new ConcurrentHashMap<>();
+                    for (OpenConnectionResponse openConnectionResponse : openConnectionResponses) {
+                        endPointOpenConnectionsStatistics.compute(openConnectionResponse.getUri().getURIAsString(), (key, value) -> {
+                            if (value == null) {
+                                return openConnectionResponse.isConnected();
+                            }
+
+                            // Sometimes different replicas can landed on the same server, that is why we could reach here
+                            // We will only create max one connection for each endpoint in openConnectionsAndInitCaches
+                            // if one failed, one succeeded, then it is still good
+                            return openConnectionResponse.isConnected() || value;
+                        });
+                    }
+
+                    long endpointConnected = endPointOpenConnectionsStatistics.values().stream().filter(isConnected -> isConnected).count();
+                    return Mono.just(String.format(
+                            "EndpointsConnected: %s, Failed: %s", endpointConnected, endPointOpenConnectionsStatistics.size() - endpointConnected));
+                });
+    }
+
     /**
      * Query for items in the current container using a string.
      * <p>

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java

@@ -790,19 +790,13 @@ public void enableGlobalThroughputControlGroup(ThroughputControlGroupConfig grou
         this.asyncContainer.enableGlobalThroughputControlGroup(groupConfig, globalControlConfig);
     }
 
-    /**
-     * Initializes the container by warming up the caches and connections for the current read region.
-     *
-     * <p><br>The execution of this method is expected to result in some RU charges to your account.
-     * The number of RU consumed by this request varies, depending on data consistency, size of the overall data in the container,
-     * item indexing, number of projections. For more information regarding RU considerations please visit
-     * <a href="https://docs.microsoft.com/en-us/azure/cosmos-db/request-units#request-unit-considerations">https://docs.microsoft.com/en-us/azure/cosmos-db/request-units#request-unit-considerations</a>.
-     * </p>
+    /***
+     *  Initializes the container by warming up the caches and connections for the current read region.
      *
-     * <p>
-     * <br>NOTE: This API ideally should be called only once during application initialization before any workload.
-     * <br>In case of any transient error, caller should consume the error and continue the regular workload.
-     * </p>
+     *  <p>
+     *  <br>NOTE: This API ideally should be called only once during application initialization before any workload.
+     *  <br>In case of any transient error, caller should consume the error and continue the regular workload.
+     *  </p>
      *
      */
     @Beta(value = Beta.SinceVersion.V4_14_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java

@@ -1617,4 +1617,11 @@ <T> Flux<FeedResponse<T>> readAllDocuments(
      * @param group the throughput control group.
      */
     void enableThroughputControlGroup(ThroughputControlGroupInternal group);
+
+    /***
+     *  Warming up the caches and connections to all replicas of the container for the current read region.
+     *
+     * @param containerLink the container link.
+     */
+    Flux<OpenConnectionResponse> openConnectionsAndInitCaches(String containerLink);
 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java

@@ -95,10 +95,6 @@ public class Configs {
     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;
-
     // whether to allow query empty page diagnostics logging
     private static final String QUERY_EMPTY_PAGE_DIAGNOSTICS_ENABLED = "COSMOS.QUERY_EMPTY_PAGE_DIAGNOSTICS_ENABLED";
     private static final boolean DEFAULT_QUERY_EMPTY_PAGE_DIAGNOSTICS_ENABLED = false;
@@ -270,12 +266,6 @@ public static boolean shouldSwitchOffIOThreadForResponse() {
             DEFAULT_SWITCH_OFF_IO_THREAD_FOR_RESPONSE);
     }
 
-    public static int getOpenConnectionsRetriesCount() {
-        return getJVMConfigAsInt(
-            OPEN_CONNECTIONS_RETRIES_COUNT_NAME,
-            DEFAULT_OPEN_CONNECTIONS_RETRIES_COUNT);
-    }
-
     public static boolean isEmptyPageDiagnosticsEnabled() {
         return getJVMConfigAsBoolean(
             QUERY_EMPTY_PAGE_DIAGNOSTICS_ENABLED,

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/IOpenConnectionsHandler.java

@@ -0,0 +1,13 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.implementation;
+
+import com.azure.cosmos.implementation.directconnectivity.Uri;
+import reactor.core.publisher.Flux;
+
+import java.util.List;
+
+public interface IOpenConnectionsHandler {
+    Flux<OpenConnectionResponse> openConnections(List<Uri> addresses);
+}

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/OpenConnectionResponse.java

@@ -0,0 +1,34 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.implementation;
+
+import com.azure.cosmos.implementation.directconnectivity.Uri;
+
+public class OpenConnectionResponse {
+    private final boolean connected;
+    private final Throwable exception;
+    private final Uri uri;
+
+    public OpenConnectionResponse(Uri uri, boolean connected) {
+        this(uri, connected, null);
+    }
+
+    public OpenConnectionResponse(Uri uri, boolean connected, Throwable exception) {
+        this.uri = uri;
+        this.connected = connected;
+        this.exception = exception;
+    }
+
+    public boolean isConnected() {
+        return connected;
+    }
+
+    public Throwable getException() {
+        return exception;
+    }
+
+    public Uri getUri() {
+        return uri;
+    }
+}

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

@@ -4129,6 +4129,13 @@ public synchronized void enableThroughputControlGroup(ThroughputControlGroupInte
         this.throughputControlStore.enableThroughputControlGroup(group);
     }
 
+    @Override
+    public Flux<OpenConnectionResponse> openConnectionsAndInitCaches(String containerLink) {
+        checkArgument(StringUtils.isNotEmpty(containerLink), "Argument 'containerLink' should not be null nor empty");
+
+        return this.storeModel.openConnectionsAndInitCaches(containerLink);
+    }
+
     private static SqlQuerySpec createLogicalPartitionScanQuerySpec(
         PartitionKey partitionKey,
         String partitionKeySelector) {

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java

@@ -521,6 +521,11 @@ public void enableThroughputControl(ThroughputControlStore throughputControlStor
         // Disable throughput control for gateway mode
     }
 
+    @Override
+    public Flux<OpenConnectionResponse> openConnectionsAndInitCaches(String containerLink) {
+        return Flux.empty();
+    }
+
     private void captureSessionToken(RxDocumentServiceRequest request, Map<String, String> responseHeaders) {
         if (request.getResourceType() == ResourceType.DocumentCollection &&
             request.getOperationType() == OperationType.Delete) {

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxStoreModel.java

@@ -6,6 +6,7 @@
 import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple;
 import com.azure.cosmos.implementation.spark.OperationListener;
 import com.azure.cosmos.implementation.throughputControl.ThroughputControlStore;
+import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 /**
@@ -50,4 +51,12 @@ default Mono<RxDocumentServiceResponse> processMessage(RxDocumentServiceRequest
      * @param throughputControlStore
      */
     void enableThroughputControl(ThroughputControlStore throughputControlStore);
+
+    /***
+     * Open connections and init caches.
+     *
+     * @param containerLink the container link.
+     * @return a flux of {@link OpenConnectionResponse}.
+     */
+    Flux<OpenConnectionResponse> openConnectionsAndInitCaches(String containerLink);
 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/AddressResolver.java

@@ -24,12 +24,15 @@
 import com.azure.cosmos.implementation.apachecommons.lang.NotImplementedException;
 import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
 import com.azure.cosmos.implementation.caches.RxCollectionCache;
+import com.azure.cosmos.implementation.IOpenConnectionsHandler;
+import com.azure.cosmos.implementation.OpenConnectionResponse;
 import com.azure.cosmos.implementation.routing.CollectionRoutingMap;
 import com.azure.cosmos.implementation.routing.PartitionKeyInternal;
 import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper;
 import com.azure.cosmos.implementation.routing.PartitionKeyRangeIdentity;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 import java.net.URI;
@@ -89,6 +92,16 @@ public int updateAddresses(URI serverKey) {
         throw new NotImplementedException("updateAddresses() is not supported in AddressResolver");
     }
 
+    @Override
+    public Flux<OpenConnectionResponse> openConnectionsAndInitCaches(String containerLink) {
+        return Flux.empty();
+    }
+
+    @Override
+    public void setOpenConnectionsHandler(IOpenConnectionsHandler openConnectionHandler) {
+        throw new NotImplementedException("setOpenConnectionsHandler is not supported on AddressResolver");
+    }
+
     private static boolean isSameCollection(PartitionKeyRange initiallyResolved, PartitionKeyRange newlyResolved) {
         if (initiallyResolved == null) {
             throw new IllegalArgumentException("parent");

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/AddressSelector.java

@@ -6,6 +6,8 @@
 import com.azure.cosmos.implementation.GoneException;
 import com.azure.cosmos.implementation.RxDocumentServiceRequest;
 import com.azure.cosmos.implementation.Strings;
+import com.azure.cosmos.implementation.OpenConnectionResponse;
+import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 import java.util.Arrays;
@@ -82,4 +84,8 @@ public Mono<List<AddressInformation>> resolveAddressesAsync(RxDocumentServiceReq
             }
         );
     }
+
+    public Flux<OpenConnectionResponse> openConnectionsAndInitCaches(String containerLink) {
+        return this.addressResolver.openConnectionsAndInitCaches(containerLink);
+    }
 }