[Internal] Per Partition Automatic Failover: Fixes Metadata Requests Retry Policy

Microsoft.Azure.Cosmos/src/MetadataRequestThrottleRetryPolicy.cs

@@ -17,11 +17,6 @@ namespace Microsoft.Azure.Cosmos
     /// </summary>
     internal sealed class MetadataRequestThrottleRetryPolicy : IDocumentClientRetryPolicy
     {
-        /// <summary>
-        /// A callback delegate to fetch the location endpoint at a later point of time.
-        /// </summary>
-        private readonly Func<Uri> locationEndpointCallbackUri;
-
         /// <summary>
         /// An instance of <see cref="GlobalEndpointManager"/>.
         /// </summary>
@@ -32,21 +27,24 @@ internal sealed class MetadataRequestThrottleRetryPolicy : IDocumentClientRetryP
         /// </summary>
         private readonly IDocumentClientRetryPolicy throttlingRetryPolicy;
 
+        /// <summary>
+        /// An instance of <see cref="Uri"/> containing the location endpoint where the partition key
+        /// range http request will be sent over.
+        /// </summary>
+        private Uri metadataLocationEndpoint;
+
         /// <summary>
         /// The constructor to initialize an instance of <see cref="MetadataRequestThrottleRetryPolicy"/>.
         /// </summary>
-        /// <param name="locationEndpointCallbackUri">A callback delegate to fetch the location endpoint at a later point of time.</param>
         /// <param name="endpointManager">An instance of <see cref="GlobalEndpointManager"/></param>
         /// <param name="maxRetryAttemptsOnThrottledRequests">An integer defining the maximum number
         /// of attempts to retry when requests are throttled.</param>
         /// <param name="maxRetryWaitTimeInSeconds">An integer defining the maximum wait time in seconds.</param>
         public MetadataRequestThrottleRetryPolicy(
-            Func<Uri> locationEndpointCallbackUri,
             GlobalEndpointManager endpointManager,
             int maxRetryAttemptsOnThrottledRequests,
             int maxRetryWaitTimeInSeconds)
         {
-            this.locationEndpointCallbackUri = locationEndpointCallbackUri;
             this.globalEndpointManager = endpointManager;
             this.throttlingRetryPolicy = new ResourceThrottleRetryPolicy(
                 maxRetryAttemptsOnThrottledRequests,
@@ -105,6 +103,7 @@ public Task<ShouldRetryResult> ShouldRetryAsync(
         /// <param name="request">The request being sent to the service.</param>
         public void OnBeforeSendRequest(DocumentServiceRequest request)
         {
+            this.metadataLocationEndpoint = this.globalEndpointManager.ResolveServiceEndpoint(request);
         }
 
         /// <summary>
@@ -113,11 +112,10 @@ public void OnBeforeSendRequest(DocumentServiceRequest request)
         /// <returns>A boolean flag indicating if the operation was successful.</returns>
         private bool MarkEndpointUnavailableForRead()
         {
-            Uri location = this.locationEndpointCallbackUri?.Invoke();
-            if (location != null)
+            if (this.metadataLocationEndpoint != null)
             {
-                DefaultTrace.TraceWarning("MetadataRequestThrottleRetryPolicy: Marking the following endpoint unavailable for reads: {0}.", location);
-                this.globalEndpointManager.MarkEndpointUnavailableForRead(location);
+                DefaultTrace.TraceWarning("MetadataRequestThrottleRetryPolicy: Marking the following endpoint unavailable for reads: {0}.", this.metadataLocationEndpoint);
+                this.globalEndpointManager.MarkEndpointUnavailableForRead(this.metadataLocationEndpoint);
                 return true;
             }
             else

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

@@ -32,8 +32,6 @@ internal class PartitionKeyRangeCache : IRoutingMapProvider, ICollectionRoutingM
         private readonly CollectionCache collectionCache;
         private readonly GlobalEndpointManager endpointManager;
 
-        private Uri pkRangeLocationEndpoint;
-
         public PartitionKeyRangeCache(
             ICosmosAuthorizationTokenProvider authorizationTokenProvider,
             IStoreModel storeModel,
@@ -201,13 +199,14 @@ private async Task<CollectionRoutingMap> GetRoutingMapForCollectionAsync(
                 }
 
                 RetryOptions retryOptions = new RetryOptions();
-                using (DocumentServiceResponse response = await BackoffRetryUtility<DocumentServiceResponse>.ExecuteAsync(
-                    () => this.ExecutePartitionKeyRangeReadChangeFeedAsync(collectionRid, headers, trace, clientSideRequestStatistics),
-                    new MetadataRequestThrottleRetryPolicy(
-                        locationEndpointCallbackUri: () => this.pkRangeLocationEndpoint,
+                MetadataRequestThrottleRetryPolicy metadataRetryPolicy = new (
                         endpointManager: this.endpointManager,
                         maxRetryAttemptsOnThrottledRequests: retryOptions.MaxRetryAttemptsOnThrottledRequests,
-                        maxRetryWaitTimeInSeconds: retryOptions.MaxRetryWaitTimeInSeconds)))
+                        maxRetryWaitTimeInSeconds: retryOptions.MaxRetryWaitTimeInSeconds);
+
+                using (DocumentServiceResponse response = await BackoffRetryUtility<DocumentServiceResponse>.ExecuteAsync(
+                    () => this.ExecutePartitionKeyRangeReadChangeFeedAsync(collectionRid, headers, trace, clientSideRequestStatistics, metadataRetryPolicy),
+                    retryPolicy: metadataRetryPolicy))
                 {
                     lastStatusCode = response.StatusCode;
                     changeFeedNextIfNoneMatch = response.Headers[HttpConstants.HttpHeaders.ETag];
@@ -254,7 +253,8 @@ private async Task<CollectionRoutingMap> GetRoutingMapForCollectionAsync(
         private async Task<DocumentServiceResponse> ExecutePartitionKeyRangeReadChangeFeedAsync(string collectionRid, 
                                                                                 INameValueCollection headers, 
                                                                                 ITrace trace,
-                                                                                IClientSideRequestStatistics clientSideRequestStatistics)
+                                                                                IClientSideRequestStatistics clientSideRequestStatistics,
+                                                                                IDocumentClientRetryPolicy retryPolicy)
         {
             using (ITrace childTrace = trace.StartChild("Read PartitionKeyRange Change Feed", TraceComponent.Transport, Tracing.TraceLevel.Info))
             {
@@ -265,7 +265,7 @@ private async Task<DocumentServiceResponse> ExecutePartitionKeyRangeReadChangeFe
                     AuthorizationTokenType.PrimaryMasterKey,
                     headers))
                 {
-                    this.pkRangeLocationEndpoint = this.endpointManager.ResolveServiceEndpoint(request);
+                    retryPolicy?.OnBeforeSendRequest(request);
                     string authorizationToken = null;
                     try
                     {