ClientRetryPolicy: Adds Cross Regional Retry Logic on 429/3092 (#4691)

* Initial code changes to throw 503 on 429/3092.

* Updated client retry policy. Added more tests to cover 429/3092.

* Code changes to update direct package version. Updating the tests.

* Code changes to refactor client retry policy.

* Minor code cleanup.

* Reverting the direct version bump up change.

* Code changes to address some of the review comments.

* Code changes to move failover logic in client retry policy.

* Minor code clean up.

* Code changes to clean up some cosmetic items.

* Further clean up.

* Code changes to address review comments.

* Minor refactor to address cosmetic update.

* Code changes to address cosmetic review comment.

Microsoft.Azure.Cosmos/src/ClientRetryPolicy.cs

@@ -35,6 +35,7 @@ internal sealed class ClientRetryPolicy : IDocumentClientRetryPolicy
  private int serviceUnavailableRetryCount;
  private bool isReadRequest;
  private bool canUseMultipleWriteLocations;
+ private bool isMultiMasterWriteRequest;
  private Uri locationEndpoint;
  private RetryContext retryContext;
  private DocumentServiceRequest documentServiceRequest;
@@ -57,6 +58,7 @@ public ClientRetryPolicy(
  this.sessionTokenRetryCount = 0;
  this.serviceUnavailableRetryCount = 0;
  this.canUseMultipleWriteLocations = false;
+ this.isMultiMasterWriteRequest = false;
  this.isPertitionLevelFailoverEnabled = isPertitionLevelFailoverEnabled;
  }
 
@@ -97,6 +99,23 @@ public async Task<ShouldRetryResult> ShouldRetryAsync(
 
  if (exception is DocumentClientException clientException)
  {
+ // Today, the only scenario where we would treat a throttling (429) exception as service unavailable is when we
+ // get 429 (TooManyRequests) 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 429/3092 will be treated as 503. The
+ // reason to keep the code out of the throttling retry policy is that in the near future, the 3092 sub status code
+ // might not be a throttling scenario at all and the status code in that case would be different than 429.
+ if (this.ShouldMarkEndpointUnavailableOnSystemResourceUnavailableForWrite(
+ clientException.StatusCode,
+ clientException.GetSubStatus()))
+ {
+ DefaultTrace.TraceError(
+ "Operation will NOT be retried on local region. Treating SystemResourceUnavailable (429/3092) as ServiceUnavailable (503). Status code: {0}, sub status code: {1}.",
+ StatusCodes.TooManyRequests, SubStatusCodes.SystemResourceUnavailable);
+
+ return this.TryMarkEndpointUnavailableForPkRangeAndRetryOnServiceUnavailable(
+ shouldMarkEndpointUnavailableForPkRange: true);
+ }
+
  ShouldRetryResult shouldRetryResult = await this.ShouldRetryInternalAsync(
  clientException?.StatusCode,
  clientException?.GetSubStatus());
@@ -143,6 +162,23 @@ public async Task<ShouldRetryResult> ShouldRetryAsync(
  return shouldRetryResult;
  }
 
+ // Today, the only scenario where we would treat a throttling (429) exception as service unavailable is when we
+ // get 429 (TooManyRequests) 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 429/3092 will be treated as 503. The
+ // reason to keep the code out of the throttling retry policy is that in the near future, the 3092 sub status code
+ // might not be a throttling scenario at all and the status code in that case would be different than 429.
+ if (this.ShouldMarkEndpointUnavailableOnSystemResourceUnavailableForWrite(
+ cosmosResponseMessage.StatusCode,
+ cosmosResponseMessage?.Headers.SubStatusCode))
+ {
+ DefaultTrace.TraceError(
+ "Operation will NOT be retried on local region. Treating SystemResourceUnavailable (429/3092) as ServiceUnavailable (503). Status code: {0}, sub status code: {1}.",
+ StatusCodes.TooManyRequests, SubStatusCodes.SystemResourceUnavailable);
+
+ return this.TryMarkEndpointUnavailableForPkRangeAndRetryOnServiceUnavailable(
+ shouldMarkEndpointUnavailableForPkRange: true);
+ }
+
  return await this.throttlingRetry.ShouldRetryAsync(cosmosResponseMessage, cancellationToken);
  }
 
@@ -156,6 +192,8 @@ public void OnBeforeSendRequest(DocumentServiceRequest request)
  this.isReadRequest = request.IsReadOnlyRequest;
  this.canUseMultipleWriteLocations = this.globalEndpointManager.CanUseMultipleWriteLocations(request);
  this.documentServiceRequest = request;
+ this.isMultiMasterWriteRequest = !this.isReadRequest
+ && (this.globalEndpointManager?.CanSupportMultipleWriteLocations(request) ?? false);
 
  // clear previous location-based routing directive
  request.RequestContext.ClearRouteToLocation();
@@ -274,16 +312,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 +436,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.
@@ -449,6 +506,24 @@ private ShouldRetryResult ShouldRetryOnServiceUnavailable()
  return ShouldRetryResult.RetryAfter(TimeSpan.Zero);
  }
 
+ /// <summary>
+ /// Returns a boolean flag indicating if the endpoint should be marked as unavailable
+ /// due to a 429 response with a sub status code of 3092 (system resource unavailable).
+ /// This is applicable for write requests targeted for multi master accounts.
+ /// </summary>
+ /// <param name="statusCode">An instance of <see cref="HttpStatusCode"/> containing the status code.</param>
+ /// <param name="subStatusCode">An instance of <see cref="SubStatusCodes"/> containing the sub status code.</param>
+ /// <returns>A boolean flag indicating is the endpoint should be marked as unavailable.</returns>
+ private bool ShouldMarkEndpointUnavailableOnSystemResourceUnavailableForWrite(
+ HttpStatusCode? statusCode,
+ SubStatusCodes? subStatusCode)
+ {
+ return this.isMultiMasterWriteRequest
+ && statusCode.HasValue
+ && (int)statusCode.Value == (int)StatusCodes.TooManyRequests
+ && subStatusCode == SubStatusCodes.SystemResourceUnavailable;
+ }
+
  private sealed class RetryContext
  {
  public int RetryLocationIndex { get; set; }

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

@@ -534,6 +534,20 @@ 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.locationCache.CanUseMultipleWriteLocations()
+ && this.locationCache.GetAvailableWriteLocations()?.Count > 1
+ && (request.ResourceType == ResourceType.Document ||
+ (request.ResourceType == ResourceType.StoredProcedure && request.OperationType == OperationType.Execute));
+ }
 
 #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. 
@@ -532,8 +537,8 @@ public bool CanUseMultipleWriteLocations(DocumentServiceRequest request)
  return this.CanUseMultipleWriteLocations() &&
  (request.ResourceType == ResourceType.Document ||
  (request.ResourceType == ResourceType.StoredProcedure && request.OperationType == Documents.OperationType.ExecuteJavaScript));
- }
-
+ }
+
  private void ClearStaleEndpointUnavailabilityInfo()
  {
  if (this.locationUnavailablityInfoByEndpoint.Any())
@@ -768,7 +773,7 @@ private ReadOnlyDictionary<string, Uri> GetEndpointByLocation(IEnumerable<Accoun
  return new ReadOnlyDictionary<string, Uri>(endpointsByLocation);
  }
 
- private bool CanUseMultipleWriteLocations()
+ internal bool CanUseMultipleWriteLocations()
  {
  return this.useMultipleWriteLocations && this.enableMultipleWriteLocations;
  }

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

@@ -86,6 +86,82 @@ public void MultimasterMetadataWriteRetryTest()
  retryPolicy.OnBeforeSendRequest(request);
  Assert.AreEqual(request.RequestContext.LocationEndpointToRoute, ClientRetryPolicyTests.Location1Endpoint);
  }
+
+ /// <summary>
+ /// Test to validate that when 429.3092 is thrown from the service, write requests on
+ /// a multi master account should be converted to 503 and retried to the next region.
+ /// </summary>
+ [TestMethod]
+ [DataRow(true, DisplayName = "Validate retry policy with multi master write account.")]
+ [DataRow(false, DisplayName = "Validate retry policy with single master write account.")]
+ public async Task ShouldRetryAsync_WhenRequestThrottledWithResourceNotAvailable_ShouldThrow503OnMultiMasterWriteAndRetryOnNextRegion(
+ bool isMultiMasterAccount)
+ {
+ // Arrange.
+ const bool enableEndpointDiscovery = true;
+ using GlobalEndpointManager endpointManager = this.Initialize(
+ useMultipleWriteLocations: isMultiMasterAccount,
+ enableEndpointDiscovery: enableEndpointDiscovery,
+ isPreferredLocationsListEmpty: false,
+ multimasterMetadataWriteRetryTest: true);
+
+ await endpointManager.RefreshLocationAsync();
+
+ ClientRetryPolicy retryPolicy = new (
+ endpointManager,
+ this.partitionKeyRangeLocationCache,
+ new RetryOptions(),
+ enableEndpointDiscovery,
+ false);
+
+ // Creates a sample write request.
+ DocumentServiceRequest request = this.CreateRequest(
+ isReadRequest: false,
+ isMasterResourceType: false);
+
+ // On first attempt should get (default/non hub) location.
+ retryPolicy.OnBeforeSendRequest(request);
+ Assert.AreEqual(request.RequestContext.LocationEndpointToRoute, ClientRetryPolicyTests.Location1Endpoint);
+
+ // Creation of 429.3092 Error.
+ HttpStatusCode throttleException = HttpStatusCode.TooManyRequests;
+ SubStatusCodes resourceNotAvailable = SubStatusCodes.SystemResourceUnavailable;
+
+ Exception innerException = new ();
+ Mock<INameValueCollection> nameValueCollection = new ();
+ DocumentClientException documentClientException = new (
+ message: "SystemResourceUnavailable: 429 with 3092 occurred.",
+ innerException: innerException,
+ statusCode: throttleException,
+ substatusCode: resourceNotAvailable,
+ requestUri: request.RequestContext.LocationEndpointToRoute,
+ responseHeaders: nameValueCollection.Object);
+
+ // Act.
+ Task<ShouldRetryResult> shouldRetry = retryPolicy.ShouldRetryAsync(
+ documentClientException,
+ new CancellationToken());
+
+ // Assert.
+ Assert.IsTrue(shouldRetry.Result.ShouldRetry);
+ retryPolicy.OnBeforeSendRequest(request);
+
+ if (isMultiMasterAccount)
+ {
+ Assert.AreEqual(
+ expected: ClientRetryPolicyTests.Location2Endpoint,
+ actual: request.RequestContext.LocationEndpointToRoute,
+ message: "The request should be routed to the next region, since the accound is a multi master write account and the request" +
+ "failed with 429.309 which got converted into 503 internally. This should trigger another retry attempt to the next region.");
+ }
+ else
+ {
+ Assert.AreEqual(
+ expected: ClientRetryPolicyTests.Location1Endpoint,
+ actual: request.RequestContext.LocationEndpointToRoute,
+ message: "Since this is asingle master account, the write request should not be retried on the next region.");
+ }
+ }
 
  /// <summary>
  /// Tests to see if different 503 substatus codes are handeled correctly