Skip to content

AddThroughputBucketSupport #46042

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 15 commits into
base: main
Choose a base branch
from
Open

AddThroughputBucketSupport #46042

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
152414b
add throughput bucket support
Jul 16, 2025
0cbf74d
reconstruct test structure for test
Jul 16, 2025
eb72c43
refactor
Jul 16, 2025
39d69bd
refactor
Jul 16, 2025
6e3025d
fix tests
Jul 17, 2025
708dce1
refactor
Jul 20, 2025
86bbe8e
merge from master and resolve conflicts
Jul 25, 2025
763f754
change
Jul 26, 2025
472109d
Merge branch 'main' into addThroughputBucketSupport
Jul 28, 2025
c713a38
change
Jul 28, 2025
c965815
add throughput control group config in diagnostics
Jul 29, 2025
4d895ee
fix compile
Jul 29, 2025
b84e992
resolve comments
Jul 29, 2025
1d9d4d8
Merge branch 'main' into addThroughputBucketSupport
Jul 29, 2025
e934d03
update javadoc
Jul 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 5 additions & 6 deletions sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ClosedClientTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
import com.azure.cosmos.implementation.OperationType;
import com.azure.cosmos.implementation.Utils;
import com.azure.cosmos.implementation.directconnectivity.rntbd.ClosedClientTransportException;
import com.azure.cosmos.implementation.throughputControl.TestItem;
import com.azure.cosmos.models.CosmosItemRequestOptions;
import com.azure.cosmos.models.CosmosPatchOperations;
import com.azure.cosmos.models.CosmosQueryRequestOptions;
Expand Down Expand Up @@ -95,7 +94,7 @@ public void failFastWhenUsingClosedClient(OperationType operationType) {

safeClose(asyncClient);

performDocumentOperation(cosmosAsyncContainer, TestItem.createNewItem(), operationType);
performDocumentOperation(cosmosAsyncContainer, TestObject.create(), operationType);
fail("Operation is expected to fail!");
} catch (Exception ex) {

Expand All @@ -121,7 +120,7 @@ public void failFastWhenUsingClosedClient(OperationType operationType) {

private void performDocumentOperation(
CosmosAsyncContainer asyncContainer,
TestItem testItem,
TestObject testItem,
OperationType operationType) {

if (operationType == OperationType.Query) {
Expand All @@ -131,15 +130,15 @@ private void performDocumentOperation(
SqlQuerySpec sqlQuerySpec = new SqlQuerySpec(query);

asyncContainer
.queryItems(sqlQuerySpec, queryRequestOptions, TestItem.class)
.queryItems(sqlQuerySpec, queryRequestOptions, TestObject.class)
.byPage()
.blockFirst();
}

if (operationType == OperationType.Read) {

asyncContainer
.readItem(testItem.getId(), new PartitionKey(testItem.getId()), TestItem.class)
.readItem(testItem.getId(), new PartitionKey(testItem.getId()), TestObject.class)
.block();
}

Expand Down Expand Up @@ -181,7 +180,7 @@ private void performDocumentOperation(
CosmosPatchOperations patchOperations = CosmosPatchOperations.create().add("/" + "newProperty", "newVal");

asyncContainer
.patchItem(testItem.getId(), new PartitionKey(testItem.getId()), patchOperations, TestItem.class)
.patchItem(testItem.getId(), new PartitionKey(testItem.getId()), patchOperations, TestObject.class)
.block();
}
}
Expand Down
39 changes: 19 additions & 20 deletions ...zure-cosmos-tests/src/test/java/com/azure/cosmos/EndToEndTimeOutWithAvailabilityTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
import com.azure.cosmos.implementation.OperationType;
import com.azure.cosmos.implementation.RxDocumentClientImpl;
import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils;
import com.azure.cosmos.implementation.throughputControl.TestItem;
import com.azure.cosmos.models.CosmosItemIdentity;
import com.azure.cosmos.models.CosmosItemRequestOptions;
import com.azure.cosmos.models.CosmosItemResponse;
Expand Down Expand Up @@ -107,7 +106,7 @@ public void testThresholdAvailabilityStrategy(
FaultInjectionConnectionType.DIRECT :
FaultInjectionConnectionType.GATEWAY;

TestItem createdItem = TestItem.createNewItem();
TestObject createdItem = TestObject.create();
CosmosItemRequestOptions options = new CosmosItemRequestOptions();
this.cosmosAsyncContainer.createItem(createdItem).block();

Expand Down Expand Up @@ -185,7 +184,7 @@ public void testThresholdAvailabilityStrategyForReadsDefaultEnablementWithPpaf(
CosmosAsyncDatabase asyncDatabase = cosmosAsyncClient.getDatabase(this.cosmosAsyncContainer.getDatabase().getId());
CosmosAsyncContainer asyncContainer = asyncDatabase.getContainer(this.cosmosAsyncContainer.getId());

TestItem createdItem = TestItem.createNewItem();
TestObject createdItem = TestObject.create();
CosmosItemRequestOptions options = new CosmosItemRequestOptions();
asyncContainer.createItem(createdItem).block();

Expand Down Expand Up @@ -306,7 +305,7 @@ public void afterClass() {
private CosmosDiagnostics performDocumentOperation(
CosmosAsyncContainer cosmosAsyncContainer,
OperationType operationType,
TestItem createdItem,
TestObject createdItem,
CosmosItemRequestOptions cosmosItemRequestOptions,
boolean ignoreE2E2LatencyCfgOnRequestOptions,
QueryFlavor queryFlavor) {
Expand All @@ -330,8 +329,8 @@ private CosmosDiagnostics performDocumentOperation(

logger.info("Running readAllItems...");

FeedResponse<TestItem> response = cosmosAsyncContainer
.readAllItems(queryRequestOptions, TestItem.class)
FeedResponse<TestObject> response = cosmosAsyncContainer
.readAllItems(queryRequestOptions, TestObject.class)
.byPage()
.blockFirst();

Expand All @@ -347,11 +346,11 @@ private CosmosDiagnostics performDocumentOperation(

logger.info("Running readMany...");

FeedResponse<TestItem> response = cosmosAsyncContainer
FeedResponse<TestObject> response = cosmosAsyncContainer
.readMany(
Arrays.asList(new CosmosItemIdentity(new PartitionKey(createdItem.getMypk()), createdItem.getId())),
readManyRequestOptions,
TestItem.class)
TestObject.class)
.block();

assertThat(response).isNotNull();
Expand All @@ -361,8 +360,8 @@ private CosmosDiagnostics performDocumentOperation(
logger.info("Running query ...");

String query = String.format("SELECT * from c where c.id = '%s'", createdItem.getId());
FeedResponse<TestItem> itemFeedResponse =
cosmosAsyncContainer.queryItems(query, queryRequestOptions, TestItem.class).byPage().blockFirst();
FeedResponse<TestObject> itemFeedResponse =
cosmosAsyncContainer.queryItems(query, queryRequestOptions, TestObject.class).byPage().blockFirst();

assertThat(itemFeedResponse).isNotNull();

Expand All @@ -378,11 +377,11 @@ private CosmosDiagnostics performDocumentOperation(

if (operationType == OperationType.Read) {

CosmosItemResponse<TestItem> response = cosmosAsyncContainer.readItem(
CosmosItemResponse<TestObject> response = cosmosAsyncContainer.readItem(
createdItem.getId(),
new PartitionKey(createdItem.getMypk()),
cosmosItemRequestOptions,
TestItem.class)
TestObject.class)
.block();

assertThat(response).isNotNull();
Expand All @@ -391,7 +390,7 @@ private CosmosDiagnostics performDocumentOperation(
}

if (operationType == OperationType.Replace) {
CosmosItemResponse<TestItem> response = cosmosAsyncContainer.replaceItem(
CosmosItemResponse<TestObject> response = cosmosAsyncContainer.replaceItem(
createdItem,
createdItem.getId(),
new PartitionKey(createdItem.getMypk()),
Expand All @@ -404,7 +403,7 @@ private CosmosDiagnostics performDocumentOperation(
}

if (operationType == OperationType.Delete) {
TestItem toBeDeletedItem = TestItem.createNewItem();
TestObject toBeDeletedItem = TestObject.create();
cosmosAsyncContainer.createItem(toBeDeletedItem, cosmosItemRequestOptions).block();
CosmosItemResponse<Object> response = cosmosAsyncContainer
.deleteItem(toBeDeletedItem, cosmosItemRequestOptions)
Expand All @@ -416,8 +415,8 @@ private CosmosDiagnostics performDocumentOperation(
}

if (operationType == OperationType.Create) {
CosmosItemResponse<TestItem> response = cosmosAsyncContainer
.createItem(TestItem.createNewItem(), cosmosItemRequestOptions)
CosmosItemResponse<TestObject> response = cosmosAsyncContainer
.createItem(TestObject.create(), cosmosItemRequestOptions)
.block();

assertThat(response).isNotNull();
Expand All @@ -426,8 +425,8 @@ private CosmosDiagnostics performDocumentOperation(
}

if (operationType == OperationType.Upsert) {
CosmosItemResponse<TestItem> response = cosmosAsyncContainer
.upsertItem(TestItem.createNewItem(), cosmosItemRequestOptions)
CosmosItemResponse<TestObject> response = cosmosAsyncContainer
.upsertItem(TestObject.create(), cosmosItemRequestOptions)
.block();

assertThat(response).isNotNull();
Expand All @@ -444,8 +443,8 @@ private CosmosDiagnostics performDocumentOperation(
CosmosPatchItemRequestOptions patchItemRequestOptions = new CosmosPatchItemRequestOptions();
patchItemRequestOptions.setNonIdempotentWriteRetryPolicy(true, true);
patchItemRequestOptions.setCosmosEndToEndOperationLatencyPolicyConfig(config);
CosmosItemResponse<TestItem> response = cosmosAsyncContainer
.patchItem(createdItem.getId(), new PartitionKey(createdItem.getMypk()), patchOperations, patchItemRequestOptions, TestItem.class)
CosmosItemResponse<TestObject> response = cosmosAsyncContainer
.patchItem(createdItem.getId(), new PartitionKey(createdItem.getMypk()), patchOperations, patchItemRequestOptions, TestObject.class)
.block();

assertThat(response).isNotNull();
Expand Down
33 changes: 16 additions & 17 deletions sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ExcludeRegionTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@
import com.azure.cosmos.implementation.OperationType;
import com.azure.cosmos.implementation.RxDocumentClientImpl;
import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils;
import com.azure.cosmos.implementation.throughputControl.TestItem;
import com.azure.cosmos.models.CosmosBatch;
import com.azure.cosmos.models.CosmosBatchRequestOptions;
import com.azure.cosmos.models.CosmosBatchResponse;
Expand Down Expand Up @@ -122,7 +121,7 @@ public void excludeRegionTest_SkipFirstPreferredRegion(OperationType operationTy
throw new SkipException("excludeRegionTest_SkipFirstPreferredRegion can only be tested for multi-master with multi-regions");
}

TestItem createdItem = TestItem.createNewItem();
TestObject createdItem = TestObject.create();
this.cosmosAsyncContainer.createItem(createdItem).block();

Thread.sleep(1000);
Expand Down Expand Up @@ -157,7 +156,7 @@ public void excludeRegionTest_readSessionNotAvailable(
throw new SkipException("excludeRegionTest_SkipFirstPreferredRegion can only be tested for multi-master with multi-regions");
}

TestItem createdItem = TestItem.createNewItem();
TestObject createdItem = TestObject.create();
this.cosmosAsyncContainer.createItem(createdItem).block();

FaultInjectionRule serverErrorRule = new FaultInjectionRuleBuilder("excludeRegionTest-" + operationType)
Expand Down Expand Up @@ -229,7 +228,7 @@ private List<String> getPreferredRegionList(CosmosAsyncClient client) {
private CosmosDiagnosticsContext performDocumentOperation(
CosmosAsyncContainer cosmosAsyncContainer,
OperationType operationType,
TestItem createdItem,
TestObject createdItem,
List<String> excludeRegions,
CosmosEndToEndOperationLatencyPolicyConfig cosmosEndToEndOperationLatencyPolicyConfig) {

Expand All @@ -240,8 +239,8 @@ private CosmosDiagnosticsContext performDocumentOperation(

String query = String.format("SELECT * from c where c.id = '%s'", createdItem.getId());
queryRequestOptions.setExcludedRegions(excludeRegions);
FeedResponse<TestItem> itemFeedResponse =
cosmosAsyncContainer.queryItems(query, queryRequestOptions, TestItem.class).byPage().blockFirst();
FeedResponse<TestObject> itemFeedResponse =
cosmosAsyncContainer.queryItems(query, queryRequestOptions, TestObject.class).byPage().blockFirst();

assertThat(itemFeedResponse).isNotNull();
CosmosDiagnostics cosmosDiagnostics = itemFeedResponse.getCosmosDiagnostics();
Expand Down Expand Up @@ -271,11 +270,11 @@ private CosmosDiagnosticsContext performDocumentOperation(
throw new RuntimeException(e);
}

CosmosItemResponse<TestItem> itemResponse = cosmosAsyncContainer.readItem(
CosmosItemResponse<TestObject> itemResponse = cosmosAsyncContainer.readItem(
createdItem.getId(),
new PartitionKey(createdItem.getMypk()),
cosmosItemRequestOptions,
TestItem.class).block();
TestObject.class).block();

assertThat(itemResponse).isNotNull();
assertThat(itemResponse.getDiagnostics()).isNotNull();
Expand All @@ -295,7 +294,7 @@ private CosmosDiagnosticsContext performDocumentOperation(
throw new RuntimeException(e);
}

CosmosItemResponse<TestItem> itemResponse = cosmosAsyncContainer.replaceItem(
CosmosItemResponse<TestObject> itemResponse = cosmosAsyncContainer.replaceItem(
createdItem,
createdItem.getId(),
new PartitionKey(createdItem.getMypk()),
Expand All @@ -313,7 +312,7 @@ private CosmosDiagnosticsContext performDocumentOperation(

if (operationType == OperationType.Delete) {

TestItem itemToBeDeleted = TestItem.createNewItem();
TestObject itemToBeDeleted = TestObject.create();

cosmosAsyncContainer.createItem(itemToBeDeleted, cosmosItemRequestOptions).block();

Expand All @@ -337,8 +336,8 @@ private CosmosDiagnosticsContext performDocumentOperation(
}

if (operationType == OperationType.Create) {
CosmosItemResponse<TestItem> itemResponse = cosmosAsyncContainer
.createItem(TestItem.createNewItem(), cosmosItemRequestOptions).block();
CosmosItemResponse<TestObject> itemResponse = cosmosAsyncContainer
.createItem(TestObject.create(), cosmosItemRequestOptions).block();

assertThat(itemResponse).isNotNull();
assertThat(itemResponse.getDiagnostics()).isNotNull();
Expand All @@ -351,8 +350,8 @@ private CosmosDiagnosticsContext performDocumentOperation(
}

if (operationType == OperationType.Upsert) {
CosmosItemResponse<TestItem> itemResponse
= cosmosAsyncContainer.upsertItem(TestItem.createNewItem(), cosmosItemRequestOptions).block();
CosmosItemResponse<TestObject> itemResponse
= cosmosAsyncContainer.upsertItem(TestObject.create(), cosmosItemRequestOptions).block();

assertThat(itemResponse).isNotNull();
assertThat(itemResponse.getDiagnostics()).isNotNull();
Expand All @@ -375,8 +374,8 @@ private CosmosDiagnosticsContext performDocumentOperation(
patchItemRequestOptions.setCosmosEndToEndOperationLatencyPolicyConfig(cosmosEndToEndOperationLatencyPolicyConfig);
patchItemRequestOptions.setExcludedRegions(excludeRegions);

CosmosItemResponse<TestItem> itemResponse = cosmosAsyncContainer
.patchItem(createdItem.getId(), new PartitionKey(createdItem.getMypk()), patchOperations, patchItemRequestOptions, TestItem.class)
CosmosItemResponse<TestObject> itemResponse = cosmosAsyncContainer
.patchItem(createdItem.getId(), new PartitionKey(createdItem.getMypk()), patchOperations, patchItemRequestOptions, TestObject.class)
.block();

assertThat(itemResponse).isNotNull();
Expand All @@ -396,7 +395,7 @@ private CosmosDiagnosticsContext performDocumentOperation(

cosmosBatchRequestOptions.setExcludedRegions(excludeRegions);

TestItem testItem = TestItem.createNewItem();
TestObject testItem = TestObject.create();
PartitionKey partitionKey = new PartitionKey(testItem.getMypk());

CosmosBatch cosmosBatch = CosmosBatch.createCosmosBatch(partitionKey);
Expand Down
Loading
Loading