Fix PKRangeGone retry scope for address resolution

Scopes `PartitionKeyRangeGoneException` retries to address-resolution stale routing-map cases, marks both address resolver paths, and updates retry-policy tests for marked vs unmarked behavior.

Author: arnabnandy7

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

@@ -337,7 +337,7 @@ public void shouldRetryWithAddressResolutionPartitionKeyRangeGoneException() {
             ResourceType.Document);
         GoneAndRetryWithRetryPolicy goneAndRetryWithRetryPolicy = new GoneAndRetryWithRetryPolicy(request, 30);
         Mono<ShouldRetryResult> singleShouldRetry = goneAndRetryWithRetryPolicy
-            .shouldRetry(new PartitionKeyRangeGoneException().setShouldRetryWithRoutingMapRefresh());
+            .shouldRetry(new PartitionKeyRangeGoneException().markRetryWithRoutingMapRefresh());
         ShouldRetryResult shouldRetryResult = singleShouldRetry.block();
         assertThat(shouldRetryResult.shouldRetry).isTrue();
         assertThat(request.forcePartitionKeyRangeRefresh).isTrue();
@@ -369,12 +369,12 @@ public void shouldWrapAddressResolutionPartitionKeyRangeGoneExceptionWithService
         GoneAndRetryWithRetryPolicy goneAndRetryWithRetryPolicy = new GoneAndRetryWithRetryPolicy(request, 0);
 
         ShouldRetryResult shouldRetryResult = goneAndRetryWithRetryPolicy
-            .shouldRetry(new PartitionKeyRangeGoneException().setShouldRetryWithRoutingMapRefresh())
+            .shouldRetry(new PartitionKeyRangeGoneException().markRetryWithRoutingMapRefresh())
             .block();
         assertThat(shouldRetryResult.shouldRetry).isTrue();
 
         shouldRetryResult = goneAndRetryWithRetryPolicy
-            .shouldRetry(new PartitionKeyRangeGoneException().setShouldRetryWithRoutingMapRefresh())
+            .shouldRetry(new PartitionKeyRangeGoneException().markRetryWithRoutingMapRefresh())
             .block();
 
         assertThat(shouldRetryResult.shouldRetry).isFalse();

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

@@ -92,7 +92,7 @@ public boolean shouldRetryWithRoutingMapRefresh() {
         return this.shouldRetryWithRoutingMapRefresh;
     }
 
-    public PartitionKeyRangeGoneException setShouldRetryWithRoutingMapRefresh() {
+    public PartitionKeyRangeGoneException markRetryWithRoutingMapRefresh() {
         this.shouldRetryWithRoutingMapRefresh = true;
         return this;
     }

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

@@ -584,7 +584,7 @@ private ResolutionResult handleRangeAddressResolutionFailure(
                 request.getPartitionKeyRangeIdentity().getPartitionKeyRangeId(),
                 request.getPartitionKeyRangeIdentity().getCollectionRid());
             throw BridgeInternal.setResourceAddress(
-                new PartitionKeyRangeGoneException(errorMessage).setShouldRetryWithRoutingMapRefresh(),
+                new PartitionKeyRangeGoneException(errorMessage).markRetryWithRoutingMapRefresh(),
                 request.requestContext.resourcePhysicalAddress);
         }
         logger.debug("handleRangeAddressResolutionFailure returns null");

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

@@ -733,7 +733,7 @@ private Mono<AddressInformation[]> getAddressesForRangeId(
                                 collectionRid);
 
                             PartitionKeyRangeGoneException e = new PartitionKeyRangeGoneException(errorMessage)
-                                .setShouldRetryWithRoutingMapRefresh();
+                                .markRetryWithRoutingMapRefresh();
                             BridgeInternal.setResourceAddress(e, collectionRid);
 
                             return Mono.error(e);

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

@@ -317,9 +317,7 @@ private Pair<Mono<ShouldRetryResult>, Boolean> handlePartitionIsMigratingExcepti
         }
 
         private Pair<Mono<ShouldRetryResult>, Boolean> handlePartitionKeyIsSplittingException(PartitionKeyRangeIsSplittingException exception) {
-            this.request.requestContext.resolvedPartitionKeyRange = null;
-            this.request.requestContext.quorumSelectedLSN = -1;
-            this.request.requestContext.quorumSelectedStoreResponse = null;
+            resetRequestContextForPartitionKeyRangeRefresh();
             logger.debug("Received partition key range splitting exception, will retry, {}", exception.toString());
             this.request.forcePartitionKeyRangeRefresh = true;
             return Pair.of(null, false);
@@ -329,13 +327,17 @@ private Pair<Mono<ShouldRetryResult>, Boolean> handlePartitionKeyRangeGoneExcept
             // 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;
+            resetRequestContextForPartitionKeyRangeRefresh();
             logger.debug("Received partition key range gone exception, will retry, {}", exception.toString());
             this.request.forcePartitionKeyRangeRefresh = true;
             return Pair.of(null, false);
         }
+
+        private void resetRequestContextForPartitionKeyRangeRefresh() {
+            this.request.requestContext.resolvedPartitionKeyRange = null;
+            this.request.requestContext.quorumSelectedLSN = -1;
+            this.request.requestContext.quorumSelectedStoreResponse = null;
+        }
     }
 
 }