Fault injection fix (Azure#34023)

* update routeToLocation in faultInjectionContext

---------

Co-authored-by: annie-mac <[email protected]>

Author: xinlian12

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

@@ -74,7 +74,7 @@ static class RegionEndpointValidator implements IFaultInjectionConditionValidato
         }
         @Override
         public boolean isApplicable(RxDocumentServiceRequest request) {
-            return this.regionEndpoints.contains(request.requestContext.locationEndpointToRoute);
+            return this.regionEndpoints.contains(request.faultInjectionRequestContext.getLocationEndpointToRoute());
         }
     }
 

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

@@ -77,7 +77,11 @@ public String getId() {
 
     @Override
     public long getHitCount() {
-        return this.hitCount.get() > this.hitLimit ? this.hitLimit : this.hitCount.get();
+        if (this.hitLimit == null) {
+            return this.hitCount.get();
+        }
+
+        return Math.min(this.hitLimit, this.hitCount.get());
     }
 
     @Override

sdk/cosmos/azure-cosmos-test/src/test/java/com/azure/cosmos/test/faultinjection/FaultInjectionServerErrorRuleTests.java

@@ -170,6 +170,7 @@ public void faultInjectionServerErrorRuleTests_OperationType(OperationType opera
 
             cosmosDiagnostics = this.performDocumentOperation(cosmosAsyncContainer, operationType, createdItem);
             if (operationType == OperationType.Read) {
+                assertThat(serverTooManyRequestsErrorRule.getHitCount()).isEqualTo(1);
                 this.validateFaultInjectionRuleApplied(
                     cosmosDiagnostics,
                     operationType,
@@ -261,6 +262,7 @@ public void faultInjectionServerErrorRuleTests_OperationTypeImpactAddresses(Oper
 
             CosmosDiagnostics cosmosDiagnostics = this.performDocumentOperation(container, operationType, createdItem);
             if (operationType.isWriteOperation()) {
+                assertThat(writeRegionServerGoneErrorRule.getHitCount()).isEqualTo(1);
                 this.validateFaultInjectionRuleApplied(
                     cosmosDiagnostics,
                     operationType,
@@ -358,6 +360,9 @@ public void faultInjectionServerErrorRuleTests_Region() throws JsonProcessingExc
 
             // Validate fault injection applied in the local region
             CosmosDiagnostics cosmosDiagnostics = this.performDocumentOperation(container, OperationType.Read, createdItem);
+            assertThat(serverErrorRuleLocalRegion.getHitCount()).isEqualTo(1);
+            assertThat(serverErrorRuleRemoteRegion.getHitCount()).isEqualTo(0);
+
             this.validateFaultInjectionRuleApplied(
                 cosmosDiagnostics,
                 OperationType.Read,
@@ -418,6 +423,8 @@ public void faultInjectionServerErrorRuleTests_Partition() throws JsonProcessing
 
         CosmosDiagnostics cosmosDiagnostics =
             cosmosAsyncContainer.queryItems(query, queryRequestOptions, TestItem.class).byPage().blockFirst().getCosmosDiagnostics();
+
+        assertThat(serverErrorRuleByFeedRange.getHitCount()).isEqualTo(1);
         this.validateFaultInjectionRuleApplied(
             cosmosDiagnostics,
             OperationType.Query,
@@ -484,6 +491,7 @@ public void faultInjectionServerErrorRuleTests_ServerResponseDelay() throws Json
             CosmosItemResponse<TestItem> itemResponse =
                 container.readItem(createdItem.getId(), new PartitionKey(createdItem.getId()), TestItem.class).block();
 
+            assertThat(timeoutRule.getHitCount()).isEqualTo(1);
             this.validateFaultInjectionRuleApplied(
                 itemResponse.getDiagnostics(),
                 OperationType.Read,
@@ -541,6 +549,7 @@ public void faultInjectionServerErrorRuleTests_ServerConnectionDelay() throws Js
             CosmosFaultInjectionHelper.configureFaultInjectionRules(container, Arrays.asList(serverConnectionDelayRule)).block();
             CosmosItemResponse<TestItem> itemResponse = container.createItem(TestItem.createNewItem()).block();
 
+            assertThat(serverConnectionDelayRule.getHitCount()).isEqualTo(1);
             this.validateFaultInjectionRuleApplied(
                 itemResponse.getDiagnostics(),
                 OperationType.Create,
@@ -612,6 +621,7 @@ public void faultInjectionServerErrorRuleTests_ServerErrorResponse(
                 }
             }
 
+            assertThat(serverErrorRule.getHitCount()).isEqualTo(1);
             this.validateFaultInjectionRuleApplied(
                 cosmosDiagnostics,
                 OperationType.Read,

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

@@ -115,7 +115,12 @@ public static Mono<DatabaseAccount> getDatabaseAccountFromAnyLocationsAsync(
     }
 
     public URI resolveServiceEndpoint(RxDocumentServiceRequest request) {
-        return this.locationCache.resolveServiceEndpoint(request);
+        URI serviceEndpoint = this.locationCache.resolveServiceEndpoint(request);
+        if (request.faultInjectionRequestContext != null) {
+            request.faultInjectionRequestContext.setLocationEndpointToRoute(serviceEndpoint);
+        }
+
+        return serviceEndpoint;
     }
 
     public URI resolveFaultInjectionServiceEndpoint(String region, boolean writeOnly) {

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

@@ -3,6 +3,7 @@
 
 package com.azure.cosmos.implementation.faultinjection;
 
+import java.net.URI;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
@@ -16,6 +17,8 @@ public class FaultInjectionRequestContext {
     private final Map<String, Integer> hitCountByRuleMap;
     private final Map<Long, String> transportRequestIdRuleIdMap;
 
+    private volatile URI locationEndpointToRoute;
+
     /***
      * This usually is called during retries.
      * The hit count cap will need to be copied over so that the total times defined in error result be honored.
@@ -50,6 +53,15 @@ public int getFaultInjectionRuleApplyCount(String ruleId) {
         return this.hitCountByRuleMap.getOrDefault(ruleId, 0);
     }
     public String getFaultInjectionRuleId(long transportRequesetId) {
-        return this.transportRequestIdRuleIdMap.getOrDefault(transportRequesetId, null); }
+        return this.transportRequestIdRuleIdMap.getOrDefault(transportRequesetId, null);
+    }
+
+    public void setLocationEndpointToRoute(URI locationEndpointToRoute) {
+        this.locationEndpointToRoute = locationEndpointToRoute;
+    }
+
+    public URI getLocationEndpointToRoute() {
+        return this.locationEndpointToRoute;
+    }
 }