Fix partitionLevelCircuitBreakerCfg missing from CosmosDiagnostics clientCfgs

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

@@ -4613,6 +4613,65 @@ public void validateHandlingOnNullPartitionKeyRangeOnSmallE2ETimeout_allOps(Oper
         }
     }
 
+    /**
+     * Regression validation for the Per-Partition Circuit Breaker (PPCB) diagnostics fix (see PR 49734).
+     *
+     * The {@code clientCfgs} section of the emitted {@link CosmosDiagnostics} must always include the
+     * {@code partitionLevelCircuitBreakerCfg} field for every client, regardless of whether PPCB is
+     * explicitly enabled. A prior regression silently dropped this field unless PPAF mandated it. This
+     * test asserts that all the expected {@code clientCfgs} keys - including
+     * {@code partitionLevelCircuitBreakerCfg} - are present in the diagnostics of a real operation
+     * without setting any PPCB configuration.
+     */
+    @Test(groups = { "circuit-breaker-misc-gateway", "circuit-breaker-misc-direct", "circuit-breaker-read-all-read-many", "multi-region", "fi-thinclient-multi-master" }, timeOut = TIMEOUT)
+    public void partitionLevelCircuitBreakerConfigIsPresentInClientCfgsDiagnostics() {
+
+        try (CosmosAsyncClient client = getClientBuilder().buildAsyncClient()) {
+
+            CosmosAsyncContainer container = client
+                .getDatabase(this.sharedAsyncDatabaseId)
+                .getContainer(this.sharedMultiPartitionAsyncContainerIdWhereIdIsPartitionKey);
+
+            TestObject item = TestObject.create();
+
+            CosmosItemResponse<TestObject> createResponse = container
+                .createItem(item, new PartitionKey(item.getId()), new CosmosItemRequestOptions())
+                .block();
+
+            assertThat(createResponse).isNotNull();
+
+            String diagnosticsString = createResponse.getDiagnostics().toString();
+
+            assertThat(diagnosticsString)
+                .as("clientCfgs section should be present in the CosmosDiagnostics")
+                .contains("\"clientCfgs\"");
+
+            // All the clientCfgs keys unconditionally emitted by DiagnosticsClientConfigSerializer,
+            // including partitionLevelCircuitBreakerCfg (the field the regression previously dropped).
+            List<String> expectedClientCfgsKeys = Arrays.asList(
+                "id",
+                "machineId",
+                "connectionMode",
+                "numberOfClients",
+                "isPpafEnabled",
+                "isFalseProgSessionTokenMergeEnabled",
+                "excrgns",
+                "clientEndpoints",
+                "connCfg",
+                "consistencyCfg",
+                "proactiveInitCfg",
+                "e2ePolicyCfg",
+                "sessionRetryCfg",
+                "partitionLevelCircuitBreakerCfg");
+
+            for (String expectedKey : expectedClientCfgsKeys) {
+                assertThat(diagnosticsString)
+                    .as("clientCfgs key '%s' should be present in the CosmosDiagnostics", expectedKey)
+                    .contains("\"" + expectedKey + "\"");
+            }
+        }
+    }
+
     private static Function<OperationInvocationParamsWrapper, ResponseWrapper<?>> resolveDataPlaneOperation(FaultInjectionOperationType faultInjectionOperationType) {
 
         switch (faultInjectionOperationType) {

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

@@ -5,6 +5,7 @@
 import com.azure.core.credential.AzureKeyCredential;
 import com.azure.core.http.ProxyOptions;
 import com.azure.cosmos.BridgeInternal;
+import com.azure.cosmos.ConnectionMode;
 import com.azure.cosmos.ConsistencyLevel;
 import com.azure.cosmos.CosmosContainerProactiveInitConfig;
 import com.azure.cosmos.CosmosDiagnostics;
@@ -35,6 +36,11 @@
 import com.azure.cosmos.models.ModelBridgeInternal;
 import com.azure.cosmos.models.PartitionKey;
 import com.azure.cosmos.models.PartitionKeyDefinition;
+import com.fasterxml.jackson.core.JsonFactory;
+import com.fasterxml.jackson.core.JsonGenerator;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.SerializerProvider;
+import com.fasterxml.jackson.databind.node.ObjectNode;
 import io.netty.buffer.ByteBufInputStream;
 import io.netty.buffer.Unpooled;
 import io.netty.handler.codec.http.HttpResponseStatus;
@@ -47,6 +53,8 @@
 import reactor.test.StepVerifier;
 
 import java.net.URI;
+import java.io.StringWriter;
+import java.lang.reflect.Method;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.ArrayList;
@@ -504,6 +512,109 @@ public void lookupCollectionRoutingMapWithRetryStopsAfterBoundedAttempts() {
         }
     }
 
+    // Regression test for the "partitionLevelCircuitBreakerCfg" diagnostics field silently disappearing from the
+    // CosmosDiagnostics "clientCfgs" section. Prior to the fix, the field was only written on the PPAF
+    // (service-mandated) path, so a client that did not have PPAF-mandated PPCB never surfaced it. The field must
+    // now be present for every client regardless of any PPCB configuration. This test constructs a real
+    // RxDocumentClientImpl (exercising the actual constructor wiring), drives the private
+    // initializePerPartitionCircuitBreaker() init path without setting any PPCB configuration, serializes the
+    // resulting DiagnosticsClientConfig, and asserts that every expected "clientCfgs" key - including
+    // partitionLevelCircuitBreakerCfg - is present (guarding against future serialization truncation as well as
+    // the specific regression).
+    @Test(groups = {"unit"})
+    public void diagnosticsClientConfigContainsAllClientCfgKeysIncludingPartitionLevelCircuitBreaker() throws Exception {
+        Mockito.when(this.connectionPolicyMock.getIdleHttpConnectionTimeout()).thenReturn(Duration.ZERO);
+        Mockito.when(this.connectionPolicyMock.getMaxConnectionPoolSize()).thenReturn(1);
+        Mockito.when(this.connectionPolicyMock.getProxy()).thenReturn(null);
+        Mockito.when(this.connectionPolicyMock.getHttpNetworkRequestTimeout()).thenReturn(Duration.ZERO);
+        Mockito.when(this.connectionPolicyMock.getHttp2ConnectionConfig()).thenReturn(new Http2ConnectionConfig());
+        // The serializer eagerly calls getConnectionMode().toString() for the very first "connectionMode" key; if this
+        // returns null (Mockito default), serialization would NPE and silently drop every subsequent key.
+        Mockito.when(this.connectionPolicyMock.getConnectionMode()).thenReturn(ConnectionMode.DIRECT);
+
+        MockedStatic<HttpClient> httpClientMock = Mockito.mockStatic(HttpClient.class);
+        httpClientMock
+            .when(() -> HttpClient.createFixed(Mockito.any(HttpClientConfig.class)))
+            .thenReturn(dummyHttpClient());
+
+        RxDocumentClientImpl rxDocumentClient = null;
+
+        try {
+            rxDocumentClient = new RxDocumentClientImpl(
+                this.serviceEndpointMock,
+                this.masterKeyOrResourceTokenMock,
+                this.permissionFeedMock,
+                this.connectionPolicyMock,
+                this.consistencyLevelMock,
+                null,
+                this.configsMock,
+                this.cosmosAuthorizationTokenResolverMock,
+                this.azureKeyCredentialMock,
+                false,
+                false,
+                false,
+                this.metadataCachesSnapshotMock,
+                this.apiTypeMock,
+                this.cosmosClientTelemetryConfigMock,
+                this.clientCorrelationIdMock,
+                this.endToEndOperationLatencyPolicyConfig,
+                this.sessionRetryOptionsMock,
+                this.containerProactiveInitConfigMock,
+                this.defaultItemSerializer,
+                false
+            );
+
+            // Drive the exact wiring that regressed: explicit (client-side) Per-Partition Circuit Breaker
+            // initialization. The constructor does not invoke init() (which would require network), so invoke the
+            // private no-arg initializer reflectively.
+            Method initPpcb = RxDocumentClientImpl.class.getDeclaredMethod("initializePerPartitionCircuitBreaker");
+            initPpcb.setAccessible(true);
+            initPpcb.invoke(rxDocumentClient);
+
+            ObjectMapper objectMapper = new ObjectMapper();
+            StringWriter jsonWriter = new StringWriter();
+            JsonGenerator jsonGenerator = new JsonFactory().createGenerator(jsonWriter);
+            SerializerProvider serializerProvider = objectMapper.getSerializerProvider();
+            DiagnosticsClientContext.DiagnosticsClientConfigSerializer.INSTANCE
+                .serialize(rxDocumentClient.getConfig(), jsonGenerator, serializerProvider);
+            jsonGenerator.flush();
+            ObjectNode clientCfgs = (ObjectNode) objectMapper.readTree(jsonWriter.toString());
+
+            String serializedJson = clientCfgs.toString();
+
+            // Every key the serializer unconditionally writes, including the (previously regressed)
+            // partitionLevelCircuitBreakerCfg which must be present for every client regardless of PPCB config.
+            String[] expectedKeys = new String[] {
+                "id",
+                "machineId",
+                "connectionMode",
+                "numberOfClients",
+                "isPpafEnabled",
+                "isFalseProgSessionTokenMergeEnabled",
+                "excrgns",
+                "clientEndpoints",
+                "connCfg",
+                "consistencyCfg",
+                "proactiveInitCfg",
+                "e2ePolicyCfg",
+                "sessionRetryCfg",
+                "partitionLevelCircuitBreakerCfg"
+            };
+
+            for (String expectedKey : expectedKeys) {
+                assertThat(clientCfgs.has(expectedKey))
+                    .withFailMessage("Expected clientCfgs key '%s' to be present. Serialized clientCfgs: %s",
+                        expectedKey, serializedJson)
+                    .isTrue();
+            }
+        } finally {
+            if (rxDocumentClient != null) {
+                rxDocumentClient.close();
+            }
+            httpClientMock.close();
+        }
+    }
+
     private static HttpClient dummyHttpClient() {
         return new HttpClient() {
             @Override

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -8,6 +8,7 @@
 
 #### Bugs Fixed
 * Unified request-level consistency override behavior across transports: invalid attempts to upgrade the request consistency level above the account default are now silently ignored instead of returning `BadRequest` in some gateway paths. - See PR [49606](https://github.com/Azure/azure-sdk-for-java/pull/49606).
+* Fixed `partitionLevelCircuitBreakerCfg` missing from the `clientCfgs` section of `CosmosDiagnostics` when Per-Partition Circuit Breaker is explicitly enabled. - See PR [49734](https://github.com/Azure/azure-sdk-for-java/pull/49734).
 * Fixed thin-client (Gateway V2) queries with a prefix (partial) hierarchical partition key returning co-located documents from other logical partitions. - See PR [49688](https://github.com/Azure/azure-sdk-for-java/pull/49688).
 
 #### Other Changes

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

@@ -8988,7 +8988,6 @@ private synchronized void initializePerPartitionFailover(DatabaseAccount databas
         checkNotNull(this.globalPartitionEndpointManagerForPerPartitionAutomaticFailover, "Argument 'globalPartitionEndpointManagerForPerPartitionAutomaticFailover' cannot be null.");
         checkNotNull(this.globalPartitionEndpointManagerForPerPartitionCircuitBreaker, "Argument 'globalPartitionEndpointManagerForPerPartitionCircuitBreaker' cannot be null.");
 
-        this.diagnosticsClientConfig.withPartitionLevelCircuitBreakerConfig(this.globalPartitionEndpointManagerForPerPartitionCircuitBreaker.getCircuitBreakerConfig());
         this.diagnosticsClientConfig.withIsPerPartitionAutomaticFailoverEnabled(this.globalPartitionEndpointManagerForPerPartitionAutomaticFailover.isPerPartitionAutomaticFailoverEnabled());
     }
 
@@ -9017,6 +9016,12 @@ private void initializePerPartitionCircuitBreaker() {
 
         this.globalPartitionEndpointManagerForPerPartitionCircuitBreaker.resetCircuitBreakerConfig(partitionLevelCircuitBreakerConfig);
         this.globalPartitionEndpointManagerForPerPartitionCircuitBreaker.init();
+
+        // Populate the circuit breaker config in the diagnostics client config here (rather than in
+        // initializePerPartitionFailover) so the "partitionLevelCircuitBreakerCfg" field appears in
+        // CosmosDiagnostics whenever the circuit breaker is configured client-side, not only when
+        // Per-Partition Automatic Failover is mandated by the service.
+        this.diagnosticsClientConfig.withPartitionLevelCircuitBreakerConfig(this.globalPartitionEndpointManagerForPerPartitionCircuitBreaker.getCircuitBreakerConfig());
     }
 
     private void enableAvailabilityStrategyForReads() {