ClientRetryPolicy: Adds Cross Regional Retry Logic on 429/3092

Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs

@@ -92,7 +92,11 @@ public virtual async Task<TransactionalBatchOperationResult> AddAsync(
             ItemBatchOperationContext context = new ItemBatchOperationContext(
                 resolvedPartitionKeyRangeId,
                 trace,
-                BatchAsyncContainerExecutor.GetRetryPolicy(this.cosmosContainer, operation.OperationType, this.retryOptions));
+                BatchAsyncContainerExecutor.GetRetryPolicy(
+                    this.cosmosContainer,
+                    this.cosmosClientContext?.DocumentClient?.GlobalEndpointManager,
+                    operation.OperationType,
+                    this.retryOptions));
 
             if (itemRequestOptions != null && itemRequestOptions.AddRequestHeaders != null)
             {
@@ -159,6 +163,7 @@ internal virtual async Task ValidateOperationAsync(
 
         private static IDocumentClientRetryPolicy GetRetryPolicy(
             ContainerInternal containerInternal,
+            GlobalEndpointManager endpointManager,
             OperationType operationType,
             RetryOptions retryOptions)
         {
@@ -167,6 +172,7 @@ private static IDocumentClientRetryPolicy GetRetryPolicy(
                operationType,
                new ResourceThrottleRetryPolicy(
                 retryOptions.MaxRetryAttemptsOnThrottledRequests,
+                endpointManager,
                 retryOptions.MaxRetryWaitTimeInSeconds));
         }
 

Microsoft.Azure.Cosmos/src/ClientRetryPolicy.cs

@@ -48,6 +48,7 @@ public ClientRetryPolicy(
         {
             this.throttlingRetry = new ResourceThrottleRetryPolicy(
                 retryOptions.MaxRetryAttemptsOnThrottledRequests,
+                endpointManager: globalEndpointManager,
                 retryOptions.MaxRetryWaitTimeInSeconds);
 
             this.globalEndpointManager = globalEndpointManager;
@@ -120,7 +121,18 @@ public async Task<ShouldRetryResult> ShouldRetryAsync(
                 }
             }
 
-            return await this.throttlingRetry.ShouldRetryAsync(exception, cancellationToken);
+            ShouldRetryResult throttleRetryResult = await this.throttlingRetry.ShouldRetryAsync(exception, cancellationToken);
+
+            // Today, the only scenario where we would receive a ServiceUnavailableException from the Throttling Retry Policy
+            // is when we get 410 (Gone) with sub status code 3092 (System Resource Not Available). Note that this is applicable
+            // for write requests targeted to a multiple master account. In such case, the 410/3092 will get converted into 503.
+            if (throttleRetryResult.ExceptionToThrow is ServiceUnavailableException)
+            {
+                return this.TryMarkEndpointUnavailableForPkRangeAndRetryOnServiceUnavailable(
+                    shouldMarkEndpointUnavailableForPkRange: true);
+            }
+
+            return throttleRetryResult;
         }
 
         /// <summary> 
@@ -143,7 +155,18 @@ public async Task<ShouldRetryResult> ShouldRetryAsync(
                 return shouldRetryResult;
             }
 
-            return await this.throttlingRetry.ShouldRetryAsync(cosmosResponseMessage, cancellationToken);
+            ShouldRetryResult throttleRetryResult = await this.throttlingRetry.ShouldRetryAsync(cosmosResponseMessage, cancellationToken);
+
+            // Today, the only scenario where we would receive a ServiceUnavailableException from the Throttling Retry Policy
+            // is when we get 410 (Gone) with sub status code 3092 (System Resource Not Available). Note that this is applicable
+            // for write requests targeted to a multiple master account. In such case, the 410/3092 will get converted into 503.
+            if (throttleRetryResult.ExceptionToThrow is ServiceUnavailableException)
+            {
+                return this.TryMarkEndpointUnavailableForPkRangeAndRetryOnServiceUnavailable(
+                    shouldMarkEndpointUnavailableForPkRange: true);
+            }
+
+            return throttleRetryResult;
         }
 
         /// <summary>
@@ -177,6 +200,7 @@ public void OnBeforeSendRequest(DocumentServiceRequest request)
             // This enables marking the endpoint unavailability on endpoint failover/unreachability
             this.locationEndpoint = this.globalEndpointManager.ResolveServiceEndpoint(request);
             request.RequestContext.RouteToLocation(this.locationEndpoint);
+            this.throttlingRetry.OnBeforeSendRequest(request);
         }
 
         private async Task<ShouldRetryResult> ShouldRetryInternalAsync(
@@ -274,16 +298,8 @@ private async Task<ShouldRetryResult> ShouldRetryInternalAsync(
             // Received 503 due to client connect timeout or Gateway
             if (statusCode == HttpStatusCode.ServiceUnavailable)
             {
-                DefaultTrace.TraceWarning("ClientRetryPolicy: ServiceUnavailable. Refresh cache and retry. Failed Location: {0}; ResourceAddress: {1}",
-                    this.documentServiceRequest?.RequestContext?.LocationEndpointToRoute?.ToString() ?? string.Empty,
-                    this.documentServiceRequest?.ResourceAddress ?? string.Empty);
-
-                // Mark the partition as unavailable.
-                // Let the ClientRetry logic decide if the request should be retried
-                this.partitionKeyRangeLocationCache.TryMarkEndpointUnavailableForPartitionKeyRange(
-                     this.documentServiceRequest);
-
-                return this.ShouldRetryOnServiceUnavailable();
+                return this.TryMarkEndpointUnavailableForPkRangeAndRetryOnServiceUnavailable(
+                    shouldMarkEndpointUnavailableForPkRange: true);
             }
 
             return null;
@@ -406,6 +422,33 @@ private ShouldRetryResult ShouldRetryOnSessionNotAvailable(DocumentServiceReques
             }
         }
 
+        /// <summary>
+        /// Attempts to mark the endpoint associated with the current partition key range as unavailable and determines if
+        /// a retry should be performed due to a ServiceUnavailable (503) response. This method is invoked when a 503
+        /// Service Unavailable response is received, indicating that the service might be temporarily unavailable.
+        /// It optionally marks the partition key range as unavailable, which will influence future routing decisions.
+        /// </summary>
+        /// <param name="shouldMarkEndpointUnavailableForPkRange">A boolean flag indicating whether the endpoint for the
+        /// current partition key range should be marked as unavailable.</param>
+        /// <returns>An instance of <see cref="ShouldRetryResult"/> indicating whether the operation should be retried.</returns>
+        private ShouldRetryResult TryMarkEndpointUnavailableForPkRangeAndRetryOnServiceUnavailable(
+            bool shouldMarkEndpointUnavailableForPkRange)
+        {
+            DefaultTrace.TraceWarning("ClientRetryPolicy: ServiceUnavailable. Refresh cache and retry. Failed Location: {0}; ResourceAddress: {1}",
+                this.documentServiceRequest?.RequestContext?.LocationEndpointToRoute?.ToString() ?? string.Empty,
+                this.documentServiceRequest?.ResourceAddress ?? string.Empty);
+
+            if (shouldMarkEndpointUnavailableForPkRange)
+            {
+                // Mark the partition as unavailable.
+                // Let the ClientRetry logic decide if the request should be retried
+                this.partitionKeyRangeLocationCache.TryMarkEndpointUnavailableForPartitionKeyRange(
+                     this.documentServiceRequest);
+            }
+
+            return this.ShouldRetryOnServiceUnavailable();
+        }
+
         /// <summary>
         /// For a ServiceUnavailable (503.0) we could be having a timeout from Direct/TCP locally or a request to Gateway request with a similar response due to an endpoint not yet available.
         /// We try and retry the request only if there are other regions available. The retry logic is applicable for single master write accounts as well.

Microsoft.Azure.Cosmos/src/DocumentClient.cs

@@ -988,8 +988,9 @@ internal virtual void Initialize(Uri serviceEndpoint,
             this.initializeTaskFactory = (_) => TaskHelper.InlineIfPossible<bool>(
                     () => this.GetInitializationTaskAsync(storeClientFactory: storeClientFactory),
                     new ResourceThrottleRetryPolicy(
-                        this.ConnectionPolicy.RetryOptions.MaxRetryAttemptsOnThrottledRequests,
-                        this.ConnectionPolicy.RetryOptions.MaxRetryWaitTimeInSeconds));
+                        maxAttemptCount: this.ConnectionPolicy.RetryOptions.MaxRetryAttemptsOnThrottledRequests,
+                        endpointManager: this.GlobalEndpointManager,
+                        maxWaitTimeInSeconds: this.ConnectionPolicy.RetryOptions.MaxRetryWaitTimeInSeconds));
 
             // Create the task to start the initialize task
             // Task will be awaited on in the EnsureValidClientAsync

Microsoft.Azure.Cosmos/src/MetadataRequestThrottleRetryPolicy.cs

@@ -72,6 +72,7 @@ public MetadataRequestThrottleRetryPolicy(
 
             this.throttlingRetryPolicy = new ResourceThrottleRetryPolicy(
                 maxRetryAttemptsOnThrottledRequests,
+                this.globalEndpointManager,
                 maxRetryWaitTimeInSeconds);
 
             this.retryContext = new MetadataRetryContext

Microsoft.Azure.Cosmos/src/ResourceThrottleRetryPolicy.cs

@@ -9,6 +9,7 @@ namespace Microsoft.Azure.Cosmos
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.Core.Trace;
+    using Microsoft.Azure.Cosmos.Routing;
     using Microsoft.Azure.Documents;
 
     // Retry when we receive the throttling from server.
@@ -19,12 +20,15 @@ internal sealed class ResourceThrottleRetryPolicy : IDocumentClientRetryPolicy
         private readonly uint backoffDelayFactor;
         private readonly int maxAttemptCount;
         private readonly TimeSpan maxWaitTimeInMilliseconds;
+        private readonly IGlobalEndpointManager globalEndpointManager;
 
         private int currentAttemptCount;
         private TimeSpan cumulativeRetryDelay;
+        private bool? isMultiMasterWriteRegion;
 
         public ResourceThrottleRetryPolicy(
             int maxAttemptCount,
+            IGlobalEndpointManager endpointManager,
             int maxWaitTimeInSeconds = DefaultMaxWaitTimeInSeconds,
             uint backoffDelayFactor = 1)
         {
@@ -33,6 +37,7 @@ public ResourceThrottleRetryPolicy(
                 throw new ArgumentException("maxWaitTimeInSeconds", "maxWaitTimeInSeconds must be less than " + (int.MaxValue / 1000));
             }
 
+            this.globalEndpointManager = endpointManager;
             this.maxAttemptCount = maxAttemptCount;
             this.backoffDelayFactor = backoffDelayFactor;
             this.maxWaitTimeInMilliseconds = TimeSpan.FromSeconds(maxWaitTimeInSeconds);
@@ -59,7 +64,9 @@ public Task<ShouldRetryResult> ShouldRetryAsync(
                     return Task.FromResult(ShouldRetryResult.NoRetry());
                 }
 
-                return this.ShouldRetryInternalAsync(dce.RetryAfter);
+                return this.ShouldRetryInternalAsync(
+                    dce?.GetSubStatus(),
+                    dce?.RetryAfter);
             }
 
             DefaultTrace.TraceError(
@@ -88,11 +95,34 @@ public Task<ShouldRetryResult> ShouldRetryAsync(
                 return Task.FromResult(ShouldRetryResult.NoRetry());
             }
 
-            return this.ShouldRetryInternalAsync(cosmosResponseMessage?.Headers.RetryAfter);
+            return this.ShouldRetryInternalAsync(
+                cosmosResponseMessage?.Headers.SubStatusCode,
+                cosmosResponseMessage?.Headers.RetryAfter,
+                cosmosResponseMessage?.CosmosException);
         }
 
-        private Task<ShouldRetryResult> ShouldRetryInternalAsync(TimeSpan? retryAfter)
+        private Task<ShouldRetryResult> ShouldRetryInternalAsync(
+            SubStatusCodes? subStatusCode,
+            TimeSpan? retryAfter,
+            Exception exception = null)
         {
+            if (this.isMultiMasterWriteRegion.HasValue
+                && this.isMultiMasterWriteRegion.Value
+                && subStatusCode != null
+                && subStatusCode == SubStatusCodes.SystemResourceUnavailable)
+            {
+                DefaultTrace.TraceError(
+                    "Operation will NOT be retried. Converting SystemResourceUnavailable (429/3092) to ServiceUnavailable (503). Current attempt {0} sub status code: {1}.",
+                    this.currentAttemptCount, SubStatusCodes.SystemResourceUnavailable);
+
+                ServiceUnavailableException exceptionToThrow = ServiceUnavailableException.Create(
+                    SubStatusCodes.SystemResourceUnavailable,
+                    innerException: exception);
+
+                return Task.FromResult(
+                    ShouldRetryResult.NoRetry(exceptionToThrow));
+            }
+
             TimeSpan retryDelay = TimeSpan.Zero;
             if (this.currentAttemptCount < this.maxAttemptCount &&
                 this.CheckIfRetryNeeded(retryAfter, out retryDelay))
@@ -133,6 +163,8 @@ private object GetExceptionMessage(Exception exception)
         /// <param name="request">The request being sent to the service.</param>
         public void OnBeforeSendRequest(DocumentServiceRequest request)
         {
+            this.isMultiMasterWriteRegion = !request.IsReadOnlyRequest
+                && (this.globalEndpointManager?.CanSupportMultipleWriteLocations(request) ?? false);
         }
 
         /// <summary>

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

@@ -534,6 +534,18 @@ public virtual async Task RefreshLocationAsync(bool forceRefresh = false)
 
             await this.RefreshDatabaseAccountInternalAsync(forceRefresh: forceRefresh);
         }
+
+        /// <summary>
+        /// Determines whether the current configuration and state of the service allow for supporting multiple write locations.
+        /// This method returns True is the AvailableWriteLocations in LocationCache is more than 1. Otherwise, it returns False.
+        /// </summary>
+        /// <param name="request">The document service request for which the write location support is being evaluated.</param>
+        /// <returns>A boolean flag indicating if the available write locations are more than one.</returns>
+        public bool CanSupportMultipleWriteLocations(DocumentServiceRequest request)
+        {
+            return this.CanUseMultipleWriteLocations(request)
+                && this.locationCache.GetAvailableWriteLocations()?.Count > 1;
+        }
 
 #pragma warning disable VSTHRD100 // Avoid async void methods
         private async void StartLocationBackgroundRefreshLoop()

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

@@ -36,5 +36,7 @@ internal interface IGlobalEndpointManager : IDisposable
         ReadOnlyDictionary<string, Uri> GetAvailableWriteEndpointsByLocation();
 
         ReadOnlyDictionary<string, Uri> GetAvailableReadEndpointsByLocation();
+
+        bool CanSupportMultipleWriteLocations(DocumentServiceRequest request);
     }
 }

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

@@ -263,6 +263,11 @@ public ReadOnlyCollection<string> GetAvailableReadLocations()
         {
             return this.locationInfo.AvailableReadLocations;
         }
+
+        public ReadOnlyCollection<string> GetAvailableWriteLocations()
+        {
+            return this.locationInfo.AvailableWriteLocations;
+        }
 
         /// <summary>
         /// Resolves request to service endpoint. 
@@ -529,11 +534,11 @@ public bool ShouldRefreshEndpoints(out bool canRefreshInBackground)
 
         public bool CanUseMultipleWriteLocations(DocumentServiceRequest request)
         {
-            return this.CanUseMultipleWriteLocations() &&
-                (request.ResourceType == ResourceType.Document ||
+            return this.CanUseMultipleWriteLocations()
+                && (request.ResourceType == ResourceType.Document ||
                 (request.ResourceType == ResourceType.StoredProcedure && request.OperationType == Documents.OperationType.ExecuteJavaScript));
-        }
-
+        }
+
         private void ClearStaleEndpointUnavailabilityInfo()
         {
             if (this.locationUnavailablityInfoByEndpoint.Any())