[Internal] Replica Validation: Refactors Code to Address Follow Ups

Microsoft.Azure.Cosmos/src/ConnectionPolicy.cs

@@ -50,6 +50,7 @@ public ConnectionPolicy()
             this.EnableReadRequestsFallback = null;
             this.EnableClientTelemetry = ClientTelemetryOptions.IsClientTelemetryEnabled();
             this.ServerCertificateCustomValidationCallback = null;
+            this.EnableReplicaValidation = false;
         }
 
         /// <summary>
@@ -459,6 +460,18 @@ public Func<HttpClient> HttpClientFactory
             set;
         }
 
+        /// <summary>
+        /// Gets or sets the flag to enable replica validation.
+        /// </summary>
+        /// <value>
+        /// The default value is false
+        /// </value>
+        public bool EnableReplicaValidation
+        {
+            get;
+            set;
+        }
+
         /// <summary>
         /// (Direct/TCP) This is an advanced setting that controls the number of TCP connections that will be opened eagerly to each Cosmos DB back-end.
         /// </summary>

Microsoft.Azure.Cosmos/src/CosmosClientOptions.cs

@@ -346,6 +346,23 @@ public ConnectionMode ConnectionMode
         /// <seealso cref="TransactionalBatchItemRequestOptions.EnableContentResponseOnWrite"/>
         public bool? EnableContentResponseOnWrite { get; set; }
 
+        /// <summary>
+        /// Gets or sets the prioritize healthy replicas flag.
+        /// Prioritizing healthy replicas helps the cosmos client to become more
+        /// resilient to connection timeouts, by choosing a healthy replica over an
+        /// unhealthy one. The default value for this parameter is false.
+        /// </summary>
+        /// <remarks>
+        /// <para>This is optimal for workloads where latency spikes are critical during upgrades.</para>
+        /// </remarks>
+        /// <seealso cref="CosmosClientBuilder.WithPrioritizeHealthyReplicas(bool)"/>
+#if PREVIEW
+        public
+#else
+        internal
+#endif
+        bool PrioritizeHealthyReplicas { get; set; }
+
         /// <summary>
         /// (Direct/TCP) Controls the amount of idle time after which unused connections are closed.
         /// </summary>
@@ -758,6 +775,7 @@ internal virtual ConnectionPolicy GetConnectionPolicy(int clientId)
                 EnablePartitionLevelFailover = this.EnablePartitionLevelFailover,
                 PortReuseMode = this.portReuseMode,
                 EnableTcpConnectionEndpointRediscovery = this.EnableTcpConnectionEndpointRediscovery,
+                EnableReplicaValidation = this.PrioritizeHealthyReplicas,
                 HttpClientFactory = this.httpClientFactory,
                 ServerCertificateCustomValidationCallback = this.ServerCertificateCustomValidationCallback
             };

Microsoft.Azure.Cosmos/src/DocumentClient.cs

@@ -6700,7 +6700,8 @@ private void CreateStoreModel(bool subscribeRntbdStatus)
                 this.ConnectionPolicy.EnableReadRequestsFallback ?? (this.accountServiceConfiguration.DefaultConsistencyLevel != Documents.ConsistencyLevel.BoundedStaleness),
                 !this.enableRntbdChannel,
                 this.UseMultipleWriteLocations && (this.accountServiceConfiguration.DefaultConsistencyLevel != Documents.ConsistencyLevel.Strong),
-                true);
+                true,
+                this.ConnectionPolicy.EnableReplicaValidation);
 
             if (subscribeRntbdStatus)
             {

Microsoft.Azure.Cosmos/src/Fluent/CosmosClientBuilder.cs

@@ -621,6 +621,26 @@ public CosmosClientBuilder WithContentResponseOnWrite(bool contentResponseOnWrit
             return this;
         }
 
+        /// <summary>
+        /// Gets or sets the prioritize healthy replicas flag.
+        /// Prioritizing healthy replicas helps the cosmos client to become more
+        /// resilient to connection timeouts, by choosing a healthy replica over an
+        /// unhealthy one. The default value for this parameter is false.
+        /// </summary>
+        /// <param name="replicaValidationEnabled">a boolean flag indicating if the feature will be enabled.</param>
+        /// <returns>The <see cref="CosmosClientBuilder"/> object</returns>
+#if PREVIEW
+        public
+#else
+        internal
+#endif
+        CosmosClientBuilder WithPrioritizeHealthyReplicas(
+            bool replicaValidationEnabled)
+        {
+            this.clientOptions.PrioritizeHealthyReplicas = replicaValidationEnabled;
+            return this;
+        }
+
         /// <summary>
         /// The event handler to be invoked before the request is sent.
         /// </summary>

Microsoft.Azure.Cosmos/src/Routing/GatewayAddressCache.cs

@@ -47,6 +47,7 @@ internal class GatewayAddressCache : IAddressCache, IDisposable
 
         private readonly CosmosHttpClient httpClient;
         private readonly bool isReplicaAddressValidationEnabled;
+        private static readonly TimeSpan WarmupCacheAndOpenConnectionTimeout = TimeSpan.FromMinutes(40);
 
         private Tuple<PartitionKeyRangeIdentity, PartitionAddressInformation> masterPartitionAddressCache;
         private DateTime suboptimalMasterPartitionTimestamp;
@@ -61,7 +62,8 @@ public GatewayAddressCache(
             CosmosHttpClient httpClient,
             IOpenConnectionsHandler openConnectionsHandler,
             long suboptimalPartitionForceRefreshIntervalInSeconds = 600,
-            bool enableTcpConnectionEndpointRediscovery = false)
+            bool enableTcpConnectionEndpointRediscovery = false,
+            bool replicaAddressValidationEnabled = false)
         {
             this.addressEndpoint = new Uri(serviceEndpoint + "/" + Paths.AddressPathSegment);
             this.protocol = protocol;
@@ -85,9 +87,7 @@ public GatewayAddressCache(
                 GatewayAddressCache.ProtocolString(this.protocol));
 
             this.openConnectionsHandler = openConnectionsHandler;
-            this.isReplicaAddressValidationEnabled = Helpers.GetEnvironmentVariable<bool>(
-                name: Constants.EnvironmentVariables.ReplicaConnectivityValidationEnabled,
-                defaultValue: false);
+            this.isReplicaAddressValidationEnabled = replicaAddressValidationEnabled;
         }
 
         public Uri ServiceEndpoint => this.serviceEndpoint;
@@ -113,7 +113,7 @@ public async Task OpenConnectionsAsync(
             bool shouldOpenRntbdChannels,
             CancellationToken cancellationToken)
         {
-            List<Task<TryCatch<DocumentServiceResponse>>> tasks = new ();
+            List<Task> tasks = new ();
             int batchSize = GatewayAddressCache.DefaultBatchSize;
 
 #if !(NETSTANDARD15 || NETSTANDARD16)
@@ -147,50 +147,29 @@ public async Task OpenConnectionsAsync(
             {
                 for (int i = 0; i < partitionKeyRangeIdentities.Count; i += batchSize)
                 {
-                    tasks
-                        .Add(this.GetAddressesAsync(
-                            request: request,
-                            collectionRid: collection.ResourceId,
-                            partitionKeyRangeIds: partitionKeyRangeIdentities.Skip(i).Take(batchSize).Select(range => range.PartitionKeyRangeId)));
+                    tasks.Add(
+                        this.WarmupCachesAndOpenConnectionsAsync(
+                                request: request,
+                                collectionRid: collection.ResourceId,
+                                partitionKeyRangeIds: partitionKeyRangeIdentities.Skip(i).Take(batchSize).Select(range => range.PartitionKeyRangeId),
+                                containerProperties: collection,
+                                shouldOpenRntbdChannels: shouldOpenRntbdChannels));
                 }
             }
 
-            foreach (TryCatch<DocumentServiceResponse> task in await Task.WhenAll(tasks))
-            {
-                if (task.Failed)
-                {
-                    continue;
-                }
-
-                using (DocumentServiceResponse response = task.Result)
-                {
-                    FeedResource<Address> addressFeed = response.GetResource<FeedResource<Address>>();
-
-                    bool inNetworkRequest = this.IsInNetworkRequest(response);
-
-                    IEnumerable<Tuple<PartitionKeyRangeIdentity, PartitionAddressInformation>> addressInfos =
-                        addressFeed.Where(addressInfo => ProtocolFromString(addressInfo.Protocol) == this.protocol)
-                            .GroupBy(address => address.PartitionKeyRangeId, StringComparer.Ordinal)
-                            .Select(group => this.ToPartitionAddressAndRange(collection.ResourceId, @group.ToList(), inNetworkRequest));
-
-                    foreach (Tuple<PartitionKeyRangeIdentity, PartitionAddressInformation> addressInfo in addressInfos)
-                    {
-                        this.serverPartitionAddressCache.Set(
-                            new PartitionKeyRangeIdentity(collection.ResourceId, addressInfo.Item1.PartitionKeyRangeId),
-                            addressInfo.Item2);
+            using CancellationTokenSource linkedToken = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+            Task timeoutTask = Task.Delay(GatewayAddressCache.WarmupCacheAndOpenConnectionTimeout, linkedToken.Token);
+            Task resultTask = await Task.WhenAny(Task.WhenAll(tasks), timeoutTask);
 
-                        // The `shouldOpenRntbdChannels` boolean flag indicates whether the SDK should establish Rntbd connections to the
-                        // backend replica nodes. For the `CosmosClient.CreateAndInitializeAsync()` flow, the flag should be passed as
-                        // `true` so that the Rntbd connections to the backend replicas could be established deterministically. For any
-                        // other flow, the flag should be passed as `false`.
-                        if (this.openConnectionsHandler != null && shouldOpenRntbdChannels)
-                        {
-                            await this.openConnectionsHandler
-                                .TryOpenRntbdChannelsAsync(
-                                    addresses: addressInfo.Item2.Get(Protocol.Tcp)?.ReplicaTransportAddressUris);
-                        }
-                    }
-                }
+            if (resultTask == timeoutTask)
+            {
+                // Operation cancelled.
+                DefaultTrace.TraceWarning("The open connection task was cancelled because the cancellation token was expired. '{0}'",
+                    System.Diagnostics.Trace.CorrelationManager.ActivityId);
+            }
+            else
+            {
+                linkedToken.Cancel();
             }
         }
 
@@ -350,6 +329,71 @@ public async Task<PartitionAddressInformation> TryGetAddressesAsync(
             }
         }
 
+        private async Task WarmupCachesAndOpenConnectionsAsync(
+            DocumentServiceRequest request,
+            string collectionRid,
+            IEnumerable<string> partitionKeyRangeIds,
+            ContainerProperties containerProperties,
+            bool shouldOpenRntbdChannels)
+        {
+            TryCatch<DocumentServiceResponse> documentServiceResponseWrapper = await this.GetAddressesAsync(
+                                request: request,
+                                collectionRid: collectionRid,
+                                partitionKeyRangeIds: partitionKeyRangeIds);
+
+            if (documentServiceResponseWrapper.Failed)
+            {
+                return;
+            }
+
+            try
+            {
+                using (DocumentServiceResponse response = documentServiceResponseWrapper.Result)
+                {
+                    FeedResource<Address> addressFeed = response.GetResource<FeedResource<Address>>();
+
+                    bool inNetworkRequest = this.IsInNetworkRequest(response);
+
+                    IEnumerable<Tuple<PartitionKeyRangeIdentity, PartitionAddressInformation>> addressInfos =
+                        addressFeed.Where(addressInfo => ProtocolFromString(addressInfo.Protocol) == this.protocol)
+                            .GroupBy(address => address.PartitionKeyRangeId, StringComparer.Ordinal)
+                            .Select(group => this.ToPartitionAddressAndRange(containerProperties.ResourceId, @group.ToList(), inNetworkRequest));
+
+                    List<Task> openConnectionTasks = new ();
+                    foreach (Tuple<PartitionKeyRangeIdentity, PartitionAddressInformation> addressInfo in addressInfos)
+                    {
+                        this.serverPartitionAddressCache.Set(
+                            new PartitionKeyRangeIdentity(containerProperties.ResourceId, addressInfo.Item1.PartitionKeyRangeId),
+                            addressInfo.Item2);
+
+                        // The `shouldOpenRntbdChannels` boolean flag indicates whether the SDK should establish Rntbd connections to the
+                        // backend replica nodes. For the `CosmosClient.CreateAndInitializeAsync()` flow, the flag should be passed as
+                        // `true` so that the Rntbd connections to the backend replicas could be established deterministically. For any
+                        // other flow, the flag should be passed as `false`.
+                        if (this.openConnectionsHandler != null && shouldOpenRntbdChannels)
+                        {
+                            openConnectionTasks
+                                .Add(this.openConnectionsHandler
+                                    .TryOpenRntbdChannelsAsync(
+                                        addresses: addressInfo.Item2.Get(Protocol.Tcp)?.ReplicaTransportAddressUris));
+                        }
+                    }
+
+                    if (openConnectionTasks.Any())
+                    {
+                        await Task.WhenAll(openConnectionTasks);
+                    }
+                }
+            }
+            catch (Exception ex)
+            {
+                DefaultTrace.TraceWarning("Failed to open connections warm-up cache for the server addresses: {0} with exception: {1}. '{2}'",
+                    collectionRid,
+                    ex,
+                    System.Diagnostics.Trace.CorrelationManager.ActivityId);
+            }
+        }
+
         private static void SetTransportAddressUrisToUnhealthy(
             PartitionAddressInformation stalePartitionAddressInformation,
             Lazy<HashSet<TransportAddressUri>> failedEndpoints)

Microsoft.Azure.Cosmos/src/Routing/GlobalAddressResolver.cs

@@ -39,6 +39,7 @@ internal sealed class GlobalAddressResolver : IAddressResolverExtension, IDispos
         private readonly CosmosHttpClient httpClient;
         private readonly ConcurrentDictionary<Uri, EndpointCache> addressCacheByEndpoint;
         private readonly bool enableTcpConnectionEndpointRediscovery;
+        private readonly bool replicaAddressValidationEnabled;
         private IOpenConnectionsHandler openConnectionsHandler;
 
         public GlobalAddressResolver(
@@ -66,6 +67,7 @@ public GlobalAddressResolver(
                 ? GlobalAddressResolver.MaxBackupReadRegions : 0;
 
             this.enableTcpConnectionEndpointRediscovery = connectionPolicy.EnableTcpConnectionEndpointRediscovery;
+            this.replicaAddressValidationEnabled = connectionPolicy.EnableReplicaValidation;
 
             this.maxEndpoints = maxBackupReadEndpoints + 2; // for write and alternate write endpoint (during failover)
 
@@ -281,7 +283,8 @@ private EndpointCache GetOrAddEndpoint(Uri endpoint)
                         this.serviceConfigReader,
                         this.httpClient,
                         this.openConnectionsHandler,
-                        enableTcpConnectionEndpointRediscovery: this.enableTcpConnectionEndpointRediscovery);
+                        enableTcpConnectionEndpointRediscovery: this.enableTcpConnectionEndpointRediscovery,
+                        replicaAddressValidationEnabled: this.replicaAddressValidationEnabled);
 
                     string location = this.endpointManager.GetLocation(endpoint);
                     AddressResolver addressResolver = new AddressResolver(null, new NullRequestSigner(), location);

Microsoft.Azure.Cosmos/src/direct/ConsistencyReader.cs

@@ -144,12 +144,13 @@ public ConsistencyReader(
             ISessionContainer sessionContainer,
             TransportClient transportClient,
             IServiceConfigurationReader serviceConfigReader,
-            IAuthorizationTokenProvider authorizationTokenProvider)
+            IAuthorizationTokenProvider authorizationTokenProvider,
+            bool enableReplicaValidation)
         {
             this.addressSelector = addressSelector;
             this.serviceConfigReader = serviceConfigReader;
             this.authorizationTokenProvider = authorizationTokenProvider;
-            this.storeReader = new StoreReader(transportClient, addressSelector, new AddressEnumerator(), sessionContainer);
+            this.storeReader = new StoreReader(transportClient, addressSelector, new AddressEnumerator(), sessionContainer, enableReplicaValidation);
             this.quorumReader = new QuorumReader(transportClient, addressSelector, this.storeReader, serviceConfigReader, authorizationTokenProvider);
         }
 

Microsoft.Azure.Cosmos/src/direct/ConsistencyWriter.cs

@@ -59,7 +59,8 @@ public ConsistencyWriter(
             TransportClient transportClient,
             IServiceConfigurationReader serviceConfigReader,
             IAuthorizationTokenProvider authorizationTokenProvider,
-            bool useMultipleWriteLocations)
+            bool useMultipleWriteLocations,
+            bool enableReplicaValidation)
         {
             this.transportClient = transportClient;
             this.addressSelector = addressSelector;
@@ -71,7 +72,8 @@ public ConsistencyWriter(
                                     transportClient,
                                     addressSelector,
                                     new AddressEnumerator(),
-                                    sessionContainer: null); //we need store reader only for global strong, no session is needed*/
+                                    sessionContainer: null,
+                                    enableReplicaValidation); //we need store reader only for global strong, no session is needed*/
         }
 
         // Test hook

Microsoft.Azure.Cosmos/src/direct/Constants.cs

@@ -2134,7 +2134,6 @@ public static class EnvironmentVariables
         {
             public const string SocketOptionTcpKeepAliveIntervalName = "AZURE_COSMOS_TCP_KEEPALIVE_INTERVAL_SECONDS";
             public const string SocketOptionTcpKeepAliveTimeName = "AZURE_COSMOS_TCP_KEEPALIVE_TIME_SECONDS";
-            public const string ReplicaConnectivityValidationEnabled = "AZURE_COSMOS_REPLICA_VALIDATION_ENABLED";
             public const string AggressiveTimeoutDetectionEnabled = "AZURE_COSMOS_AGGRESSIVE_TIMEOUT_DETECTION_ENABLED";
             public const string TimeoutDetectionTimeLimit = "AZURE_COSMOS_TIMEOUT_DETECTION_TIME_LIMIT_IN_SECONDS";
             public const string TimeoutDetectionOnWriteThreshold = "AZURE_COSMOS_TIMEOUT_DETECTION_ON_WRITE_THRESHOLD";

Microsoft.Azure.Cosmos/src/direct/IStoreClientFactory.cs

@@ -17,6 +17,7 @@ StoreClient CreateStoreClient(
             bool enableReadRequestsFallback = false,
             bool useFallbackClient = true,
             bool useMultipleWriteLocations = false,
-            bool detectClientConnectivityIssues = false);
+            bool detectClientConnectivityIssues = false,
+            bool enableReplicaValidation = false);
     }
 }

Microsoft.Azure.Cosmos/src/direct/ReplicatedResourceClient.cs

@@ -68,6 +68,7 @@ public ReplicatedResourceClient(
             bool useMultipleWriteLocations,
             bool detectClientConnectivityIssues,
             bool disableRetryWithRetryPolicy,
+            bool enableReplicaValidation,
             RetryWithConfiguration retryWithConfiguration = null)
         {
             this.addressResolver = addressResolver;
@@ -86,14 +87,16 @@ public ReplicatedResourceClient(
                 sessionContainer,
                 transportClient,
                 serviceConfigReader,
-                authorizationTokenProvider);
+                authorizationTokenProvider,
+                enableReplicaValidation);
             this.consistencyWriter = new ConsistencyWriter(
                 this.addressSelector,
                 sessionContainer,
                 transportClient,
                 serviceConfigReader,
                 authorizationTokenProvider,
-                useMultipleWriteLocations);
+                useMultipleWriteLocations,
+                enableReplicaValidation);
             this.enableReadRequestsFallback = enableReadRequestsFallback;
             this.useMultipleWriteLocations = useMultipleWriteLocations;
             this.detectClientConnectivityIssues = detectClientConnectivityIssues;