Early yield on 429 throttling on barrier requests

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java

@@ -7,6 +7,7 @@
 import com.azure.cosmos.ConsistencyLevel;
 import com.azure.cosmos.CosmosAsyncClient;
 import com.azure.cosmos.CosmosAsyncContainer;
+import com.azure.cosmos.CosmosException;
 import com.azure.cosmos.CosmosClientBuilder;
 import com.azure.cosmos.CosmosDiagnostics;
 import com.azure.cosmos.DirectConnectionConfig;
@@ -1915,6 +1916,249 @@ private static boolean isOperationAWriteOperation(OperationType operationType) {
             || operationType == OperationType.Batch;
     }
 
+    @Test(groups = {"multi-region-strong"}, timeOut = 2 * TIMEOUT)
+    public void faultInjection_readBarrierThrottled_yieldsEarly() throws JsonProcessingException {
+        // Validates that when barrier HEAD requests are throttled (429) during a strong read,
+        // the early yield mechanism fires and the 429 is propagated to the caller.
+        // This test requires a multi-region strong consistency account because read barriers
+        // are only triggered when numberOfReadRegions > 0.
+
+        if (this.databaseAccount.getConsistencyPolicy().getDefaultConsistencyLevel() != ConsistencyLevel.STRONG) {
+            throw new SkipException("Test only applicable to STRONG consistency level.");
+        }
+
+        if (this.accountLevelReadRegions.size() <= 1) {
+            throw new SkipException("Test requires multi-region account for read barriers to be triggered.");
+        }
+
+        CosmosAsyncClient newClient = null;
+        String faultInjectionRuleId = "barrier-429-read-yield-" + UUID.randomUUID();
+        FaultInjectionRule barrierThrottleRule =
+            new FaultInjectionRuleBuilder(faultInjectionRuleId)
+                .condition(
+                    new FaultInjectionConditionBuilder()
+                        .operationType(FaultInjectionOperationType.HEAD_COLLECTION)
+                        .build()
+                )
+                .result(
+                    FaultInjectionResultBuilders
+                        .getResultBuilder(FaultInjectionServerErrorType.TOO_MANY_REQUEST)
+                        .times(10)
+                        .build()
+                )
+                .duration(Duration.ofMinutes(5))
+                .build();
+
+        try {
+            System.setProperty("COSMOS.ENABLE_BARRIER_EARLY_YIELD_ON_429", "true");
+            newClient = new CosmosClientBuilder()
+                .endpoint(TestConfigurations.HOST)
+                .key(TestConfigurations.MASTER_KEY)
+                .contentResponseOnWriteEnabled(true)
+                .buildAsyncClient();
+
+            CosmosAsyncContainer container =
+                newClient
+                    .getDatabase(cosmosAsyncContainer.getDatabase().getId())
+                    .getContainer(cosmosAsyncContainer.getId());
+
+            TestObject testItem = TestObject.create();
+            container.createItem(testItem).block();
+
+            CosmosFaultInjectionHelper.configureFaultInjectionRules(container, Arrays.asList(barrierThrottleRule)).block();
+
+            // Force barrier requests by making GCLSN < LSN (simulate replication lag)
+            CosmosInterceptorHelper.registerTransportClientInterceptor(
+                newClient,
+                (request, storeResponse) -> {
+                    if (request.getResourceType() == ResourceType.Document
+                        && request.getOperationType() == OperationType.Read) {
+                        storeResponse.setGCLSN(storeResponse.getLSN() - 2L);
+                    }
+                    return storeResponse;
+                }
+            );
+
+            CosmosDiagnostics cosmosDiagnostics = this.performDocumentOperation(container, OperationType.Read, testItem, false);
+
+            // Validate: the 429 fault injection rule was applied on barrier requests
+            validateFaultInjectionRuleAppliedForBarrier(
+                cosmosDiagnostics,
+                OperationType.Read,
+                HttpConstants.StatusCodes.TOO_MANY_REQUESTS,
+                HttpConstants.SubStatusCodes.USER_REQUEST_RATE_TOO_LARGE,
+                faultInjectionRuleId);
+
+            assertThat(barrierThrottleRule.getHitCount()).isGreaterThanOrEqualTo(1);
+
+        } finally {
+            barrierThrottleRule.disable();
+            safeClose(newClient);
+            System.clearProperty("COSMOS.ENABLE_BARRIER_EARLY_YIELD_ON_429");
+        }
+    }
+
+    @Test(groups = {"multi-region-strong"}, timeOut = 2 * TIMEOUT)
+    public void faultInjection_writeBarrierThrottled_returns408() throws JsonProcessingException {
+        // Validates that when barrier HEAD requests are throttled (429) during a strong write,
+        // the writer exhausts retries and returns 408 with SERVER_WRITE_BARRIER_THROTTLED substatus.
+
+        if (this.databaseAccount.getConsistencyPolicy().getDefaultConsistencyLevel() != ConsistencyLevel.STRONG) {
+            throw new SkipException("Test only applicable to STRONG consistency level.");
+        }
+
+        if (this.accountLevelReadRegions.size() <= 1) {
+            throw new SkipException("Test requires multi-region account for write barriers to be triggered.");
+        }
+
+        CosmosAsyncClient newClient = null;
+        String faultInjectionRuleId = "barrier-429-write-408-" + UUID.randomUUID();
+        FaultInjectionRule barrierThrottleRule =
+            new FaultInjectionRuleBuilder(faultInjectionRuleId)
+                .condition(
+                    new FaultInjectionConditionBuilder()
+                        .operationType(FaultInjectionOperationType.HEAD_COLLECTION)
+                        .build()
+                )
+                .result(
+                    FaultInjectionResultBuilders
+                        .getResultBuilder(FaultInjectionServerErrorType.TOO_MANY_REQUEST)
+                        .times(100) // Enough to exhaust all barrier retries
+                        .build()
+                )
+                .duration(Duration.ofMinutes(5))
+                .build();
+
+        try {
+            System.setProperty("COSMOS.ENABLE_BARRIER_EARLY_YIELD_ON_429", "true");
+            newClient = new CosmosClientBuilder()
+                .endpoint(TestConfigurations.HOST)
+                .key(TestConfigurations.MASTER_KEY)
+                .contentResponseOnWriteEnabled(true)
+                .buildAsyncClient();
+
+            CosmosAsyncContainer container =
+                newClient
+                    .getDatabase(cosmosAsyncContainer.getDatabase().getId())
+                    .getContainer(cosmosAsyncContainer.getId());
+
+            TestObject testItem = TestObject.create();
+            container.createItem(testItem).block();
+
+            CosmosFaultInjectionHelper.configureFaultInjectionRules(container, Arrays.asList(barrierThrottleRule)).block();
+
+            // Force barrier requests by making GCLSN < LSN (simulate replication lag)
+            CosmosInterceptorHelper.registerTransportClientInterceptor(
+                newClient,
+                (request, storeResponse) -> {
+                    if (request.getResourceType() == ResourceType.Document
+                        && request.getOperationType() == OperationType.Create) {
+                        storeResponse.setGCLSN(storeResponse.getLSN() - 2L);
+                    }
+                    return storeResponse;
+                }
+            );
+
+            TestObject newItem = TestObject.create();
+            try {
+                container.createItem(newItem).block();
+                // If we get here, the barrier was met before throttle exhaustion (possible if
+                // GCLSN catches up between retries). That's acceptable — skip the assertion.
+            } catch (Exception e) {
+                // Validate: the write barrier failed due to throttle exhaustion and surfaced as a
+                // 408 (RequestTimeout) with the SERVER_WRITE_BARRIER_THROTTLED (21013) substatus.
+                assertThat(barrierThrottleRule.getHitCount()).isGreaterThanOrEqualTo(1);
+                assertThat(e).isInstanceOf(CosmosException.class);
+                CosmosException cosmosException = (CosmosException) e;
+                assertThat(cosmosException.getStatusCode())
+                    .isEqualTo(HttpConstants.StatusCodes.REQUEST_TIMEOUT);
+                assertThat(cosmosException.getSubStatusCode())
+                    .isEqualTo(HttpConstants.SubStatusCodes.SERVER_WRITE_BARRIER_THROTTLED);
+            }
+
+        } finally {
+            barrierThrottleRule.disable();
+            safeClose(newClient);
+            System.clearProperty("COSMOS.ENABLE_BARRIER_EARLY_YIELD_ON_429");
+        }
+    }
+
+    @Test(groups = {"multi-region-strong"}, timeOut = 2 * TIMEOUT)
+    public void faultInjection_readBarrierThrottled_thenRecovers() throws JsonProcessingException {
+        // Validates that when barrier requests are initially throttled but the rule has a hitLimit,
+        // the read eventually succeeds after the throttle clears.
+
+        if (this.databaseAccount.getConsistencyPolicy().getDefaultConsistencyLevel() != ConsistencyLevel.STRONG) {
+            throw new SkipException("Test only applicable to STRONG consistency level.");
+        }
+
+        if (this.accountLevelReadRegions.size() <= 1) {
+            throw new SkipException("Test requires multi-region account for read barriers to be triggered.");
+        }
+
+        CosmosAsyncClient newClient = null;
+        String faultInjectionRuleId = "barrier-429-recover-" + UUID.randomUUID();
+        FaultInjectionRule barrierThrottleRule =
+            new FaultInjectionRuleBuilder(faultInjectionRuleId)
+                .condition(
+                    new FaultInjectionConditionBuilder()
+                        .operationType(FaultInjectionOperationType.HEAD_COLLECTION)
+                        .build()
+                )
+                .result(
+                    FaultInjectionResultBuilders
+                        .getResultBuilder(FaultInjectionServerErrorType.TOO_MANY_REQUEST)
+                        .times(1)
+                        .build()
+                )
+                .hitLimit(2) // Only throttle the first 2 barrier attempts
+                .duration(Duration.ofMinutes(5))
+                .build();
+
+        try {
+            System.setProperty("COSMOS.ENABLE_BARRIER_EARLY_YIELD_ON_429", "true");
+            newClient = new CosmosClientBuilder()
+                .endpoint(TestConfigurations.HOST)
+                .key(TestConfigurations.MASTER_KEY)
+                .contentResponseOnWriteEnabled(true)
+                .buildAsyncClient();
+
+            CosmosAsyncContainer container =
+                newClient
+                    .getDatabase(cosmosAsyncContainer.getDatabase().getId())
+                    .getContainer(cosmosAsyncContainer.getId());
+
+            TestObject testItem = TestObject.create();
+            container.createItem(testItem).block();
+
+            CosmosFaultInjectionHelper.configureFaultInjectionRules(container, Arrays.asList(barrierThrottleRule)).block();
+
+            // Force barrier requests by making GCLSN < LSN
+            CosmosInterceptorHelper.registerTransportClientInterceptor(
+                newClient,
+                (request, storeResponse) -> {
+                    if (request.getResourceType() == ResourceType.Document
+                        && request.getOperationType() == OperationType.Read) {
+                        storeResponse.setGCLSN(storeResponse.getLSN() - 2L);
+                    }
+                    return storeResponse;
+                }
+            );
+
+            // The read should eventually succeed after the throttle clears (hitLimit=2)
+            CosmosDiagnostics cosmosDiagnostics = this.performDocumentOperation(container, OperationType.Read, testItem, false);
+
+            // The rule should have been applied at least once
+            assertThat(barrierThrottleRule.getHitCount()).isGreaterThanOrEqualTo(1);
+            assertThat(barrierThrottleRule.getHitCount()).isLessThanOrEqualTo(2);
+
+        } finally {
+            barrierThrottleRule.disable();
+            safeClose(newClient);
+            System.clearProperty("COSMOS.ENABLE_BARRIER_EARLY_YIELD_ON_429");
+        }
+    }
+
     private static class AccountLevelLocationContext {
         private final List<String> serviceOrderedReadableRegions;
         private final List<String> serviceOrderedWriteableRegions;