MaxRetryCountTest+[BUG]404/1002CyclicRegionRetry+[BUG]449RetryTooAggressively +[BUG]MissingCrossRegionRetryForWritesWithServer410 #37050)

sdk/cosmos/azure-cosmos-test/src/main/java/com/azure/cosmos/test/implementation/faultinjection/FaultInjectionServerErrorResultInternal.java

@@ -90,6 +90,8 @@ public CosmosException getInjectedServerError(RxDocumentServiceRequest request)
                 responseHeaders.put(
                     HttpConstants.HttpHeaders.RETRY_AFTER_IN_MILLISECONDS,
                     String.valueOf(500));
+                responseHeaders.put(WFConstants.BackendHeaders.SUB_STATUS,
+                    Integer.toString(HttpConstants.SubStatusCodes.USER_REQUEST_RATE_TOO_LARGE));
                 cosmosException = new RequestRateTooLargeException(null, lsn, partitionKeyRangeId, responseHeaders);
 
                 break;
@@ -137,7 +139,10 @@ public CosmosException getInjectedServerError(RxDocumentServiceRequest request)
             case SERVICE_UNAVAILABLE:
                 responseHeaders.put(WFConstants.BackendHeaders.SUB_STATUS,
                     Integer.toString(HttpConstants.SubStatusCodes.SERVER_GENERATED_503));
-                cosmosException = new ServiceUnavailableException(null, lsn, null, responseHeaders, HttpConstants.SubStatusCodes.SERVER_GENERATED_503);
+                ServiceUnavailableException serviceUnavailableException =
+                    new ServiceUnavailableException(null, lsn, null, responseHeaders, HttpConstants.SubStatusCodes.SERVER_GENERATED_503);
+                serviceUnavailableException.setIsBasedOn503ResponseFromService();
+                cosmosException = serviceUnavailableException;
                 break;
 
             case STALED_ADDRESSES_SERVER_GONE:

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientRetryPolicy.java

@@ -162,7 +162,8 @@ public Mono<ShouldRetryResult> shouldRetry(Exception e) {
             return this.shouldRetryOnBackendServiceUnavailableAsync(
                 this.isReadRequest,
                 isWebExceptionRetriable,
-                this.request.getNonIdempotentWriteRetriesEnabled());
+                this.request.getNonIdempotentWriteRetriesEnabled(),
+                clientException);
         }
 
         return this.throttlingRetry.shouldRetry(e);
@@ -204,7 +205,7 @@ private ShouldRetryResult shouldRetryOnSessionNotAvailable(RxDocumentServiceRequ
                     this.isReadRequest ?
                         this.globalEndpointManager.getApplicableReadEndpoints(request) : this.globalEndpointManager.getApplicableWriteEndpoints(request);
 
-                if (this.sessionTokenRetryCount > endpoints.size()) {
+                if (this.sessionTokenRetryCount >= endpoints.size()) {
                     // When use multiple write locations is true and the request has been tried
                     // on all locations, then don't retry the request
                     return ShouldRetryResult.noRetry();
@@ -299,11 +300,27 @@ private Mono<Void> refreshLocation(boolean isReadRequest, boolean forceRefresh,
     private Mono<ShouldRetryResult> shouldRetryOnBackendServiceUnavailableAsync(
         boolean isReadRequest,
         boolean isWebExceptionRetriable,
-        boolean nonIdempotentWriteRetriesEnabled) {
-
-        if (!isReadRequest && !nonIdempotentWriteRetriesEnabled && !isWebExceptionRetriable) {
+        boolean nonIdempotentWriteRetriesEnabled,
+        CosmosException cosmosException) {
+
+        // The request has failed with 503, SDK need to decide whether it is safe to retry for write operations
+        // For server generated retries, it is safe to retry
+        // For SDK generated 503, it will be more tricky as we have to decide the cause of it. For any causes that SDK not sure whether the request
+        // has reached/processed from server side, unless customer has specifically opted in for nonIdempotentWriteRetries, SDK should not retry.
+        // When SDK would generate 503:
+        //    - When server return 410, SDK may internally retry multiple times, when all the retries exhausted, SDK will bubble up 503 with corresponding subStatusCode
+        //      (Note: currently, subStatus code for read may get lost during the conversion, but for writes, the subStatus code will be reserved)
+        //    - when SDK generated 410 due to different reason (like connectionTimeout, transient timeout etc), SDK will internally retry multiple times
+        //      when all the retries exhausted, SDK will bubble up 503
+        //
+        // Fow now, without nonIdempotentWriteRetries being enabled, SDK will only retry for the following situation:
+        // 1. For any connection related errors, it will be covered under isWebExceptionRetriable -> which SDK will retry
+        // 2. For any server returned 503s, SDK will retry
+        // 3. For SDK generated 503, SDK will only retry if the subStatusCode is SERVER_GENERATED_410
+        if (!isReadRequest && !shouldRetryWriteOnServiceUnavailable(nonIdempotentWriteRetriesEnabled, isWebExceptionRetriable, cosmosException)) {
             logger.warn(
-                "shouldRetryOnBackendServiceUnavailableAsync() Not retrying on write with non retriable exception. Retry count = {}",
+                "shouldRetryOnBackendServiceUnavailableAsync() Not retrying" +
+                    " on write with non retriable exception and non server returned service unavailable. Retry count = {}",
                 this.serviceUnavailableRetryCount);
             return Mono.just(ShouldRetryResult.noRetry());
         }
@@ -374,6 +391,24 @@ CosmosDiagnostics getCosmosDiagnostics() {
         return cosmosDiagnostics;
     }
 
+    private boolean shouldRetryWriteOnServiceUnavailable(
+        boolean nonIdempotentWriteRetriesEnabled,
+        boolean isWebExceptionRetriable,
+        CosmosException cosmosException) {
+
+        if (nonIdempotentWriteRetriesEnabled || isWebExceptionRetriable) {
+            return true;
+        }
+
+        if (cosmosException instanceof ServiceUnavailableException) {
+            ServiceUnavailableException serviceUnavailableException = (ServiceUnavailableException) cosmosException;
+            return serviceUnavailableException.isBasedOn503ResponseFromService()
+                || serviceUnavailableException.getStatusCode() == HttpConstants.SubStatusCodes.SERVER_GENERATED_410;
+        }
+
+        return false;
+    }
+
     private static class RetryContext {
 
         public int retryCount;

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java

@@ -100,15 +100,15 @@ public class Configs {
     private static final int DEFAULT_ADDRESS_REFRESH_RESPONSE_TIMEOUT_IN_SECONDS = 5;
 
     // SessionTokenMismatchRetryPolicy Constants
-    private static final String DEFAULT_SESSION_TOKEN_MISMATCH_WAIT_TIME_IN_MILLISECONDS_NAME =
+    public static final String DEFAULT_SESSION_TOKEN_MISMATCH_WAIT_TIME_IN_MILLISECONDS_NAME =
         "COSMOS.DEFAULT_SESSION_TOKEN_MISMATCH_WAIT_TIME_IN_MILLISECONDS";
     private static final int DEFAULT_SESSION_TOKEN_MISMATCH_WAIT_TIME_IN_MILLISECONDS = 5000;
 
-    private static final String DEFAULT_SESSION_TOKEN_MISMATCH_INITIAL_BACKOFF_TIME_IN_MILLISECONDS_NAME =
+    public static final String DEFAULT_SESSION_TOKEN_MISMATCH_INITIAL_BACKOFF_TIME_IN_MILLISECONDS_NAME =
         "COSMOS.DEFAULT_SESSION_TOKEN_MISMATCH_INITIAL_BACKOFF_TIME_IN_MILLISECONDS";
     private static final int DEFAULT_SESSION_TOKEN_MISMATCH_INITIAL_BACKOFF_TIME_IN_MILLISECONDS = 5;
 
-    private static final String DEFAULT_SESSION_TOKEN_MISMATCH_MAXIMUM_BACKOFF_TIME_IN_MILLISECONDS_NAME =
+    public static final String DEFAULT_SESSION_TOKEN_MISMATCH_MAXIMUM_BACKOFF_TIME_IN_MILLISECONDS_NAME =
         "COSMOS.DEFAULT_SESSION_TOKEN_MISMATCH_MAXIMUM_BACKOFF_TIME_IN_MILLISECONDS";
     private static final int DEFAULT_SESSION_TOKEN_MISMATCH_MAXIMUM_BACKOFF_TIME_IN_MILLISECONDS = 500;
 
@@ -141,7 +141,7 @@ public class Configs {
     private static final String OPEN_CONNECTIONS_CONCURRENCY = "COSMOS.OPEN_CONNECTIONS_CONCURRENCY";
     private static final int DEFAULT_OPEN_CONNECTIONS_CONCURRENCY = 1;
 
-    private static final String MAX_RETRIES_IN_LOCAL_REGION_WHEN_REMOTE_REGION_PREFERRED = "COSMOS.MAX_RETRIES_IN_LOCAL_REGION_WHEN_REMOTE_REGION_PREFERRED";
+    public static final String MAX_RETRIES_IN_LOCAL_REGION_WHEN_REMOTE_REGION_PREFERRED = "COSMOS.MAX_RETRIES_IN_LOCAL_REGION_WHEN_REMOTE_REGION_PREFERRED";
     private static final int DEFAULT_MAX_RETRIES_IN_LOCAL_REGION_WHEN_REMOTE_REGION_PREFERRED = 1;
 
     public Configs() {

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ServiceUnavailableException.java

@@ -15,6 +15,8 @@
  * The type Service unavailable exception.
  */
 public class ServiceUnavailableException extends CosmosException {
+    private boolean basedOn503ResponseFromService = false;
+
     ServiceUnavailableException() {
         this(RMResources.ServiceUnavailable, HttpConstants.SubStatusCodes.UNKNOWN);
     }
@@ -89,6 +91,14 @@ public ServiceUnavailableException(String message,
         setSubStatus(subStatusCode);
     }
 
+    public void setIsBasedOn503ResponseFromService() {
+        this.basedOn503ResponseFromService = true;
+    }
+
+    public boolean isBasedOn503ResponseFromService() {
+        return this.basedOn503ResponseFromService;
+    }
+
     private void setSubStatus(int subStatusCode) {
         this.getResponseHeaders().put(
             HttpConstants.HttpHeaders.SUB_STATUS,

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/SessionTokenMismatchRetryPolicy.java

@@ -144,6 +144,6 @@ private boolean shouldRetryLocally(SessionRetryOptions sessionRetryOptions, int
         // hence to curb the retry attempts on a region,
         // compare sessionTokenMismatchRetryAttempts with max retry attempts allowed on the region - 1
         return !(regionSwitchHint == CosmosRegionSwitchHint.REMOTE_REGION_PREFERRED
-            && sessionTokenMismatchRetryAttempts == (this.maxRetryAttemptsInCurrentRegion.get() - 1));
+            && sessionTokenMismatchRetryAttempts >= (this.maxRetryAttemptsInCurrentRegion.get() - 1));
     }
 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/GoneAndRetryWithRetryPolicy.java

@@ -360,7 +360,15 @@ public Mono<ShouldRetryResult> shouldRetry(Exception exception) {
                 Math.min(
                     Math.min(this.currentBackoffMilliseconds.get() + random.nextInt(RANDOM_SALT_IN_MS), remainingMilliseconds),
                     RetryWithRetryPolicy.MAXIMUM_BACKOFF_TIME_IN_MS));
-            this.currentBackoffMilliseconds.accumulateAndGet(RetryWithRetryPolicy.BACK_OFF_MULTIPLIER, (left, right) -> left * right);
+
+            this.currentBackoffMilliseconds.set(
+                Math.max(
+                    RetryWithRetryPolicy.INITIAL_BACKOFF_TIME_MS,
+                    Math.min(
+                        RetryWithRetryPolicy.MAXIMUM_BACKOFF_TIME_IN_MS,
+                        this.currentBackoffMilliseconds.get() * RetryWithRetryPolicy.BACK_OFF_MULTIPLIER))
+            );
+
             logger.debug("BackoffTime: {} ms.", backoffTime.toMillis());
 
             // Calculate the remaining time based after accounting for the backoff that we

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/HttpTransportClient.java

@@ -953,9 +953,11 @@ private Mono<StoreResponse> createErrorResponseFromHttpResponse(String resourceA
 
                         case HttpConstants.StatusCodes.SERVICE_UNAVAILABLE:
                             int subStatusCode = getSubStatusCodeFromHeader(response);
-                            exception = new ServiceUnavailableException(errorMessage, response.headers(), request.uri(),
+                            ServiceUnavailableException serviceUnavailableException = new ServiceUnavailableException(errorMessage, response.headers(), request.uri(),
                                 (subStatusCode == 0) ? HttpConstants.SubStatusCodes.SERVER_GENERATED_503
                                     : HttpConstants.SubStatusCodes.UNKNOWN);
+                            serviceUnavailableException.setIsBasedOn503ResponseFromService();
+                            exception = serviceUnavailableException;
                             break;
 
                         case HttpConstants.StatusCodes.REQUEST_TIMEOUT:

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestManager.java

@@ -1113,8 +1113,15 @@ private void messageReceived(final ChannelHandlerContext context, final RntbdRes
                     break;
 
                 case StatusCodes.SERVICE_UNAVAILABLE:
-                    cause = new ServiceUnavailableException(error, lsn, partitionKeyRangeId, responseHeaders,
-                        SubStatusCodes.SERVER_GENERATED_503);
+                    ServiceUnavailableException serviceUnavailableException =
+                        new ServiceUnavailableException(
+                            error,
+                            lsn,
+                            partitionKeyRangeId,
+                            responseHeaders,
+                            SubStatusCodes.SERVER_GENERATED_503);
+                    serviceUnavailableException.setIsBasedOn503ResponseFromService();
+                    cause = serviceUnavailableException;
                     break;
 
                 case StatusCodes.TOO_MANY_REQUESTS: