diff --git a/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncBenchmark.java b/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncBenchmark.java index 2d645d5ad548..3044dc80e61e 100644 --- a/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncBenchmark.java +++ b/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncBenchmark.java @@ -210,6 +210,7 @@ abstract class AsyncBenchmark implements Benchmark { }); CosmosBulkExecutionOptions bulkExecutionOptions = new CosmosBulkExecutionOptions(); + bulkExecutionOptions.setExcludedRegions(cfg.getExcludedRegionsList()); List> failedResponses = Collections.synchronizedList(new ArrayList<>()); cosmosAsyncContainer .executeBulkOperations(bulkOperationFlux, bulkExecutionOptions) diff --git a/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncQueryBenchmark.java b/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncQueryBenchmark.java index a7327f78255e..838351c19f6d 100644 --- a/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncQueryBenchmark.java +++ b/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncQueryBenchmark.java @@ -40,6 +40,7 @@ protected Mono> performWorkload(long i) { Flux> obs; Random r = new Random(); CosmosQueryRequestOptions options = new CosmosQueryRequestOptions(); + options.setExcludedRegions(workloadConfig.getExcludedRegionsList()); if (workloadConfig.getOperationType() == Operation.QueryCross) { diff --git a/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncQuerySinglePartitionMultiple.java b/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncQuerySinglePartitionMultiple.java index 55e381948967..eb41473a5a26 100644 --- a/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncQuerySinglePartitionMultiple.java +++ b/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncQuerySinglePartitionMultiple.java @@ -19,6 +19,7 @@ class AsyncQuerySinglePartitionMultiple extends AsyncBenchmark { protected Mono performWorkload(long i) { int index = (int) (i % docsToRead.size()); PojoizedJson doc = docsToRead.get(index); + CosmosItemRequestOptions options = new CosmosItemRequestOptions(); + options.setExcludedRegions(workloadConfig.getExcludedRegionsList()); return cosmosAsyncContainer.readItem(doc.getId(), - new PartitionKey(doc.getId()), PojoizedJson.class) + new PartitionKey(doc.getId()), options, PojoizedJson.class) .map(CosmosItemResponse::getItem); } } diff --git a/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncReadManyBenchmark.java b/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncReadManyBenchmark.java index ffd64f88a7d1..f9c00d91da41 100644 --- a/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncReadManyBenchmark.java +++ b/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncReadManyBenchmark.java @@ -4,6 +4,7 @@ package com.azure.cosmos.benchmark; import com.azure.cosmos.models.CosmosItemIdentity; +import com.azure.cosmos.models.CosmosReadManyRequestOptions; import com.azure.cosmos.models.FeedResponse; import com.azure.cosmos.models.PartitionKey; @@ -37,6 +38,8 @@ protected Mono> performWorkload(long i) { cosmosItemIdentities.add(new CosmosItemIdentity(partitionKey, doc.getId())); } - return cosmosAsyncContainer.readMany(cosmosItemIdentities, PojoizedJson.class); + CosmosReadManyRequestOptions options = new CosmosReadManyRequestOptions(); + options.setExcludedRegions(workloadConfig.getExcludedRegionsList()); + return cosmosAsyncContainer.readMany(cosmosItemIdentities, options, PojoizedJson.class); } } \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncWriteBenchmark.java b/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncWriteBenchmark.java index 53c01d097a2c..d388a4400c6f 100644 --- a/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncWriteBenchmark.java +++ b/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncWriteBenchmark.java @@ -3,6 +3,7 @@ package com.azure.cosmos.benchmark; +import com.azure.cosmos.models.CosmosItemRequestOptions; import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.PartitionKey; @@ -28,18 +29,21 @@ class AsyncWriteBenchmark extends AsyncBenchmark { protected Mono performWorkload(long i) { String id = uuid + i; Mono> result; + CosmosItemRequestOptions options = new CosmosItemRequestOptions(); + options.setExcludedRegions(workloadConfig.getExcludedRegionsList()); if (workloadConfig.isDisablePassingPartitionKeyAsOptionOnWrite()) { result = cosmosAsyncContainer.createItem(BenchmarkHelper.generateDocument(id, dataFieldValue, partitionKey, - workloadConfig.getDocumentDataFieldCount())); + workloadConfig.getDocumentDataFieldCount()), + options); } else { result = cosmosAsyncContainer.createItem(BenchmarkHelper.generateDocument(id, dataFieldValue, partitionKey, workloadConfig.getDocumentDataFieldCount()), new PartitionKey(id), - null); + options); } // Raw type cast is required because CosmosItemResponse uses wildcard generics // that cannot be expressed in the class type parameter without propagating diff --git a/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/SyncBenchmark.java b/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/SyncBenchmark.java index c7fea47c1dcb..5f48e0411ac9 100644 --- a/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/SyncBenchmark.java +++ b/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/SyncBenchmark.java @@ -186,6 +186,7 @@ private static ImplementationBridgeHelpers.CosmosClientHelper.CosmosClientAccess }); CosmosBulkExecutionOptions bulkExecutionOptions = new CosmosBulkExecutionOptions(); + bulkExecutionOptions.setExcludedRegions(workloadCfg.getExcludedRegionsList()); List> failedResponses = Collections.synchronizedList(new ArrayList<>()); asyncContainer .executeBulkOperations(bulkOperationFlux, bulkExecutionOptions) diff --git a/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/SyncReadBenchmark.java b/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/SyncReadBenchmark.java index 064b8a84c4fe..a60cdf06d2eb 100644 --- a/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/SyncReadBenchmark.java +++ b/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/SyncReadBenchmark.java @@ -21,7 +21,9 @@ protected CosmosItemResponse performWorkload(long i) throws Exception { PojoizedJson doc = docsToRead.get(index); String partitionKeyValue = doc.getId(); + CosmosItemRequestOptions options = new CosmosItemRequestOptions(); + options.setExcludedRegions(workloadConfig.getExcludedRegionsList()); return cosmosContainer.readItem(doc.getId(), new PartitionKey(partitionKeyValue), - new CosmosItemRequestOptions(), InternalObjectNode.class); + options, InternalObjectNode.class); } } diff --git a/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/SyncWriteBenchmark.java b/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/SyncWriteBenchmark.java index 8f0d085a1167..2066826745a0 100644 --- a/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/SyncWriteBenchmark.java +++ b/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/SyncWriteBenchmark.java @@ -3,6 +3,7 @@ package com.azure.cosmos.benchmark; +import com.azure.cosmos.models.CosmosItemRequestOptions; import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.PartitionKey; import org.apache.commons.lang3.RandomStringUtils; @@ -26,13 +27,15 @@ class SyncWriteBenchmark extends SyncBenchmark { @Override protected CosmosItemResponse performWorkload(long i) throws Exception { String id = uuid + i; - CosmosItemResponse response; + CosmosItemRequestOptions options = new CosmosItemRequestOptions(); + options.setExcludedRegions(workloadConfig.getExcludedRegionsList()); if (workloadConfig.isDisablePassingPartitionKeyAsOptionOnWrite()) { // require parsing partition key from the doc return cosmosContainer.createItem(BenchmarkHelper.generateDocument(id, dataFieldValue, partitionKey, - workloadConfig.getDocumentDataFieldCount())); + workloadConfig.getDocumentDataFieldCount()), + options); } // more optimized for write as partition key is already passed as config @@ -41,6 +44,6 @@ protected CosmosItemResponse performWorkload(long i) throws Exception { partitionKey, workloadConfig.getDocumentDataFieldCount()), new PartitionKey(id), - null); + options); } } diff --git a/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/TenantDefaultConfig.java b/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/TenantDefaultConfig.java index 11aa456e650a..d91d50949638 100644 --- a/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/TenantDefaultConfig.java +++ b/sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/TenantDefaultConfig.java @@ -140,6 +140,9 @@ class TenantDefaultConfig { @JsonProperty("preferredRegionsList") protected String preferredRegionsList; + @JsonProperty("excludedRegionsList") + protected String excludedRegionsList; + @JsonProperty("manageDatabase") protected Boolean manageDatabase; @@ -233,6 +236,13 @@ public List getPreferredRegionsList() { return regions; } + public List getExcludedRegionsList() { + if (excludedRegionsList == null || excludedRegionsList.isEmpty()) return null; + List regions = new ArrayList<>(); + for (String r : excludedRegionsList.split(",")) { regions.add(r.trim()); } + return regions; + } + public boolean shouldManageDatabase() { return manageDatabase != null ? manageDatabase : false; } public String getApplicationName() { return applicationName != null ? applicationName : ""; } public boolean isManagedIdentityRequired() { return isManagedIdentityRequired != null && isManagedIdentityRequired; } @@ -278,6 +288,7 @@ void applyTo(TenantWorkloadConfig tenant) { if (http2Enabled != null && tenant.http2Enabled == null) tenant.http2Enabled = http2Enabled; if (http2MaxConcurrentStreams != null && tenant.http2MaxConcurrentStreams == null) tenant.http2MaxConcurrentStreams = http2MaxConcurrentStreams; if (preferredRegionsList != null && tenant.preferredRegionsList == null) tenant.preferredRegionsList = preferredRegionsList; + if (excludedRegionsList != null && tenant.excludedRegionsList == null) tenant.excludedRegionsList = excludedRegionsList; if (manageDatabase != null && tenant.manageDatabase == null) tenant.manageDatabase = manageDatabase; if (isManagedIdentityRequired != null && tenant.isManagedIdentityRequired == null) tenant.isManagedIdentityRequired = isManagedIdentityRequired; if (applicationName != null && tenant.applicationName == null) tenant.applicationName = applicationName;