Performance: Fix CPU consumption with EnableTcpConnectionEndpointRediscovery

Microsoft.Azure.Cosmos/Directory.Build.props

@@ -3,7 +3,7 @@
   <PropertyGroup>
     <ClientOfficialVersion>3.10.1</ClientOfficialVersion>
     <ClientPreviewVersion>3.9.1</ClientPreviewVersion>
-    <DirectVersion>3.11.1</DirectVersion>
+    <DirectVersion>3.11.2</DirectVersion>
     <EncryptionVersion>1.0.0-preview4</EncryptionVersion>
     <HybridRowVersion>1.0.0-preview</HybridRowVersion>
     <AboveDirBuildProps>$([MSBuild]::GetPathOfFileAbove('Directory.Build.props', '$(MSBuildThisFileDirectory)../'))</AboveDirBuildProps>

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

@@ -4,6 +4,7 @@
 namespace Microsoft.Azure.Cosmos
 {
     using System;
+    using System.Collections.Generic;
     using System.Diagnostics;
     using System.Globalization;
     using System.Linq;
@@ -13,6 +14,7 @@ namespace Microsoft.Azure.Cosmos
     using Microsoft.Azure.Cosmos.Core.Trace;
     using Microsoft.Azure.Cosmos.Routing;
     using Microsoft.Azure.Documents;
+    using Microsoft.Azure.Documents.Rntbd;
     using Microsoft.Azure.Documents.Routing;
     using Newtonsoft.Json;
 
@@ -709,6 +711,16 @@ private PartitionKeyRange TryResolveServerPartitionByPartitionKey(
             return null;
         }
 
+        public Task UpdateAsync(ServerKey serverKey, CancellationToken cancellationToken = default)
+        {
+            throw new NotImplementedException();
+        }
+
+        public Task UpdateAsync(IReadOnlyList<AddressCacheToken> addressCacheTokens, CancellationToken cancellationToken = default)
+        {
+            throw new NotImplementedException();
+        }
+
         private class ResolutionResult
         {
             public PartitionKeyRange TargetPartitionKeyRange { get; private set; }

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

@@ -19,6 +19,7 @@ namespace Microsoft.Azure.Cosmos.Routing
     using Microsoft.Azure.Documents;
     using Microsoft.Azure.Documents.Client;
     using Microsoft.Azure.Documents.Collections;
+    using Microsoft.Azure.Documents.Rntbd;
     using Microsoft.Azure.Documents.Routing;
 
     internal class GatewayAddressCache : IAddressCache, IDisposable
@@ -33,12 +34,14 @@ internal class GatewayAddressCache : IAddressCache, IDisposable
 
         private readonly AsyncCache<PartitionKeyRangeIdentity, PartitionAddressInformation> serverPartitionAddressCache;
         private readonly ConcurrentDictionary<PartitionKeyRangeIdentity, DateTime> suboptimalServerPartitionTimestamps;
+        private readonly ConcurrentDictionary<ServerKey, HashSet<PartitionKeyRangeIdentity>> serverPartitionAddressToPkRangeIdMap;
         private readonly IServiceConfigurationReader serviceConfigReader;
         private readonly long suboptimalPartitionForceRefreshIntervalInSeconds;
 
         private readonly Protocol protocol;
         private readonly string protocolFilter;
         private readonly IAuthorizationTokenProvider tokenProvider;
+        private readonly bool enableTcpConnectionEndpointRediscovery;
 
         private HttpClient httpClient;
 
@@ -54,7 +57,8 @@ public GatewayAddressCache(
             TimeSpan requestTimeout,
             long suboptimalPartitionForceRefreshIntervalInSeconds = 600,
             HttpMessageHandler messageHandler = null,
-            ApiType apiType = ApiType.None)
+            ApiType apiType = ApiType.None,
+            bool enableTcpConnectionEndpointRediscovery = false)
         {
             this.addressEndpoint = new Uri(serviceEndpoint + "/" + Paths.AddressPathSegment);
             this.protocol = protocol;
@@ -63,7 +67,9 @@ public GatewayAddressCache(
             this.serviceConfigReader = serviceConfigReader;
             this.serverPartitionAddressCache = new AsyncCache<PartitionKeyRangeIdentity, PartitionAddressInformation>();
             this.suboptimalServerPartitionTimestamps = new ConcurrentDictionary<PartitionKeyRangeIdentity, DateTime>();
+            this.serverPartitionAddressToPkRangeIdMap = new ConcurrentDictionary<ServerKey, HashSet<PartitionKeyRangeIdentity>>();
             this.suboptimalMasterPartitionTimestamp = DateTime.MaxValue;
+            this.enableTcpConnectionEndpointRediscovery = enableTcpConnectionEndpointRediscovery;
 
             this.suboptimalPartitionForceRefreshIntervalInSeconds = suboptimalPartitionForceRefreshIntervalInSeconds;
 
@@ -255,6 +261,46 @@ public async Task<PartitionAddressInformation> TryGetAddressesAsync(
             }
         }
 
+        public async Task TryUpdateAddressAsync(
+            ServerKey serverKey,
+            CancellationToken cancellationToken)
+        {
+            if (serverKey == null)
+            {
+                throw new ArgumentNullException(nameof(serverKey));
+            }
+
+            List<Task> tasks = new List<Task>();
+            HashSet<PartitionKeyRangeIdentity> pkRangeIds;
+            if (this.serverPartitionAddressToPkRangeIdMap.TryGetValue(serverKey, out pkRangeIds))
+            {
+                foreach (PartitionKeyRangeIdentity pkRangeId in pkRangeIds)
+                {
+                    DefaultTrace.TraceInformation("Refresh addresses for collectionRid :{0}, pkRangeId: {1}, serviceEndpoint: {2}",
+                       pkRangeId.CollectionRid,
+                       pkRangeId.PartitionKeyRangeId,
+                       this.serviceEndpoint);
+
+                    tasks.Add(this.serverPartitionAddressCache.GetAsync(
+                       pkRangeId,
+                       null,
+                       () => this.GetAddressesForRangeIdAsync(
+                           null,
+                           pkRangeId.CollectionRid,
+                           pkRangeId.PartitionKeyRangeId,
+                           forceRefresh: true),
+                       cancellationToken,
+                       forceRefresh: true));
+                }
+
+                // remove the server key from the map since we are updating the addresses
+                HashSet<PartitionKeyRangeIdentity> ignorePkRanges;
+                this.serverPartitionAddressToPkRangeIdMap.TryRemove(serverKey, out ignorePkRanges);
+            }
+
+            await Task.WhenAll(tasks);
+        }
+
         public async Task<PartitionAddressInformation> UpdateAsync(
             PartitionKeyRangeIdentity partitionKeyRangeIdentity,
             CancellationToken cancellationToken)
@@ -528,12 +574,29 @@ internal Tuple<PartitionKeyRangeIdentity, PartitionAddressInformation> ToPartiti
 
             PartitionKeyRangeIdentity partitionKeyRangeIdentity = new PartitionKeyRangeIdentity(collectionRid, address.PartitionKeyRangeId);
 
+            if (this.enableTcpConnectionEndpointRediscovery && partitionKeyRangeIdentity.PartitionKeyRangeId != PartitionKeyRange.MasterPartitionKeyRangeId)
+            {
+                // add serverKey-pkRangeIdentity mapping only for addresses retrieved from gateway
+                foreach (AddressInformation addressInfo in addressInfos)
+                {
+                    DefaultTrace.TraceInformation("Added address to serverPartitionAddressToPkRangeIdMap, collectionRid :{0}, pkRangeId: {1}, address: {2}",
+                       partitionKeyRangeIdentity.CollectionRid,
+                       partitionKeyRangeIdentity.PartitionKeyRangeId,
+                       addressInfo.PhysicalUri);
+
+                    this.serverPartitionAddressToPkRangeIdMap.AddOrUpdate(
+                        new ServerKey(new Uri(addressInfo.PhysicalUri)), new HashSet<PartitionKeyRangeIdentity>() { partitionKeyRangeIdentity },
+                        (serverKey, pkRangeIds) =>
+                        {
+                            pkRangeIds.Add(partitionKeyRangeIdentity);
+                            return pkRangeIds;
+                        });
+                }
+            }
+
             return Tuple.Create(
                 partitionKeyRangeIdentity,
-                new PartitionAddressInformation(
-                    addressInfos,
-                    partitionKeyRangeIdentity.PartitionKeyRangeId == PartitionKeyRange.MasterPartitionKeyRangeId ? null : partitionKeyRangeIdentity,
-                    this.serviceEndpoint));
+                new PartitionAddressInformation(addressInfos));
         }
 
         private static string LogAddressResolutionStart(DocumentServiceRequest request, Uri targetEndpoint)

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

@@ -14,12 +14,13 @@ namespace Microsoft.Azure.Cosmos.Routing
     using Microsoft.Azure.Cosmos.Common;
     using Microsoft.Azure.Documents;
     using Microsoft.Azure.Documents.Client;
+    using Microsoft.Azure.Documents.Rntbd;
 
     /// <summary>
     /// AddressCache implementation for client SDK. Supports cross region address routing based on
     /// avaialbility and preference list.
     /// </summary>
-    internal sealed class GlobalAddressResolver : IAddressResolverExtension, IDisposable
+    internal sealed class GlobalAddressResolver : IAddressResolver, IDisposable
     {
         private const int MaxBackupReadRegions = 3;
 
@@ -35,6 +36,7 @@ internal sealed class GlobalAddressResolver : IAddressResolverExtension, IDispos
         private readonly ConcurrentDictionary<Uri, EndpointCache> addressCacheByEndpoint;
         private readonly TimeSpan requestTimeout;
         private readonly ApiType apiType;
+        private readonly bool enableTcpConnectionEndpointRediscovery;
 
         public GlobalAddressResolver(
             GlobalEndpointManager endpointManager,
@@ -63,6 +65,8 @@ public GlobalAddressResolver(
                 !connectionPolicy.EnableReadRequestsFallback.HasValue || connectionPolicy.EnableReadRequestsFallback.Value
                 ? GlobalAddressResolver.MaxBackupReadRegions : 0;
 
+            this.enableTcpConnectionEndpointRediscovery = connectionPolicy.EnableTcpConnectionEndpointRediscovery;
+
             this.maxEndpoints = maxBackupReadEndpoints + 2; // for write and alternate write endpoint (during failover)
 
             this.addressCacheByEndpoint = new ConcurrentDictionary<Uri, EndpointCache>();
@@ -131,6 +135,21 @@ public async Task UpdateAsync(
             await Task.WhenAll(tasks);
         }
 
+        public async Task UpdateAsync(
+            ServerKey serverKey,
+            CancellationToken cancellationToken)
+        {
+            List<Task> tasks = new List<Task>();
+
+            foreach (KeyValuePair<Uri, EndpointCache> addressCache in this.addressCacheByEndpoint)
+            {
+                // since we don't know which address cache contains the pkRanges mapped to this node, we do a tryUpdate on all AddressCaches of all regions
+                tasks.Add(addressCache.Value.AddressCache.TryUpdateAddressAsync(serverKey, cancellationToken));
+            }
+
+            await Task.WhenAll(tasks);
+        }
+
         /// <summary>
         /// ReplicatedResourceClient will use this API to get the direct connectivity AddressCache for given request.
         /// </summary>
@@ -164,7 +183,8 @@ private EndpointCache GetOrAddEndpoint(Uri endpoint)
                         this.serviceConfigReader,
                         this.requestTimeout,
                         messageHandler: this.messageHandler,
-                        apiType: this.apiType);
+                        apiType: this.apiType,
+                        enableTcpConnectionEndpointRediscovery: this.enableTcpConnectionEndpointRediscovery);
 
                     string location = this.endpointManager.GetLocation(endpoint);
                     AddressResolver addressResolver = new AddressResolver(null, new NullRequestSigner(), location);

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/RequestEventHandlerTests.cs

@@ -81,9 +81,8 @@ private StoreClient GetMockStoreClient()
             AddressSelector addressSelector = new AddressSelector(mockAddressCache.Object, Protocol.Tcp);
             TransportClient mockTransportClient = this.GetMockTransportClient();
             ISessionContainer sessionContainer = new SessionContainer(string.Empty);
-            var connectionStateListener = new ConnectionStateListener(null);
 
-            StoreReader storeReader = new StoreReader(mockTransportClient, addressSelector, sessionContainer, connectionStateListener);
+            StoreReader storeReader = new StoreReader(mockTransportClient, addressSelector, sessionContainer);
 
             Mock<IAuthorizationTokenProvider> mockAuthorizationTokenProvider = new Mock<IAuthorizationTokenProvider>();
             mockAuthorizationTokenProvider.Setup(provider => provider.AddSystemAuthorizationHeaderAsync(
@@ -149,7 +148,7 @@ private Mock<IAddressResolver> GetMockAddressCache()
                     It.IsAny<DocumentServiceRequest>(),
                     It.IsAny<bool>(),
                     new CancellationToken()))
-                    .ReturnsAsync(new PartitionAddressInformation(addressInformation, null, null));
+                    .ReturnsAsync(new PartitionAddressInformation(addressInformation));
 
             return mockAddressCache;
         }