Initial code changes to throw 503 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;

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.AadTokenExpired)
+            {
+                DefaultTrace.TraceError(
+                    "Operation will NOT be retried. Converting 429/3092 to 503. Current attempt {0} sub status code: {1}.",
+                    this.currentAttemptCount, SubStatusCodes.AadTokenExpired);
+
+                ServiceUnavailableException exceptionToThrow = ServiceUnavailableException.Create(
+                    SubStatusCodes.AadTokenExpired,
+                    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?.CanUseMultipleWriteLocations(request) ?? false);
         }
 
         /// <summary>

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

@@ -529,11 +529,12 @@ public bool ShouldRefreshEndpoints(out bool canRefreshInBackground)
 
         public bool CanUseMultipleWriteLocations(DocumentServiceRequest request)
         {
-            return this.CanUseMultipleWriteLocations() &&
-                (request.ResourceType == ResourceType.Document ||
+            return this.CanUseMultipleWriteLocations()
+                && this.locationInfo.AvailableWriteLocations.Count > 1
+                && (request.ResourceType == ResourceType.Document ||
                 (request.ResourceType == ResourceType.StoredProcedure && request.OperationType == Documents.OperationType.ExecuteJavaScript));
-        }
-
+        }
+
         private void ClearStaleEndpointUnavailabilityInfo()
         {
             if (this.locationUnavailablityInfoByEndpoint.Any())

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

@@ -23,6 +23,23 @@ public class BatchAsyncBatcherTests
     {
         private static readonly Exception expectedException = new Exception();
         private static readonly BatchPartitionMetric metric = new BatchPartitionMetric();
+        private GlobalEndpointManager mockedEndpointManager;
+
+        [TestInitialize]
+        public void Initialize()
+        {
+            Mock<IDocumentClientInternal> mockedClient = new();
+
+            this.mockedEndpointManager = new(
+                mockedClient.Object,
+                new ConnectionPolicy());
+        }
+
+        [TestCleanup]
+        public void Cleanup()
+        {
+            this.mockedEndpointManager.Dispose();
+        }
 
         private ItemBatchOperation CreateItemBatchOperation(bool withContext = false)
         {
@@ -565,12 +582,12 @@ public async Task RetrierGetsCalledOnSplit()
             IDocumentClientRetryPolicy retryPolicy1 = new BulkExecutionRetryPolicy(
                 GetSplitEnabledContainer(),
                 OperationType.Read,
-                new ResourceThrottleRetryPolicy(1));
+                new ResourceThrottleRetryPolicy(1, this.mockedEndpointManager));
 
             IDocumentClientRetryPolicy retryPolicy2 = new BulkExecutionRetryPolicy(
                 GetSplitEnabledContainer(),
                 OperationType.Read,
-                new ResourceThrottleRetryPolicy(1));
+                new ResourceThrottleRetryPolicy(1, this.mockedEndpointManager));
 
             ItemBatchOperation operation1 = this.CreateItemBatchOperation();
             ItemBatchOperation operation2 = this.CreateItemBatchOperation();
@@ -594,12 +611,12 @@ public async Task RetrierGetsCalledOnCompletingSplit()
             IDocumentClientRetryPolicy retryPolicy1 = new BulkExecutionRetryPolicy(
                 GetSplitEnabledContainer(),
                 OperationType.Read,
-                new ResourceThrottleRetryPolicy(1));
+                new ResourceThrottleRetryPolicy(1, this.mockedEndpointManager));
 
             IDocumentClientRetryPolicy retryPolicy2 = new BulkExecutionRetryPolicy(
                 GetSplitEnabledContainer(),
                 OperationType.Read,
-                new ResourceThrottleRetryPolicy(1));
+                new ResourceThrottleRetryPolicy(1, this.mockedEndpointManager));
 
             ItemBatchOperation operation1 = this.CreateItemBatchOperation();
             ItemBatchOperation operation2 = this.CreateItemBatchOperation();
@@ -623,12 +640,12 @@ public async Task RetrierGetsCalledOnCompletingPartitionMigration()
             IDocumentClientRetryPolicy retryPolicy1 = new BulkExecutionRetryPolicy(
                 GetSplitEnabledContainer(),
                 OperationType.Read,
-                new ResourceThrottleRetryPolicy(1));
+                new ResourceThrottleRetryPolicy(1, this.mockedEndpointManager));
 
             IDocumentClientRetryPolicy retryPolicy2 = new BulkExecutionRetryPolicy(
                 GetSplitEnabledContainer(),
                 OperationType.Read,
-                new ResourceThrottleRetryPolicy(1));
+                new ResourceThrottleRetryPolicy(1, this.mockedEndpointManager));
 
             ItemBatchOperation operation1 = this.CreateItemBatchOperation();
             ItemBatchOperation operation2 = this.CreateItemBatchOperation();
@@ -672,17 +689,17 @@ public async Task RetrierGetsCalledOn413_3402()
             IDocumentClientRetryPolicy retryPolicy1 = new BulkExecutionRetryPolicy(
                 GetSplitEnabledContainer(),
                 OperationType.Read,
-                new ResourceThrottleRetryPolicy(1));
+                new ResourceThrottleRetryPolicy(1, this.mockedEndpointManager));
 
             IDocumentClientRetryPolicy retryPolicy2 = new BulkExecutionRetryPolicy(
                 GetSplitEnabledContainer(),
                 OperationType.Read,
-                new ResourceThrottleRetryPolicy(1));
+                new ResourceThrottleRetryPolicy(1, this.mockedEndpointManager));
 
             IDocumentClientRetryPolicy retryPolicy3 = new BulkExecutionRetryPolicy(
                 GetSplitEnabledContainer(),
                 OperationType.Create,
-                new ResourceThrottleRetryPolicy(1));
+                new ResourceThrottleRetryPolicy(1, this.mockedEndpointManager));
 
             ItemBatchOperation operation1 = this.CreateItemBatchOperation();
             ItemBatchOperation operation2 = this.CreateItemBatchOperation();
@@ -710,17 +727,17 @@ public async Task RetrierGetsCalledOn413_NoSubstatus()
             IDocumentClientRetryPolicy retryPolicy1 = new BulkExecutionRetryPolicy(
                 GetSplitEnabledContainer(),
                 OperationType.Read,
-                new ResourceThrottleRetryPolicy(1));
+                new ResourceThrottleRetryPolicy(1, this.mockedEndpointManager));
 
             IDocumentClientRetryPolicy retryPolicy2 = new BulkExecutionRetryPolicy(
                 GetSplitEnabledContainer(),
                 OperationType.Read,
-                new ResourceThrottleRetryPolicy(1));
+                new ResourceThrottleRetryPolicy(1, this.mockedEndpointManager));
 
             IDocumentClientRetryPolicy retryPolicy3 = new BulkExecutionRetryPolicy(
                 GetSplitEnabledContainer(),
                 OperationType.Create,
-                new ResourceThrottleRetryPolicy(1));
+                new ResourceThrottleRetryPolicy(1, this.mockedEndpointManager));
 
             ItemBatchOperation operation1 = this.CreateItemBatchOperation();
             ItemBatchOperation operation2 = this.CreateItemBatchOperation();