Fix retry policy for GoneException when writes fail with 410/0 from Service (Azure#17172)

Author: FabianMeiswinkel

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

@@ -18,6 +18,8 @@
  */
 public class GoneException extends CosmosException {
 
+    private boolean basedOn410ResponseFromService = false;
+
     /**
      * Instantiates a new Gone exception.
      *
@@ -159,4 +161,12 @@ public GoneException(String message,
                          String requestUriString) {
         super(message, innerException, headers, HttpConstants.StatusCodes.GONE, requestUriString);
     }
+
+    public boolean isBasedOn410ResponseFromService() {
+        return this.basedOn410ResponseFromService;
+    }
+
+    public void setIsBasedOn410ResponseFromService() {
+        this.basedOn410ResponseFromService = true;
+    }
 }

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

@@ -163,10 +163,12 @@ public Mono<ShouldRetryResult> shouldRetry(Exception exception) {
                 return Mono.just(ShouldRetryResult.noRetry());
             } else if (exception instanceof GoneException &&
                 !request.isReadOnly() &&
-                BridgeInternal.hasSendingRequestStarted((CosmosException)exception)) {
+                BridgeInternal.hasSendingRequestStarted((CosmosException)exception) &&
+                !((GoneException)exception).isBasedOn410ResponseFromService()) {
 
                 logger.warn(
-                    "Operation will NOT be retried. Write operations can not be retried safely when sending the request " +
+                    "Operation will NOT be retried. Write operations which failed due to transient transport errors " +
+                        "can not be retried safely when sending the request " +
                         "to the service because they aren't idempotent. Current attempt {}, Exception: ",
                     this.attemptCount,
                     exception);

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

@@ -863,15 +863,19 @@ private Mono<StoreResponse> createErrorResponseFromHttpResponse(String resourceA
                                 break;
                             } else {
                                 // Have the request URL in the exception message for debugging purposes.
-                                exception = new GoneException(
+                                GoneException goneExceptionFromService = new GoneException(
                                         String.format(
                                                 RMResources.ExceptionMessage,
                                                 RMResources.Gone),
                                         response.headers(),
                                         request.uri());
+                                goneExceptionFromService.setIsBasedOn410ResponseFromService();
 
-                                exception.getResponseHeaders().put(HttpConstants.HttpHeaders.ACTIVITY_ID,
-                                        activityId);
+                                goneExceptionFromService.getResponseHeaders().put(
+                                    HttpConstants.HttpHeaders.ACTIVITY_ID,
+                                    activityId);
+
+                                exception = goneExceptionFromService;
                                 break;
                             }
                         }

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

@@ -791,7 +791,10 @@ private void messageReceived(final ChannelHandlerContext context, final RntbdRes
                             cause = new PartitionKeyRangeGoneException(error, lsn, partitionKeyRangeId, responseHeaders);
                             break;
                         default:
-                            cause = new GoneException(error, lsn, partitionKeyRangeId, responseHeaders);
+                            GoneException goneExceptionFromService =
+                                new GoneException(error, lsn, partitionKeyRangeId, responseHeaders);
+                            goneExceptionFromService.setIsBasedOn410ResponseFromService();
+                            cause = goneExceptionFromService;
                             break;
                     }
                     break;

sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/directconnectivity/GoneAndRetryWithRetryPolicyTest.java

@@ -151,6 +151,58 @@ public void shouldNotRetryFlushedWriteWithGoneExceptionButForceAddressRefresh()
         assertThat(shouldRetryResult.backOffTime).isNull();
     }
 
+    /**
+     * GoneException for write which is already sent to the wire but based on receiving
+     * an actual response from the Service with 410 Status Code and SubStatusCode 0
+     * should result in retry
+     * shouldRetryResult. ShouldRetryResult
+     */
+    @Test(groups = { "unit" }, timeOut = TIMEOUT)
+    public void shouldRetryFlushedWriteWithGoneExceptionFromService() {
+        RxDocumentServiceRequest request = RxDocumentServiceRequest.create(
+            mockDiagnosticsClientContext(),
+            OperationType.Create,
+            ResourceType.Document);
+        GoneAndRetryWithRetryPolicy goneAndRetryWithRetryPolicy =
+            new GoneAndRetryWithRetryPolicy(request, 30);
+
+        Supplier<GoneException> goneExceptionForFlushedRequestSupplier = () -> {
+            GoneException goneExceptionForFlushedRequest = new GoneException();
+            BridgeInternal.setSendingRequestStarted(goneExceptionForFlushedRequest, true);
+            goneExceptionForFlushedRequest.setIsBasedOn410ResponseFromService();
+            return goneExceptionForFlushedRequest;
+        };
+
+        Mono<IRetryPolicy.ShouldRetryResult> singleShouldRetry = goneAndRetryWithRetryPolicy
+            .shouldRetry(goneExceptionForFlushedRequestSupplier.get());
+        IRetryPolicy.ShouldRetryResult shouldRetryResult = singleShouldRetry.block();
+        assertThat(shouldRetryResult.shouldRetry).isTrue();
+        assertThat(shouldRetryResult.policyArg.getValue0()).isTrue();
+        assertThat(shouldRetryResult.policyArg.getValue3()).isEqualTo(1);
+        assertThat(shouldRetryResult.backOffTime.getSeconds()).isEqualTo(0);
+
+        singleShouldRetry = goneAndRetryWithRetryPolicy.shouldRetry(goneExceptionForFlushedRequestSupplier.get());
+        shouldRetryResult = singleShouldRetry.block();
+        assertThat(shouldRetryResult.shouldRetry).isTrue();
+        assertThat(shouldRetryResult.policyArg.getValue0()).isTrue();
+        assertThat(shouldRetryResult.policyArg.getValue3()).isEqualTo(2);
+        assertThat(shouldRetryResult.backOffTime.getSeconds()).isEqualTo(1);
+
+        singleShouldRetry = goneAndRetryWithRetryPolicy.shouldRetry(goneExceptionForFlushedRequestSupplier.get());
+        shouldRetryResult = singleShouldRetry.block();
+        assertThat(shouldRetryResult.shouldRetry).isTrue();
+        assertThat(shouldRetryResult.policyArg.getValue0()).isTrue();
+        assertThat(shouldRetryResult.policyArg.getValue3()).isEqualTo(3);
+        assertThat(shouldRetryResult.backOffTime.getSeconds()).isEqualTo(2);
+
+        singleShouldRetry = goneAndRetryWithRetryPolicy.shouldRetry(goneExceptionForFlushedRequestSupplier.get());
+        shouldRetryResult = singleShouldRetry.block();
+        assertThat(shouldRetryResult.shouldRetry).isTrue();
+        assertThat(shouldRetryResult.policyArg.getValue0()).isTrue();
+        assertThat(shouldRetryResult.policyArg.getValue3()).isEqualTo(4);
+        assertThat(shouldRetryResult.backOffTime.getSeconds()).isEqualTo(4);
+    }
+
     /**
      * RequestTimeoutExceptions should not be retried for read or write - no address cache refresh expected
      * shouldRetryResult. ShouldRetryResult