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

Microsoft.Azure.Cosmos/src/ClientRetryPolicy.cs

@@ -98,6 +98,17 @@ public async Task<ShouldRetryResult> ShouldRetryAsync(
                 }
             }
 
+            if (exception is CosmosException cosmosException)
+            {
+                ShouldRetryResult shouldRetryResult = await this.ShouldRetryInternalAsync(
+                    cosmosException.StatusCode,
+                    cosmosException.Headers.SubStatusCode);
+                if (shouldRetryResult != null)
+                {
+                    return shouldRetryResult;
+                }
+            }
+
             return await this.throttlingRetry.ShouldRetryAsync(exception, cancellationToken);
         }
 

Microsoft.Azure.Cosmos/src/DocumentClient.cs

@@ -662,7 +662,7 @@ private async Task OpenPrivateAsync(CancellationToken cancellationToken)
                     tokenProvider: this, 
                     retryPolicy: this.retryPolicy,
                     telemetryToServiceHelper: this.telemetryToServiceHelper);
-                this.partitionKeyRangeCache = new PartitionKeyRangeCache(this, this.GatewayStoreModel, this.collectionCache);
+                this.partitionKeyRangeCache = new PartitionKeyRangeCache(this, this.GatewayStoreModel, this.collectionCache, this.GlobalEndpointManager);
 
                 DefaultTrace.TraceWarning("{0} occurred while OpenAsync. Exception Message: {1}", ex.ToString(), ex.Message);
             }
@@ -1033,7 +1033,7 @@ private async Task<bool> GetInitializationTaskAsync(IStoreClientFactory storeCli
                     tokenProvider: this, 
                     retryPolicy: this.retryPolicy,
                     telemetryToServiceHelper: this.telemetryToServiceHelper);
-            this.partitionKeyRangeCache = new PartitionKeyRangeCache(this, this.GatewayStoreModel, this.collectionCache);
+            this.partitionKeyRangeCache = new PartitionKeyRangeCache(this, this.GatewayStoreModel, this.collectionCache, this.GlobalEndpointManager);
             this.ResetSessionTokenRetryPolicy = new ResetSessionTokenRetryPolicyFactory(this.sessionContainer, this.collectionCache, this.retryPolicy);
 
             gatewayStoreModel.SetCaches(this.partitionKeyRangeCache, this.collectionCache);

Microsoft.Azure.Cosmos/src/HttpClient/CosmosHttpClientCore.cs

@@ -395,7 +395,9 @@ private async Task<HttpResponseMessage> SendHttpHelperAsync(
                                             headers: new Headers()
                                             {
                                                 ActivityId = System.Diagnostics.Trace.CorrelationManager.ActivityId.ToString(),
-                                                SubStatusCode = SubStatusCodes.TransportGenerated503
+                                                SubStatusCode = resourceType == ResourceType.PartitionKeyRange
+                                                ? SubStatusCodes.PartitionKeyRangeGone
+                                                : SubStatusCodes.TransportGenerated503
                                             },
                                             trace: trace,
                                             innerException: e);

Microsoft.Azure.Cosmos/src/HttpClient/HttpTimeoutPolicy.cs

@@ -30,7 +30,7 @@ public static HttpTimeoutPolicy GetTimeoutPolicy(
             }
 
             //Partition Key Requests
-            if (documentServiceRequest.ResourceType == ResourceType.PartitionKeyRange)
+            if (documentServiceRequest.ResourceType == ResourceType.Address)
             {
                 return HttpTimeoutPolicyControlPlaneRetriableHotPath.Instance;
             }

Microsoft.Azure.Cosmos/src/HttpClient/HttpTimeoutPolicyDefault.cs

@@ -22,9 +22,9 @@ private HttpTimeoutPolicyDefault(bool shouldThrow503OnTimeout)
 
         private readonly IReadOnlyList<(TimeSpan requestTimeout, TimeSpan delayForNextRequest)> TimeoutsAndDelays = new List<(TimeSpan requestTimeout, TimeSpan delayForNextRequest)>()
         {
-            (TimeSpan.FromSeconds(65), TimeSpan.Zero),
-            (TimeSpan.FromSeconds(65), TimeSpan.FromSeconds(1)),
-            (TimeSpan.FromSeconds(65), TimeSpan.Zero),
+            (TimeSpan.FromSeconds(3), TimeSpan.Zero),
+            (TimeSpan.FromSeconds(5), TimeSpan.FromSeconds(1)),
+            (TimeSpan.FromSeconds(10), TimeSpan.Zero),
         };
 
         public override string TimeoutPolicyName => HttpTimeoutPolicyDefault.Name;

Microsoft.Azure.Cosmos/src/MetadataRequestThrottleRetryPolicy.cs

@@ -0,0 +1,128 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos
+{
+    using System;
+    using System.Net;
+    using System.Threading;
+    using System.Threading.Tasks;
+    using Microsoft.Azure.Cosmos.Core.Trace;
+    using Microsoft.Azure.Cosmos.Routing;
+    using Microsoft.Azure.Documents;
+
+    /// <summary>
+    /// Metadata Request Throttle Retry Policy is combination of endpoint change retry + throttling retry.
+    /// </summary>
+    internal sealed class MetadataRequestThrottleRetryPolicy : IDocumentClientRetryPolicy
+    {
+        /// <summary>
+        /// An instance of <see cref="GlobalEndpointManager"/>.
+        /// </summary>
+        private readonly GlobalEndpointManager globalEndpointManager;
+
+        /// <summary>
+        /// Defines the throttling retry policy that is used as the underlying retry policy.
+        /// </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="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(
+            GlobalEndpointManager endpointManager,
+            int maxRetryAttemptsOnThrottledRequests,
+            int maxRetryWaitTimeInSeconds)
+        {
+            this.globalEndpointManager = endpointManager;
+            this.throttlingRetryPolicy = new ResourceThrottleRetryPolicy(
+                maxRetryAttemptsOnThrottledRequests,
+                maxRetryWaitTimeInSeconds);
+        }
+
+        /// <summary> 
+        /// Should the caller retry the operation.
+        /// </summary>
+        /// <param name="exception">Exception that occured when the operation was tried</param>
+        /// <param name="cancellationToken">An instance of <see cref="CancellationToken"/>.</param>
+        /// <returns>True indicates caller should retry, False otherwise</returns>
+        public Task<ShouldRetryResult> ShouldRetryAsync(
+            Exception exception,
+            CancellationToken cancellationToken)
+        {
+            if (exception is CosmosException cosmosException
+                && cosmosException.StatusCode == HttpStatusCode.ServiceUnavailable
+                && cosmosException.Headers.SubStatusCode == SubStatusCodes.PartitionKeyRangeGone)
+            {
+                if (!this.MarkEndpointUnavailableForRead())
+                {
+                    return Task.FromResult(ShouldRetryResult.NoRetry());
+                }
+            }
+
+            return this.throttlingRetryPolicy.ShouldRetryAsync(exception, cancellationToken);
+        }
+
+        /// <summary> 
+        /// Should the caller retry the operation.
+        /// </summary>
+        /// <param name="cosmosResponseMessage"><see cref="ResponseMessage"/> in return of the request</param>
+        /// <param name="cancellationToken">An instance of <see cref="CancellationToken"/>.</param>
+        /// <returns>True indicates caller should retry, False otherwise</returns>
+        public Task<ShouldRetryResult> ShouldRetryAsync(
+            ResponseMessage cosmosResponseMessage,
+            CancellationToken cancellationToken)
+        {
+            if (cosmosResponseMessage?.StatusCode == HttpStatusCode.ServiceUnavailable
+                && cosmosResponseMessage?.Headers.SubStatusCode == SubStatusCodes.PartitionKeyRangeGone)
+            {
+                if (!this.MarkEndpointUnavailableForRead())
+                {
+                    return Task.FromResult(ShouldRetryResult.NoRetry());
+                }
+            }
+
+            return this.throttlingRetryPolicy.ShouldRetryAsync(cosmosResponseMessage, cancellationToken);
+        }
+
+        /// <summary>
+        /// Method that is called before a request is sent to allow the retry policy implementation
+        /// to modify the state of the request.
+        /// </summary>
+        /// <param name="request">The request being sent to the service.</param>
+        public void OnBeforeSendRequest(DocumentServiceRequest request)
+        {
+            this.metadataLocationEndpoint = this.globalEndpointManager.ResolveServiceEndpoint(request);
+        }
+
+        /// <summary>
+        /// Marks an endpoint unavailable in the global endpoint manager, for any future read requests.
+        /// </summary>
+        /// <returns>A boolean flag indicating if the operation was successful.</returns>
+        private bool MarkEndpointUnavailableForRead()
+        {
+            if (this.metadataLocationEndpoint != null)
+            {
+                DefaultTrace.TraceWarning("MetadataRequestThrottleRetryPolicy: Marking the following endpoint unavailable for reads: {0}.", this.metadataLocationEndpoint);
+                this.globalEndpointManager.MarkEndpointUnavailableForRead(this.metadataLocationEndpoint);
+                return true;
+            }
+            else
+            {
+                DefaultTrace.TraceWarning("MetadataRequestThrottleRetryPolicy: location endpoint couldn't be resolved. Skip marking endpoint unavailable for reads.");
+                return false;
+            }
+        }
+    }
+}

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

@@ -30,17 +30,20 @@ internal class PartitionKeyRangeCache : IRoutingMapProvider, ICollectionRoutingM
         private readonly ICosmosAuthorizationTokenProvider authorizationTokenProvider;
         private readonly IStoreModel storeModel;
         private readonly CollectionCache collectionCache;
+        private readonly GlobalEndpointManager endpointManager;
 
         public PartitionKeyRangeCache(
             ICosmosAuthorizationTokenProvider authorizationTokenProvider,
             IStoreModel storeModel,
-            CollectionCache collectionCache)
+            CollectionCache collectionCache,
+            GlobalEndpointManager endpointManager)
         {
             this.routingMapCache = new AsyncCacheNonBlocking<string, CollectionRoutingMap>(
                     keyEqualityComparer: StringComparer.Ordinal);
             this.authorizationTokenProvider = authorizationTokenProvider;
             this.storeModel = storeModel;
             this.collectionCache = collectionCache;
+            this.endpointManager = endpointManager;
         }
 
         public virtual async Task<IReadOnlyList<PartitionKeyRange>> TryGetOverlappingRangesAsync(
@@ -121,10 +124,10 @@ public virtual async Task<CollectionRoutingMap> TryLookupAsync(
                 return await this.routingMapCache.GetAsync(
                     key: collectionRid,
                     singleValueInitFunc: (_) => this.GetRoutingMapForCollectionAsync(
-                        collectionRid, 
-                        previousValue, 
-                        trace,
-                        request?.RequestContext?.ClientRequestStatistics),
+                        collectionRid: collectionRid,
+                        previousRoutingMap: previousValue,
+                        trace: trace,
+                        clientSideRequestStatistics: request?.RequestContext?.ClientRequestStatistics),
                     forceRefresh: (currentValue) => PartitionKeyRangeCache.ShouldForceRefresh(previousValue, currentValue));
             }
             catch (DocumentClientException ex)
@@ -174,35 +177,6 @@ private static bool ShouldForceRefresh(
             return previousValue.ChangeFeedNextIfNoneMatch == currentValue.ChangeFeedNextIfNoneMatch; 
         }
 
-        public async Task<PartitionKeyRange> TryGetRangeByPartitionKeyRangeIdAsync(string collectionRid, 
-                            string partitionKeyRangeId, 
-                            ITrace trace,
-                            IClientSideRequestStatistics clientSideRequestStatistics)
-        {
-            try
-            {
-                CollectionRoutingMap routingMap = await this.routingMapCache.GetAsync(
-                    key: collectionRid,
-                    singleValueInitFunc: (_) => this.GetRoutingMapForCollectionAsync(
-                        collectionRid: collectionRid,
-                        previousRoutingMap: null,
-                        trace: trace,
-                        clientSideRequestStatistics: clientSideRequestStatistics),
-                    forceRefresh: (_) => false);
-
-                return routingMap.TryGetRangeByPartitionKeyRangeId(partitionKeyRangeId);
-            }
-            catch (DocumentClientException ex)
-            {
-                if (ex.StatusCode == HttpStatusCode.NotFound)
-                {
-                    return null;
-                }
-
-                throw;
-            }
-        }
-
         private async Task<CollectionRoutingMap> GetRoutingMapForCollectionAsync(
             string collectionRid,
             CollectionRoutingMap previousRoutingMap,
@@ -225,9 +199,14 @@ private async Task<CollectionRoutingMap> GetRoutingMapForCollectionAsync(
                 }
 
                 RetryOptions retryOptions = new RetryOptions();
+                MetadataRequestThrottleRetryPolicy metadataRetryPolicy = new (
+                        endpointManager: this.endpointManager,
+                        maxRetryAttemptsOnThrottledRequests: retryOptions.MaxRetryAttemptsOnThrottledRequests,
+                        maxRetryWaitTimeInSeconds: retryOptions.MaxRetryWaitTimeInSeconds);
+
                 using (DocumentServiceResponse response = await BackoffRetryUtility<DocumentServiceResponse>.ExecuteAsync(
-                    () => this.ExecutePartitionKeyRangeReadChangeFeedAsync(collectionRid, headers, trace, clientSideRequestStatistics),
-                    new ResourceThrottleRetryPolicy(retryOptions.MaxRetryAttemptsOnThrottledRequests, retryOptions.MaxRetryWaitTimeInSeconds)))
+                    () => this.ExecutePartitionKeyRangeReadChangeFeedAsync(collectionRid, headers, trace, clientSideRequestStatistics, metadataRetryPolicy),
+                    retryPolicy: metadataRetryPolicy))
                 {
                     lastStatusCode = response.StatusCode;
                     changeFeedNextIfNoneMatch = response.Headers[HttpConstants.HttpHeaders.ETag];
@@ -274,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))
             {
@@ -285,6 +265,7 @@ private async Task<DocumentServiceResponse> ExecutePartitionKeyRangeReadChangeFe
                     AuthorizationTokenType.PrimaryMasterKey,
                     headers))
                 {
+                    retryPolicy?.OnBeforeSendRequest(request);
                     string authorizationToken = null;
                     try
                     {

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Mocks/MockDocumentClient.cs

@@ -228,7 +228,7 @@ private void Init()
                     },
                 string.Empty);
 
-            this.partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(null, null, null);
+            this.partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(null, null, null, null);
             this.partitionKeyRangeCache.Setup(
                         m => m.TryLookupAsync(
                             It.IsAny<string>(),

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs

@@ -798,7 +798,7 @@ private class ClientWithSplitDetection : MockDocumentClient
 
             public ClientWithSplitDetection()
             {
-                this.partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(MockBehavior.Strict, null, null, null);
+                this.partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(MockBehavior.Strict, null, null, null, null);
                 this.partitionKeyRangeCache.Setup(
                         m => m.TryGetOverlappingRangesAsync(
                             It.IsAny<string>(),

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncContainerExecutorTests.cs

@@ -376,7 +376,7 @@ private class ClientWithSplitDetection : MockDocumentClient
 
             public ClientWithSplitDetection()
             {
-                this.partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(MockBehavior.Strict, null, null, null);
+                this.partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(MockBehavior.Strict, null, null, null, null);
                 this.partitionKeyRangeCache.Setup(
                         m => m.TryGetOverlappingRangesAsync(
                             It.IsAny<string>(),

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncOperationContextTests.cs

@@ -287,7 +287,7 @@ private class ClientWithSplitDetection : MockDocumentClient
 
             public ClientWithSplitDetection()
             {
-                this.partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(MockBehavior.Strict, null, null, null);
+                this.partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(MockBehavior.Strict, null, null, null, null);
                 this.partitionKeyRangeCache.Setup(
                         m => m.TryGetOverlappingRangesAsync(
                             It.IsAny<string>(),

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

@@ -145,7 +145,7 @@ private class ClientWithSplitDetection : MockDocumentClient
 
             public ClientWithSplitDetection()
             {
-                this.partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(MockBehavior.Strict, null, null, null);
+                this.partitionKeyRangeCache = new Mock<PartitionKeyRangeCache>(MockBehavior.Strict, null, null, null, null);
                 this.partitionKeyRangeCache.Setup(
                         m => m.TryGetOverlappingRangesAsync(
                             It.IsAny<string>(),