Fix direct retry handling for partition key range gone

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

@@ -12,6 +12,7 @@
 import com.azure.cosmos.implementation.LeaseNotFoundException;
 import com.azure.cosmos.implementation.OperationType;
 import com.azure.cosmos.implementation.PartitionIsMigratingException;
+import com.azure.cosmos.implementation.PartitionKeyRangeGoneException;
 import com.azure.cosmos.implementation.PartitionKeyRangeIsSplittingException;
 import com.azure.cosmos.implementation.RequestTimeoutException;
 import com.azure.cosmos.implementation.ResourceType;
@@ -325,6 +326,52 @@ public void shouldRetryWithPartitionKeyRangeIsSplittingException() {
 
     }
 
+    /**
+     * Retry with PartitionKeyRangeGoneException
+     */
+    @Test(groups = { "unit" }, timeOut = TIMEOUT)
+    public void shouldRetryWithPartitionKeyRangeGoneException() {
+        RxDocumentServiceRequest request = RxDocumentServiceRequest.create(
+            mockDiagnosticsClientContext(),
+            OperationType.Read,
+            ResourceType.Document);
+        GoneAndRetryWithRetryPolicy goneAndRetryWithRetryPolicy = new GoneAndRetryWithRetryPolicy(request, 30);
+        Mono<ShouldRetryResult> singleShouldRetry = goneAndRetryWithRetryPolicy
+            .shouldRetry(new PartitionKeyRangeGoneException());
+        ShouldRetryResult shouldRetryResult = singleShouldRetry.block();
+        assertThat(shouldRetryResult.shouldRetry).isTrue();
+        assertThat(request.forcePartitionKeyRangeRefresh).isTrue();
+        assertThat(request.requestContext.resolvedPartitionKeyRange).isNull();
+        assertThat(request.requestContext.quorumSelectedLSN).isEqualTo(-1);
+        assertThat(shouldRetryResult.policyArg.getValue0()).isFalse();
+    }
+
+    @Test(groups = { "unit" }, timeOut = TIMEOUT)
+    public void shouldWrapPartitionKeyRangeGoneExceptionWithServiceUnavailableWhenRetryBudgetExhausted() {
+        RxDocumentServiceRequest request = RxDocumentServiceRequest.create(
+            mockDiagnosticsClientContext(),
+            OperationType.Read,
+            ResourceType.Document);
+        GoneAndRetryWithRetryPolicy goneAndRetryWithRetryPolicy = new GoneAndRetryWithRetryPolicy(request, 0);
+
+        ShouldRetryResult shouldRetryResult = goneAndRetryWithRetryPolicy
+            .shouldRetry(new PartitionKeyRangeGoneException())
+            .block();
+        assertThat(shouldRetryResult.shouldRetry).isTrue();
+
+        shouldRetryResult = goneAndRetryWithRetryPolicy
+            .shouldRetry(new PartitionKeyRangeGoneException())
+            .block();
+
+        assertThat(shouldRetryResult.shouldRetry).isFalse();
+        assertThat(shouldRetryResult.exception).isInstanceOf(CosmosException.class);
+
+        CosmosException cosmosException = (CosmosException) shouldRetryResult.exception;
+        assertThat(cosmosException.getStatusCode()).isEqualTo(HttpConstants.StatusCodes.SERVICE_UNAVAILABLE);
+        assertThat(cosmosException.getSubStatusCode())
+            .isEqualTo(HttpConstants.SubStatusCodes.PARTITION_KEY_RANGE_GONE_EXCEEDED_RETRY_LIMIT);
+    }
+
     /**
      * No retry on bad request exception
      */

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -7,6 +7,7 @@
 #### Breaking Changes
 
 #### Bugs Fixed
+* Fixed direct connectivity retries for `PartitionKeyRangeGoneException` thrown during address resolution so stale partition key ranges trigger routing-map refresh and exhausted retries surface as 503 with partition-key-range-gone retry-limit substatus. - See [Issue 49381](https://github.com/Azure/azure-sdk-for-java/issues/49381).
 
 #### Other Changes
 * Reduced memory footprint of deserialized `PartitionKeyRange` instances by stripping unused fields in the `PartitionKeyRange(ObjectNode)` constructor - See PR [49513](https://github.com/Azure/azure-sdk-for-java/pull/49513).

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

@@ -127,6 +127,7 @@ private boolean isNonRetryableException(Exception exception) {
             if (exception instanceof GoneException ||
                 exception instanceof PartitionIsMigratingException ||
                 exception instanceof PartitionKeyRangeIsSplittingException ||
+                exception instanceof PartitionKeyRangeGoneException ||
                 exception instanceof LeaseNotFoundException) {
 
                 return false;
@@ -292,6 +293,8 @@ private Pair<Mono<ShouldRetryResult>, Boolean> handleException(Exception excepti
                 return handlePartitionIsMigratingException((PartitionIsMigratingException)exception);
             } else if (exception instanceof PartitionKeyRangeIsSplittingException) {
                 return handlePartitionKeyIsSplittingException((PartitionKeyRangeIsSplittingException) exception);
+            } else if (exception instanceof PartitionKeyRangeGoneException) {
+                return handlePartitionKeyRangeGoneException((PartitionKeyRangeGoneException) exception);
             }
 
             throw new IllegalStateException("Invalid exception type", exception);
@@ -316,6 +319,18 @@ private Pair<Mono<ShouldRetryResult>, Boolean> handlePartitionKeyIsSplittingExce
             this.request.forcePartitionKeyRangeRefresh = true;
             return Pair.of(null, false);
         }
+
+        private Pair<Mono<ShouldRetryResult>, Boolean> handlePartitionKeyRangeGoneException(PartitionKeyRangeGoneException exception) {
+            // PartitionKeyRangeGoneException is generally treated as non-retriable, but when it is thrown while resolving
+            // addresses in direct mode it typically indicates stale routing/partition state; clear the cached target and
+            // force a routing-map (partition key range) refresh to allow the request to be re-routed.
+            this.request.requestContext.resolvedPartitionKeyRange = null;
+            this.request.requestContext.quorumSelectedLSN = -1;
+            this.request.requestContext.quorumSelectedStoreResponse = null;
+            logger.debug("Received partition key range gone exception, will retry, {}", exception.toString());
+            this.request.forcePartitionKeyRangeRefresh = true;
+            return Pair.of(null, false);
+        }
     }
 
 }