Partition level circuit breaker

sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncBenchmark.java

@@ -88,6 +88,16 @@ abstract class AsyncBenchmark<T> {
         logger = LoggerFactory.getLogger(this.getClass());
         configuration = cfg;
 
+        if (configuration.isPartitionLevelCircuitBreakerEnabled()) {
+            System.setProperty(
+                "COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG",
+                "{\"isPartitionLevelCircuitBreakerEnabled\": true, "
+                    + "\"circuitBreakerType\": \"CONSECUTIVE_EXCEPTION_COUNT_BASED\","
+                    + "\"consecutiveExceptionCountToleratedForReads\": 10,"
+                    + "\"consecutiveExceptionCountToleratedForWrites\": 5,"
+                    + "}");
+        }
+
         CosmosClientBuilder cosmosClientBuilder = new CosmosClientBuilder()
             .endpoint(cfg.getServiceEndpoint())
             .key(cfg.getMasterKey())
@@ -138,8 +148,11 @@ abstract class AsyncBenchmark<T> {
             cosmosClientBuilder = cosmosClientBuilder.gatewayMode(gatewayConnectionConfig);
         }
 
-        CosmosClient syncClient = cosmosClientBuilder.buildClient();
         cosmosClient = cosmosClientBuilder.buildAsyncClient();
+        CosmosClient syncClient = cosmosClientBuilder
+            .endpoint(configuration.getServiceEndpointForRunResultsUploadAccount())
+            .key(configuration.getMasterKeyForRunResultsUploadAccount())
+            .buildClient();
 
         try {
             cosmosAsyncDatabase = cosmosClient.getDatabase(this.configuration.getDatabaseId());
@@ -168,6 +181,17 @@ abstract class AsyncBenchmark<T> {
                 ).block();
 
                 cosmosAsyncContainer = cosmosAsyncDatabase.getContainer(this.configuration.getCollectionId());
+
+                // add some delay to allow container to be created across multiple regions
+                // container creation across regions is an async operation
+                // without the delay a container may not be available to process reads / writes
+
+                try {
+                    Thread.sleep(30_000);
+                } catch (Exception exception) {
+                    throw new RuntimeException(exception);
+                }
+
                 logger.info("Collection {} is created for this test", this.configuration.getCollectionId());
                 collectionCreated = true;
             } else {

sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/Configuration.java

@@ -47,6 +47,12 @@ public class Configuration {
     @Parameter(names = "-masterKey", description = "Master Key")
     private String masterKey;
 
+    @Parameter(names = "-serviceEndpointForResultsUploadAccount", description = "Service Endpoint for run results upload account")
+    private String serviceEndpointForRunResultsUploadAccount;
+
+    @Parameter(names = "-masterKeyForResultsUploadAccount", description = "Master Key for run results upload account")
+    private String masterKeyForRunResultsUploadAccount;
+
     @Parameter(names = "-databaseId", description = "Database ID")
     private String databaseId;
 
@@ -137,6 +143,9 @@ public class Configuration {
     @Parameter(names = "-isRegionScopedSessionContainerEnabled", description = "A flag to denote whether region scoped session container is enabled")
     private String isRegionScopedSessionContainerEnabled = String.valueOf(false);
 
+    @Parameter(names = "isPartitionLevelCircuitBreakerEnabled", description = "A flag to denote whether partition level circuit breaker is enabled.")
+    private String isPartitionLevelCircuitBreakerEnabled = String.valueOf(true);
+
     @Parameter(names = "-operation", description = "Type of Workload:\n"
         + "\tReadThroughput- run a READ workload that prints only throughput *\n"
         + "\tReadThroughputWithMultipleClients - run a READ workload that prints throughput and latency for multiple client read.*\n"
@@ -397,6 +406,14 @@ public String getMasterKey() {
         return masterKey;
     }
 
+    public String getServiceEndpointForRunResultsUploadAccount() {
+        return serviceEndpointForRunResultsUploadAccount;
+    }
+
+    public String getMasterKeyForRunResultsUploadAccount() {
+        return masterKeyForRunResultsUploadAccount;
+    }
+
     public String getApplicationName() {
         return applicationName;
     }
@@ -639,6 +656,10 @@ public boolean isRegionScopedSessionContainerEnabled() {
         return Boolean.parseBoolean(isRegionScopedSessionContainerEnabled);
     }
 
+    public boolean isPartitionLevelCircuitBreakerEnabled() {
+        return Boolean.parseBoolean(isPartitionLevelCircuitBreakerEnabled);
+    }
+
     public void tryGetValuesFromSystem() {
         serviceEndpoint = StringUtils.defaultString(Strings.emptyToNull(System.getenv().get("SERVICE_END_POINT")),
                                                     serviceEndpoint);

sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/SyncBenchmark.java

@@ -114,6 +114,16 @@ public T apply(T o, Throwable throwable) {
         configuration = cfg;
         logger = LoggerFactory.getLogger(this.getClass());
 
+        if (configuration.isPartitionLevelCircuitBreakerEnabled()) {
+            System.setProperty(
+                "COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG",
+                "{\"isPartitionLevelCircuitBreakerEnabled\": true, "
+                    + "\"circuitBreakerType\": \"CONSECUTIVE_EXCEPTION_COUNT_BASED\","
+                    + "\"consecutiveExceptionCountToleratedForReads\": 10,"
+                    + "\"consecutiveExceptionCountToleratedForWrites\": 5,"
+                    + "}");
+        }
+
         CosmosClientBuilder cosmosClientBuilder = new CosmosClientBuilder()
             .endpoint(cfg.getServiceEndpoint())
             .preferredRegions(cfg.getPreferredRegionsList())
@@ -147,6 +157,11 @@ public T apply(T o, Throwable throwable) {
         }
 
         cosmosClient = cosmosClientBuilder.buildClient();
+        CosmosClient syncClient = cosmosClientBuilder
+            .endpoint(configuration.getServiceEndpointForRunResultsUploadAccount())
+            .key(configuration.getMasterKeyForRunResultsUploadAccount())
+            .buildClient();
+
         try {
             cosmosDatabase = cosmosClient.getDatabase(this.configuration.getDatabaseId());
             cosmosDatabase.read();
@@ -171,6 +186,16 @@ public T apply(T o, Throwable throwable) {
                     ThroughputProperties.createManualThroughput(this.configuration.getThroughput()));
                 cosmosContainer = cosmosDatabase.getContainer(this.configuration.getCollectionId());
                 logger.info("Collection {} is created for this test", this.configuration.getCollectionId());
+
+                // add some delay to allow container to be created across multiple regions
+                // container creation across regions is an async operation
+                // without the delay a container may not be available to process reads / writes
+                try {
+                    Thread.sleep(30_000);
+                } catch (Exception exception) {
+                    throw new RuntimeException(exception);
+                }
+
                 collectionCreated = true;
             } else {
                 throw e;
@@ -236,7 +261,7 @@ public T apply(T o, Throwable throwable) {
             resultReporter = CosmosTotalResultReporter
                 .forRegistry(
                     metricsRegistry,
-                    cosmosClient.getDatabase(configuration.getResultUploadDatabase()).getContainer(configuration.getResultUploadContainer()),
+                    syncClient.getDatabase(configuration.getResultUploadDatabase()).getContainer(configuration.getResultUploadContainer()),
                     configuration)
                 .convertRatesTo(TimeUnit.SECONDS)
                 .convertDurationsTo(TimeUnit.MILLISECONDS).build();

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

@@ -11,6 +11,7 @@
 import com.azure.cosmos.CosmosRegionSwitchHint;
 import com.azure.cosmos.SessionRetryOptions;
 import com.azure.cosmos.SessionRetryOptionsBuilder;
+import com.azure.cosmos.implementation.circuitBreaker.PartitionLevelCircuitBreakerConfig;
 import com.azure.cosmos.implementation.directconnectivity.RntbdTransportClient;
 import com.azure.cosmos.implementation.guava25.collect.ImmutableList;
 import com.azure.cosmos.implementation.http.HttpClientConfig;
@@ -73,7 +74,8 @@ public Object[][] clientCfgProvider() {
                 aggressiveWarmUpDuration1,
                 proactiveConnectionRegionCount1,
                 cosmosContainerIdentities,
-                false // is region scoped session capturing enabled
+                false, // is region scoped session capturing enabled
+                false // is partition-level circuit breaking enabled
             },
             {
                 new CosmosContainerProactiveInitConfigBuilder(cosmosContainerIdentities)
@@ -83,7 +85,8 @@ public Object[][] clientCfgProvider() {
                 aggressiveWarmUpDuration2,
                 proactiveConnectionRegionCount2,
                 cosmosContainerIdentities,
-                true
+                true,  // is region scoped session capturing enabled
+                false // is partition-level circuit breaking enabled
             },
             {
                 new CosmosContainerProactiveInitConfigBuilder(cosmosContainerIdentities)
@@ -92,7 +95,8 @@ public Object[][] clientCfgProvider() {
                 null,
                 proactiveConnectionRegionCount3,
                 cosmosContainerIdentities,
-                false
+                false, // is region scoped session capturing enabled
+                true // is partition-level circuit breaking enabled
             }
         };
     }
@@ -243,7 +247,8 @@ public void full(
         Duration aggressiveWarmupDuration,
         int proactiveConnectionRegionCount,
         List<CosmosContainerIdentity> cosmosContainerIdentities,
-        boolean isRegionScopedSessionCapturingEnabled) throws Exception {
+        boolean isRegionScopedSessionCapturingEnabled,
+        boolean isPartitionLevelCircuitBreakerEnabled) throws Exception {
 
         DiagnosticsClientContext clientContext = Mockito.mock(DiagnosticsClientContext.class);
         System.setProperty("COSMOS.REPLICA_ADDRESS_VALIDATION_ENABLED", "false");
@@ -277,6 +282,19 @@ public void full(
             diagnosticsClientConfig.withRegionScopedSessionContainerOptions(regionScopedSessionContainer);
         }
 
+        if (isPartitionLevelCircuitBreakerEnabled) {
+            System.setProperty(
+                "COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG",
+                "{\"isPartitionLevelCircuitBreakerEnabled\": true, "
+                    + "\"circuitBreakerType\": \"CONSECUTIVE_EXCEPTION_COUNT_BASED\","
+                    + "\"consecutiveExceptionCountToleratedForReads\": 10,"
+                    + "\"consecutiveExceptionCountToleratedForWrites\": 5,"
+                    + "}");
+
+            PartitionLevelCircuitBreakerConfig partitionLevelCircuitBreakerConfig = Configs.getPartitionLevelCircuitBreakerConfig();
+            diagnosticsClientConfig.withPartitionLevelCircuitBreakerConfig(partitionLevelCircuitBreakerConfig);
+        }
+
         Mockito.doReturn(diagnosticsClientConfig).when(clientContext).getConfig();
 
         StringWriter jsonWriter = new StringWriter();
@@ -301,11 +319,18 @@ public void full(
             assertThat(objectNode.get("regionScopedSessionCfg")).isNull();
         }
 
+        if (isPartitionLevelCircuitBreakerEnabled) {
+            assertThat(objectNode.get("partitionLevelCircuitBreakerCfg").asText()).isEqualTo("(cb: true, type: CONSECUTIVE_EXCEPTION_COUNT_BASED, rexcntt: 10, wexcntt: 5)");
+        } else {
+            assertThat(objectNode.get("partitionLevelCircuitBreakerCfg")).isNull();
+        }
+
         String expectedProactiveInitConfigString = reconstructProactiveInitConfigString(cosmosContainerIdentities, aggressiveWarmupDuration, proactiveConnectionRegionCount);
 
         assertThat(objectNode.get("proactiveInitCfg").asText()).isEqualTo(expectedProactiveInitConfigString);
 
         System.clearProperty("COSMOS.REPLICA_ADDRESS_VALIDATION_ENABLED");
+        System.clearProperty("COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG");
     }
 
     @Test(groups = {"unit"}, dataProvider = "sessionRetryOptionsConfigProvider")