[Internal] Direct: Adds Direct stack 3.37.6 upgrade (#4961)

# Pull Request Template ## Description This PR is for upgrading Direct stack version 3.37.6 the V3 codebase with 3.37.6. It contains two changes, https://msdata.visualstudio.com/CosmosDB/_git/CosmosDB/pullrequest/1505615 : QuorumReader: Fixes bug with BoundedStaleness https://msdata.visualstudio.com/CosmosDB/_git/CosmosDB/pullrequest/1529101 ConnectionStateListener multiplexing: support multiple CosmosClients (Sharing StoreClient) It also brings in a dependent change done in OSS codebase https://msdata.visualstudio.com/CosmosDB/_git/CosmosSdkOSSClone/pullrequest/1532400?_a=files&path=/Microsoft.Azure.Cosmos/src/DocumentClient.cs ## Type of change Please delete options that are not relevant. - [x] New feature (non-breaking change which adds functionality)
Azure · Jan 16, 2025 · b4a4ac0 · b4a4ac0
1 parent dc8b9f9
commit b4a4ac0
Show file tree

Hide file tree

Showing 6 changed files with 396 additions and 43 deletions.
diff --git a/Directory.Build.props b/Directory.Build.props
@@ -3,7 +3,7 @@
 		<ClientOfficialVersion>3.46.1</ClientOfficialVersion>
 		<ClientPreviewVersion>3.47.0</ClientPreviewVersion>
 		<ClientPreviewSuffixVersion>preview.1</ClientPreviewSuffixVersion>
-		<DirectVersion>3.37.5</DirectVersion>
+		<DirectVersion>3.37.6</DirectVersion>
 		<FaultInjectionVersion>1.0.0</FaultInjectionVersion>
 		<FaultInjectionSuffixVersion>beta.0</FaultInjectionSuffixVersion>
 		<EncryptionOfficialVersion>2.0.4</EncryptionOfficialVersion>

diff --git a/Microsoft.Azure.Cosmos/src/DocumentClient.cs b/Microsoft.Azure.Cosmos/src/DocumentClient.cs
@@ -6652,17 +6652,6 @@ private bool IsValidConsistency(
 
         private void InitializeDirectConnectivity(IStoreClientFactory storeClientFactory)
         {
-            this.AddressResolver = new GlobalAddressResolver(
-                this.GlobalEndpointManager,
-                this.PartitionKeyRangeLocation,
-                this.ConnectionPolicy.ConnectionProtocol,
-                this,
-                this.collectionCache,
-                this.partitionKeyRangeCache,
-                this.accountServiceConfiguration,
-                this.ConnectionPolicy,
-                this.httpClient);
-
             // Check if we have a store client factory in input and if we do, do not initialize another store client
             // The purpose is to reuse store client factory across all document clients inside compute gateway
             if (storeClientFactory != null)
@@ -6704,7 +6693,6 @@ private void InitializeDirectConnectivity(IStoreClientFactory storeClientFactory
                     sendHangDetectionTimeSeconds: this.rntbdSendHangDetectionTimeSeconds,
                     retryWithConfiguration: this.ConnectionPolicy.RetryOptions?.GetRetryWithConfiguration(),
                     enableTcpConnectionEndpointRediscovery: this.ConnectionPolicy.EnableTcpConnectionEndpointRediscovery,
-                    addressResolver: this.AddressResolver,
                     rntbdMaxConcurrentOpeningConnectionCount: this.rntbdMaxConcurrentOpeningConnectionCount,
                     remoteCertificateValidationCallback: this.remoteCertificateValidationCallback,
                     distributedTracingOptions: distributedTracingOptions,
@@ -6719,6 +6707,18 @@ private void InitializeDirectConnectivity(IStoreClientFactory storeClientFactory
                 this.isStoreClientFactoryCreatedInternally = true;
             }
 
+            this.AddressResolver = new GlobalAddressResolver(
+                this.GlobalEndpointManager,
+                this.PartitionKeyRangeLocation,
+                this.ConnectionPolicy.ConnectionProtocol,
+                this,
+                this.collectionCache,
+                this.partitionKeyRangeCache,
+                this.accountServiceConfiguration,
+                this.ConnectionPolicy,
+                this.httpClient,
+                this.storeClientFactory.GetConnectionStateListener());
+
             this.CreateStoreModel(subscribeRntbdStatus: true);
         }
 

diff --git a/Microsoft.Azure.Cosmos/src/Routing/GatewayAddressCache.cs b/Microsoft.Azure.Cosmos/src/Routing/GatewayAddressCache.cs
@@ -53,6 +53,7 @@ internal class GatewayAddressCache : IAddressCache, IDisposable
         private readonly SemaphoreSlim semaphore;
         private readonly CosmosHttpClient httpClient;
         private readonly bool isReplicaAddressValidationEnabled;
+        private readonly IConnectionStateListener connectionStateListener;
 
         private Tuple<PartitionKeyRangeIdentity, PartitionAddressInformation> masterPartitionAddressCache;
         private DateTime suboptimalMasterPartitionTimestamp;
@@ -67,6 +68,7 @@ public GatewayAddressCache(
             IServiceConfigurationReader serviceConfigReader,
             CosmosHttpClient httpClient,
             IOpenConnectionsHandler openConnectionsHandler,
+            IConnectionStateListener connectionStateListener,
             long suboptimalPartitionForceRefreshIntervalInSeconds = 600,
             bool enableTcpConnectionEndpointRediscovery = false,
             bool replicaAddressValidationEnabled = false)
@@ -81,6 +83,7 @@ public GatewayAddressCache(
             this.serverPartitionAddressToPkRangeIdMap = new ConcurrentDictionary<ServerKey, HashSet<PartitionKeyRangeIdentity>>();
             this.suboptimalMasterPartitionTimestamp = DateTime.MaxValue;
             this.enableTcpConnectionEndpointRediscovery = enableTcpConnectionEndpointRediscovery;
+            this.connectionStateListener = connectionStateListener;
 
             this.suboptimalPartitionForceRefreshIntervalInSeconds = suboptimalPartitionForceRefreshIntervalInSeconds;
 
@@ -496,6 +499,13 @@ private static void LogPartitionCacheRefresh(
         public async Task MarkAddressesToUnhealthyAsync(
             ServerKey serverKey)
         {
+            if (this.disposedValue)
+            {
+                // Will enable Listener to un-register in-case of un-graceful dispose
+                // <see cref="ConnectionStateMuxListener.NotifyAsync(ServerKey, ConcurrentDictionary{Func{ServerKey, Task}, object})"/>
+                throw new ObjectDisposedException(nameof(GatewayAddressCache));
+            }
+
             if (serverKey == null)
             {
                 throw new ArgumentNullException(nameof(serverKey));
@@ -538,6 +548,9 @@ where serverKey.Equals(transportAddress.ReplicaServerKey)
 
                         address.SetUnhealthy();
                     }
+
+                    // Update the health status
+                    this.CaptureTransportAddressUriHealthStates(addressInfo, transportAddresses);
                 }
             }
         }
@@ -828,9 +841,21 @@ internal Tuple<PartitionKeyRangeIdentity, PartitionAddressInformation> ToPartiti
                        partitionKeyRangeIdentity.PartitionKeyRangeId,
                        addressInfo.PhysicalUri);
 
+                    HashSet<PartitionKeyRangeIdentity> createdValue = null;
+                    ServerKey serverKey = new ServerKey(new Uri(addressInfo.PhysicalUri));
                     HashSet<PartitionKeyRangeIdentity> pkRangeIdSet = this.serverPartitionAddressToPkRangeIdMap.GetOrAdd(
-                        new ServerKey(new Uri(addressInfo.PhysicalUri)),
-                        (_) => new HashSet<PartitionKeyRangeIdentity>());
+                        serverKey,
+                        (_) =>
+                        {
+                            createdValue = new HashSet<PartitionKeyRangeIdentity>();
+                            return createdValue;
+                        });
+
+                    if (object.ReferenceEquals(pkRangeIdSet, createdValue))
+                    {
+                        this.connectionStateListener.Register(serverKey, this.MarkAddressesToUnhealthyAsync);
+                    }
+
                     lock (pkRangeIdSet)
                     {
                         pkRangeIdSet.Add(partitionKeyRangeIdentity);
@@ -893,7 +918,7 @@ private static void LogAddressResolutionEnd(DocumentServiceRequest request, stri
 
         private static Protocol ProtocolFromString(string protocol)
         {
-            return (protocol.ToLowerInvariant()) switch
+            return protocol.ToLowerInvariant() switch
             {
                 RuntimeConstants.Protocols.HTTPS => Protocol.Https,
                 RuntimeConstants.Protocols.RNTBD => Protocol.Tcp,
@@ -903,7 +928,7 @@ private static Protocol ProtocolFromString(string protocol)
 
         private static string ProtocolString(Protocol protocol)
         {
-            return ((int)protocol) switch
+            return (int)protocol switch
             {
                 (int)Protocol.Https => RuntimeConstants.Protocols.HTTPS,
                 (int)Protocol.Tcp => RuntimeConstants.Protocols.RNTBD,
@@ -1071,11 +1096,18 @@ protected virtual void Dispose(bool disposing)
         {
             if (this.disposedValue)
             {
+                DefaultTrace.TraceInformation("GatewayAddressCache is already disposed {0}", this.GetHashCode());
                 return;
             }
 
             if (disposing)
             {
+                // Unregister the server-key
+                foreach (ServerKey serverKey in this.serverPartitionAddressToPkRangeIdMap.Keys)
+                {
+                    this.connectionStateListener.UnRegister(serverKey, this.MarkAddressesToUnhealthyAsync);
+                }
+
                 this.serverPartitionAddressCache?.Dispose();
             }
 

diff --git a/Microsoft.Azure.Cosmos/src/Routing/GlobalAddressResolver.cs b/Microsoft.Azure.Cosmos/src/Routing/GlobalAddressResolver.cs
@@ -39,6 +39,7 @@ internal sealed class GlobalAddressResolver : IAddressResolverExtension, IDispos
         private readonly ConcurrentDictionary<Uri, EndpointCache> addressCacheByEndpoint;
         private readonly bool enableTcpConnectionEndpointRediscovery;
         private readonly bool isReplicaAddressValidationEnabled;
+        private readonly IConnectionStateListener connectionStateListener;
         private IOpenConnectionsHandler openConnectionsHandler;
 
         public GlobalAddressResolver(
@@ -50,7 +51,8 @@ public GlobalAddressResolver(
             PartitionKeyRangeCache routingMapProvider,
             IServiceConfigurationReader serviceConfigReader,
             ConnectionPolicy connectionPolicy,
-            CosmosHttpClient httpClient)
+            CosmosHttpClient httpClient,
+            IConnectionStateListener connectionStateListener)
         {
             this.endpointManager = endpointManager;
             this.partitionKeyRangeLocationCache = partitionKeyRangeLocationCache;
@@ -60,6 +62,7 @@ public GlobalAddressResolver(
             this.routingMapProvider = routingMapProvider;
             this.serviceConfigReader = serviceConfigReader;
             this.httpClient = httpClient;
+            this.connectionStateListener = connectionStateListener;
 
             int maxBackupReadEndpoints =
                 !connectionPolicy.EnableReadRequestsFallback.HasValue || connectionPolicy.EnableReadRequestsFallback.Value
@@ -229,19 +232,6 @@ public async Task<PartitionAddressInformation> ResolveAsync(
             return await resolver.ResolveAsync(request, forceRefresh, cancellationToken);
         }
 
-        public async Task UpdateAsync(
-           ServerKey serverKey,
-           CancellationToken cancellationToken)
-        {
-            foreach (KeyValuePair<Uri, EndpointCache> addressCache in this.addressCacheByEndpoint)
-            {
-                // since we don't know which address cache contains the pkRanges mapped to this node,
-                // we mark all transport uris that has the same server key to unhealthy status in the
-                // AddressCaches of all regions.
-                await addressCache.Value.AddressCache.MarkAddressesToUnhealthyAsync(serverKey);
-            }
-        }
-
         /// <summary>
         /// ReplicatedResourceClient will use this API to get the direct connectivity AddressCache for given request.
         /// </summary>
@@ -283,6 +273,7 @@ private EndpointCache GetOrAddEndpoint(Uri endpoint)
                         this.serviceConfigReader,
                         this.httpClient,
                         this.openConnectionsHandler,
+                        this.connectionStateListener,
                         enableTcpConnectionEndpointRediscovery: this.enableTcpConnectionEndpointRediscovery,
                         replicaAddressValidationEnabled: this.isReplicaAddressValidationEnabled);