Improve PPCB failback diagnostics

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PerPartitionCircuitBreakerE2ETests.java

@@ -58,6 +58,7 @@
 import com.azure.cosmos.test.faultinjection.FaultInjectionRuleBuilder;
 import com.azure.cosmos.test.faultinjection.FaultInjectionServerErrorResult;
 import com.azure.cosmos.test.faultinjection.FaultInjectionServerErrorType;
+import com.fasterxml.jackson.databind.JsonNode;
 import org.testng.SkipException;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.AfterMethod;
@@ -3807,6 +3808,7 @@ private void execute(
                         testId,
                         executeDataPlaneOperation,
                         operationInvocationParamsWrapper);
+                    List<JsonNode> ppcbStateByRegionNodes = getPpcbStateByRegionNodes(response);
 
                     ConsecutiveExceptionBasedCircuitBreaker consecutiveExceptionBasedCircuitBreaker
                         = globalPartitionEndpointManagerForPerPartitionCircuitBreaker.getConsecutiveExceptionBasedCircuitBreaker();
@@ -3832,6 +3834,9 @@ private void execute(
 
                     if (executionCountAfterCircuitBreakingThresholdBreached > 1) {
                         validateResponseInAbsenceOfFailures.accept(response);
+                        assertPpcbHealthStatus(
+                            ppcbStateByRegionNodes,
+                            LocationHealthStatus.Unavailable);
                     }
 
                     if (response.cosmosItemResponse != null) {
@@ -3898,6 +3903,10 @@ private void execute(
                         executeDataPlaneOperation,
                         operationInvocationParamsWrapper);
                     validateResponseInAbsenceOfFailures.accept(response);
+                    assertPpcbHealthStatus(
+                        getPpcbStateByRegionNodes(response),
+                        LocationHealthStatus.HealthyTentative,
+                        LocationHealthStatus.Healthy);
 
                     if (response.cosmosItemResponse != null) {
                         assertThat(response.cosmosItemResponse).isNotNull();
@@ -3958,6 +3967,57 @@ private static CosmosDiagnosticsContext getDiagnosticsContext(ResponseWrapper<?>
         return null;
     }
 
+    private static List<JsonNode> getPpcbStateByRegionNodes(ResponseWrapper<?> response) {
+        CosmosDiagnosticsContext diagnosticsContext = getDiagnosticsContext(response);
+        assertThat(diagnosticsContext).isNotNull();
+
+        try {
+            JsonNode diagnostics = Utils.getSimpleObjectMapper().readTree(diagnosticsContext.toJson());
+            List<JsonNode> stateByRegionNodes = new ArrayList<>();
+            for (JsonNode ppcbNode : diagnostics.findValues("ppcb")) {
+                JsonNode stateByRegion = ppcbNode.get("stateByRegion");
+                if (stateByRegion != null && stateByRegion.isObject()) {
+                    stateByRegionNodes.add(stateByRegion);
+                }
+            }
+
+            assertThat(stateByRegionNodes)
+                .as("Expected every PPCB-enabled data-plane operation to include ppcb.stateByRegion. Diagnostics: %s", diagnostics)
+                .isNotEmpty();
+            return stateByRegionNodes;
+        } catch (Exception e) {
+            throw new AssertionError("Failed to parse CosmosDiagnostics for PPCB state", e);
+        }
+    }
+
+    private static void assertPpcbHealthStatus(
+        List<JsonNode> stateByRegionNodes,
+        LocationHealthStatus... expectedStatuses) {
+
+        List<String> actualStatuses = new ArrayList<>();
+        for (JsonNode stateByRegion : stateByRegionNodes) {
+            Iterator<JsonNode> regionStates = stateByRegion.elements();
+            while (regionStates.hasNext()) {
+                JsonNode healthStatus = regionStates.next().get("locationHealthStatus");
+                if (healthStatus != null) {
+                    actualStatuses.add(healthStatus.asText());
+                }
+            }
+        }
+
+        boolean expectedStatusFound = false;
+        for (LocationHealthStatus expectedStatus : expectedStatuses) {
+            if (actualStatuses.contains(expectedStatus.toString())) {
+                expectedStatusFound = true;
+                break;
+            }
+        }
+
+        assertThat(expectedStatusFound)
+            .as("Expected PPCB health status to be one of %s but found %s", Arrays.toString(expectedStatuses), actualStatuses)
+            .isTrue();
+    }
+
     private ResponseWrapper<?> executeDataPlaneOperationWithTransient4041002Retry(
         String testId,
         Function<OperationInvocationParamsWrapper, ResponseWrapper<?>> executeDataPlaneOperation,

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/PerPartitionCircuitBreakerInfoHolderTest.java

@@ -0,0 +1,188 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.implementation.perPartitionCircuitBreaker;
+
+import com.azure.cosmos.implementation.ClientSideRequestStatistics;
+import com.azure.cosmos.implementation.CrossRegionAvailabilityContextForRxDocumentServiceRequest;
+import com.azure.cosmos.implementation.DiagnosticsClientContext;
+import com.azure.cosmos.implementation.GlobalEndpointManager;
+import com.azure.cosmos.implementation.OperationType;
+import com.azure.cosmos.implementation.PartitionKeyRange;
+import com.azure.cosmos.implementation.ResourceType;
+import com.azure.cosmos.implementation.RxDocumentServiceRequest;
+import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList;
+import com.azure.cosmos.implementation.directconnectivity.StoreResponseDiagnostics;
+import com.azure.cosmos.implementation.perPartitionAutomaticFailover.PerPartitionAutomaticFailoverInfoHolder;
+import com.azure.cosmos.implementation.routing.RegionalRoutingContext;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.module.SimpleModule;
+import org.mockito.Mockito;
+import org.testng.annotations.Test;
+
+import java.time.Instant;
+import java.net.URI;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.Mockito.doReturn;
+
+public class PerPartitionCircuitBreakerInfoHolderTest {
+
+    @Test(groups = {"unit"})
+    public void storesImmutableStateSnapshot() {
+        LocationSpecificHealthContext healthContext = createHealthContext(LocationHealthStatus.Unavailable);
+        Map<String, LocationSpecificHealthContext> currentState = new LinkedHashMap<>();
+        currentState.put("eastus", healthContext);
+
+        PerPartitionCircuitBreakerInfoHolder holder = new PerPartitionCircuitBreakerInfoHolder();
+        holder.setPerPartitionCircuitBreakerInfoHolder(currentState);
+        currentState.clear();
+
+        assertThat(holder.getPerPartitionCircuitBreakerInfoHolder())
+            .containsOnlyKeys("eastus")
+            .containsValue(healthContext);
+        assertThatThrownBy(() -> holder.getPerPartitionCircuitBreakerInfoHolder().clear())
+            .isInstanceOf(UnsupportedOperationException.class);
+    }
+
+    @Test(groups = {"unit"})
+    public void initializedEmptyStateIsSerialized() throws Exception {
+        PerPartitionCircuitBreakerInfoHolder holder = new PerPartitionCircuitBreakerInfoHolder();
+        holder.setPerPartitionCircuitBreakerInfoHolder(Collections.emptyMap());
+
+        ObjectMapper objectMapper = new ObjectMapper();
+        objectMapper.registerModule(new SimpleModule().addSerializer(
+            PerPartitionCircuitBreakerInfoHolder.class,
+            new PerPartitionCircuitBreakerInfoHolder.PerPartitionCircuitBreakerInfoHolderSerializer()));
+
+        assertThat(objectMapper.writeValueAsString(holder))
+            .isEqualTo("{\"stateByRegion\":{}}");
+        assertThat(objectMapper.writeValueAsString(PerPartitionCircuitBreakerInfoHolder.EMPTY))
+            .isEqualTo("null");
+    }
+
+    @Test(groups = {"unit"})
+    public void responseStatisticsRetainStateAtRecordTime() {
+        DiagnosticsClientContext diagnosticsClientContext = Mockito.mock(DiagnosticsClientContext.class);
+        PerPartitionCircuitBreakerInfoHolder holder = new PerPartitionCircuitBreakerInfoHolder();
+        holder.setPerPartitionCircuitBreakerInfoHolder(Collections.singletonMap(
+            "eastus",
+            createHealthContext(LocationHealthStatus.Unavailable)));
+
+        RxDocumentServiceRequest request = RxDocumentServiceRequest.create(
+            diagnosticsClientContext,
+            OperationType.Read,
+            ResourceType.Document);
+        request.requestContext.setCrossRegionAvailabilityContext(
+            new CrossRegionAvailabilityContextForRxDocumentServiceRequest(
+                null,
+                null,
+                null,
+                new AtomicBoolean(false),
+                holder,
+                new PerPartitionAutomaticFailoverInfoHolder()));
+
+        ClientSideRequestStatistics statistics = new ClientSideRequestStatistics(diagnosticsClientContext);
+        statistics.recordResponse(request, null, null);
+        holder.setPerPartitionCircuitBreakerInfoHolder(Collections.singletonMap(
+            "westus",
+            createHealthContext(LocationHealthStatus.Healthy)));
+
+        PerPartitionCircuitBreakerInfoHolder recordedHolder = statistics.getResponseStatisticsList()
+            .iterator()
+            .next()
+            .getPerPartitionCircuitBreakerInfoHolder();
+        assertThat(recordedHolder.getPerPartitionCircuitBreakerInfoHolder()).containsOnlyKeys("eastus");
+    }
+
+    @Test(groups = {"unit"})
+    public void gatewayStatisticsRetainStateAtRecordTime() throws Exception {
+        DiagnosticsClientContext diagnosticsClientContext = Mockito.mock(DiagnosticsClientContext.class);
+        PerPartitionCircuitBreakerInfoHolder holder = new PerPartitionCircuitBreakerInfoHolder();
+        holder.setPerPartitionCircuitBreakerInfoHolder(Collections.singletonMap(
+            "eastus",
+            createHealthContext(LocationHealthStatus.Unavailable)));
+        RxDocumentServiceRequest request = createRequest(diagnosticsClientContext, holder);
+
+        ClientSideRequestStatistics statistics = new ClientSideRequestStatistics(diagnosticsClientContext);
+        statistics.recordGatewayResponse(request, Mockito.mock(StoreResponseDiagnostics.class), null);
+        holder.setPerPartitionCircuitBreakerInfoHolder(Collections.singletonMap(
+            "westus",
+            createHealthContext(LocationHealthStatus.Healthy)));
+
+        PerPartitionCircuitBreakerInfoHolder recordedHolder = statistics.getGatewayStatisticsList()
+            .get(0)
+            .getPerPartitionCircuitBreakerInfoHolder();
+        assertThat(recordedHolder.getPerPartitionCircuitBreakerInfoHolder()).containsOnlyKeys("eastus");
+        assertThat(new ObjectMapper().writeValueAsString(statistics))
+            .contains("\"ppcb\":{\"stateByRegion\":{\"eastus\":");
+    }
+
+    @Test(groups = {"unit"})
+    public void routingLookupInitializesEmptyStateWhenNoCircuitExists() throws Exception {
+        DiagnosticsClientContext diagnosticsClientContext = Mockito.mock(DiagnosticsClientContext.class);
+        PerPartitionCircuitBreakerInfoHolder holder = new PerPartitionCircuitBreakerInfoHolder();
+        RxDocumentServiceRequest request = createRequest(diagnosticsClientContext, holder);
+        request.setResourceId("collectionRid");
+        PartitionKeyRange partitionKeyRange = new PartitionKeyRange("0", "AA", "BB");
+        request.requestContext.resolvedPartitionKeyRange = partitionKeyRange;
+        request.requestContext.resolvedPartitionKeyRangeForCircuitBreaker = partitionKeyRange;
+
+        RegionalRoutingContext eastUs = new RegionalRoutingContext(URI.create("https://eastus.documents.azure.com"));
+        RegionalRoutingContext westUs = new RegionalRoutingContext(URI.create("https://westus.documents.azure.com"));
+        GlobalEndpointManager globalEndpointManager = Mockito.mock(GlobalEndpointManager.class);
+        doReturn(false).when(globalEndpointManager).canUseMultipleWriteLocations(request);
+        doReturn(UnmodifiableList.unmodifiableList(Arrays.asList(eastUs, westUs)))
+            .when(globalEndpointManager)
+            .getApplicableReadRegionalRoutingContexts(Collections.emptyList());
+
+        GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker manager
+            = new GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker(globalEndpointManager);
+        manager.resetCircuitBreakerConfig(PartitionLevelCircuitBreakerConfig.fromJsonString(
+            "{\"isPartitionLevelCircuitBreakerEnabled\":true,"
+                + "\"consecutiveExceptionCountToleratedForReads\":10,"
+                + "\"consecutiveExceptionCountToleratedForWrites\":5}"));
+
+        assertThat(manager.getUnavailableRegionsForPartitionKeyRange(request, "collectionRid", partitionKeyRange))
+            .isEmpty();
+        assertThat(holder.isInitialized()).isTrue();
+        assertThat(holder.getPerPartitionCircuitBreakerInfoHolder()).isEmpty();
+
+        ClientSideRequestStatistics statistics = new ClientSideRequestStatistics(diagnosticsClientContext);
+        statistics.recordResponse(request, null, null);
+        assertThat(new ObjectMapper().writeValueAsString(statistics))
+            .contains("\"ppcb\":{\"stateByRegion\":{}}");
+    }
+
+    private static RxDocumentServiceRequest createRequest(
+        DiagnosticsClientContext diagnosticsClientContext,
+        PerPartitionCircuitBreakerInfoHolder holder) {
+
+        RxDocumentServiceRequest request = RxDocumentServiceRequest.create(
+            diagnosticsClientContext,
+            OperationType.Read,
+            ResourceType.Document);
+        request.requestContext.setCrossRegionAvailabilityContext(
+            new CrossRegionAvailabilityContextForRxDocumentServiceRequest(
+                null,
+                null,
+                null,
+                new AtomicBoolean(false),
+                holder,
+                new PerPartitionAutomaticFailoverInfoHolder()));
+        return request;
+    }
+
+    private static LocationSpecificHealthContext createHealthContext(LocationHealthStatus healthStatus) {
+        return new LocationSpecificHealthContext.Builder()
+            .withLocationHealthStatus(healthStatus)
+            .withUnavailableSince(Instant.EPOCH)
+            .build();
+    }
+}