Skip to content

Commit c713a38

Browse files
author
annie-mac
committed
change
1 parent 472109d commit c713a38

File tree

15 files changed

+159
-31
lines changed

15 files changed

+159
-31
lines changed

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -920,7 +920,7 @@ public void faultInjectionServerErrorRuleTests_LeaseNotFound(OperationType opera
920920
shouldRetryCrossRegion = true;
921921
}
922922

923-
TestItem createdItem = TestItem.createNewItem();
923+
TestObject createdItem = TestObject.create();
924924

925925
String ruleId = "serverErrorRule-" + FaultInjectionServerErrorType.LEASE_NOT_FOUND + "-" + UUID.randomUUID();
926926
FaultInjectionRule serverErrorRule =
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
package com.azure.cosmos.implementation.throughputControl;
5+
6+
import com.azure.cosmos.ThroughputControlGroupConfig;
7+
import com.azure.cosmos.ThroughputControlGroupConfigBuilder;
8+
import org.testng.annotations.Test;
9+
10+
import java.util.UUID;
11+
12+
import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
13+
import static org.assertj.core.api.AssertionsForClassTypes.assertThatThrownBy;
14+
15+
public class ThroughputControlGroupConfigTests {
16+
17+
@Test(groups = "unit")
18+
public void throughputControlGroup_throughputBucket_priorityLevel() {
19+
// validate throughputBucket >= 0
20+
assertThatThrownBy(
21+
() -> new ThroughputControlGroupConfigBuilder()
22+
.groupName("throughputBucket-" + UUID.randomUUID())
23+
.throughputBucket(-1)
24+
.build())
25+
.isInstanceOf(IllegalArgumentException.class)
26+
.hasMessage("Throughput bucket should be no smaller than 0");
27+
28+
// valid config
29+
ThroughputControlGroupConfig throughputControlGroupConfig =
30+
new ThroughputControlGroupConfigBuilder()
31+
.groupName("throughputBucket-" + UUID.randomUUID())
32+
.throughputBucket(1)
33+
.build();
34+
35+
assertThat(throughputControlGroupConfig.getThroughputBucket()).isEqualTo(1);
36+
37+
// neither priorityLevel nor throughput bucket configured
38+
assertThatThrownBy(
39+
() -> new ThroughputControlGroupConfigBuilder()
40+
.groupName("throughputBucket-" + UUID.randomUUID())
41+
.build())
42+
.isInstanceOf(IllegalArgumentException.class)
43+
.hasMessage("All targetThroughput, targetThroughputThreshold, priorityLevel and throughput bucket cannot be null or empty.");
44+
}
45+
}

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

Lines changed: 63 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -672,9 +672,9 @@ public void enableSameGroupMultipleTimes() {
672672
}
673673

674674
@Test(groups = {"emulator"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT)
675-
public void serverThroughputControl(OperationType operationType) {
675+
public void serverThroughputControl_throughputBucket(OperationType operationType) {
676676
// TODO: currently there is no easy way to do e2e testing, so the testing here is to just verify that
677-
// server throughput can be enabled on the container
677+
// server throughput can be enabled on the container with throughput control
678678
this.ensureContainer();
679679

680680
ThroughputControlGroupConfig serverThroughputControlGroup =
@@ -701,6 +701,67 @@ public void serverThroughputControl(OperationType operationType) {
701701
serverThroughputControlGroup.getGroupName());
702702
}
703703

704+
@Test(groups = {"emulator"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT)
705+
public void serverThroughputControl_priorityLevel(OperationType operationType) {
706+
// TODO: currently there is no easy way to do e2e testing, so the testing here is to just verify that
707+
// server throughput can be enabled on the container with priority level
708+
this.ensureContainer();
709+
710+
ThroughputControlGroupConfig serverThroughputControlGroup =
711+
new ThroughputControlGroupConfigBuilder()
712+
.groupName("serverThroughputControl")
713+
.priorityLevel(PriorityLevel.LOW)
714+
.build();
715+
container.enableServerThroughputControlGroup(serverThroughputControlGroup);
716+
717+
CosmosItemRequestOptions requestOptions = new CosmosItemRequestOptions();
718+
requestOptions.setContentResponseOnWriteEnabled(true);
719+
requestOptions.setThroughputControlGroupName(serverThroughputControlGroup.getGroupName());
720+
721+
CosmosItemResponse<TestItem> createItemResponse = container.createItem(getDocumentDefinition(), requestOptions).block();
722+
TestItem createdItem = createItemResponse.getItem();
723+
this.validateRequestNotThrottled(
724+
createItemResponse.getDiagnostics().toString(),
725+
BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode());
726+
727+
performDocumentOperation(
728+
this.container,
729+
operationType,
730+
createdItem,
731+
serverThroughputControlGroup.getGroupName());
732+
}
733+
734+
@Test(groups = {"emulator"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT)
735+
public void serverThroughputControl_priorityLevel_throughputBucket(OperationType operationType) {
736+
// TODO: currently there is no easy way to do e2e testing, so the testing here is to just verify that
737+
// server throughput can be enabled on the container with priority level and throughput bucket
738+
this.ensureContainer();
739+
740+
ThroughputControlGroupConfig serverThroughputControlGroup =
741+
new ThroughputControlGroupConfigBuilder()
742+
.groupName("serverThroughputControl")
743+
.priorityLevel(PriorityLevel.LOW)
744+
.throughputBucket(3)
745+
.build();
746+
container.enableServerThroughputControlGroup(serverThroughputControlGroup);
747+
748+
CosmosItemRequestOptions requestOptions = new CosmosItemRequestOptions();
749+
requestOptions.setContentResponseOnWriteEnabled(true);
750+
requestOptions.setThroughputControlGroupName(serverThroughputControlGroup.getGroupName());
751+
752+
CosmosItemResponse<TestItem> createItemResponse = container.createItem(getDocumentDefinition(), requestOptions).block();
753+
TestItem createdItem = createItemResponse.getItem();
754+
this.validateRequestNotThrottled(
755+
createItemResponse.getDiagnostics().toString(),
756+
BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode());
757+
758+
performDocumentOperation(
759+
this.container,
760+
operationType,
761+
createdItem,
762+
serverThroughputControlGroup.getGroupName());
763+
}
764+
704765
@Test(groups = {"emulator"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT)
705766
public void throughputControl_LocalAndServer_requestOptions(OperationType operationType) {
706767
this.ensureContainer();

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ClientRetryPolicyE2ETests.java

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -687,7 +687,6 @@ private Mono<CosmosDiagnostics> performDocumentOperation(
687687
Function<TestObject, PartitionKey> extractPartitionKeyFunc,
688688
boolean isReadMany) {
689689

690-
691690
try {
692691
if (operationType == OperationType.Query && isReadMany) {
693692
CosmosQueryRequestOptions queryRequestOptions = new CosmosQueryRequestOptions();

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@
7171
import com.azure.cosmos.models.SqlQuerySpec;
7272
import com.azure.cosmos.models.ThroughputProperties;
7373
import com.azure.cosmos.models.ThroughputResponse;
74+
import com.azure.cosmos.util.Beta;
7475
import com.azure.cosmos.util.CosmosPagedFlux;
7576
import com.azure.cosmos.util.UtilBridgeInternal;
7677
import org.slf4j.Logger;
@@ -2825,12 +2826,14 @@ void enableGlobalThroughputControlGroup(
28252826

28262827
/***
28272828
* Enable the server throughput bucket control group.
2828-
* <p>
2829-
* For more information about throughput bucket please visit
2830-
* <a href="https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/throughput-buckets?tabs=dotnet">Throughput buckets in Azure Cosmos DB</a>
2829+
* </br>
2830+
*
2831+
* <!-- src_embed com.azure.cosmos.throughputControl.serverControl -->
2832+
* <!-- end com.azure.cosmos.throughputControl.serverControl -->
28312833
*
28322834
* @param groupConfig the throughput control group config, see {@link ThroughputControlGroupConfig}.
28332835
*/
2836+
@Beta(value = Beta.SinceVersion.V4_74_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
28342837
public void enableServerThroughputControlGroup(ThroughputControlGroupConfig groupConfig) {
28352838
if (groupConfig.getPriorityLevel() == null && groupConfig.getThroughputBucket() == null) {
28362839
throw new IllegalArgumentException("Config 'priorityLevel' and 'throughputBucket' can not be null for both.");

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import com.azure.cosmos.models.SqlQuerySpec;
2929
import com.azure.cosmos.models.ThroughputProperties;
3030
import com.azure.cosmos.models.ThroughputResponse;
31+
import com.azure.cosmos.util.Beta;
3132
import com.azure.cosmos.util.CosmosPagedFlux;
3233
import com.azure.cosmos.util.CosmosPagedIterable;
3334
import org.slf4j.Logger;
@@ -948,7 +949,7 @@ public CosmosScripts getScripts() {
948949
}
949950

950951
// TODO: should make partitionkey public in CosmosAsyncItem and fix the below call
951-
952+
952953
private <T> CosmosPagedIterable<T> getCosmosPagedIterable(CosmosPagedFlux<T> cosmosPagedFlux) {
953954
return new CosmosPagedIterable<>(cosmosPagedFlux);
954955
}
@@ -1031,13 +1032,15 @@ public void enableGlobalThroughputControlGroup(ThroughputControlGroupConfig grou
10311032
}
10321033

10331034
/***
1034-
* Enable the server throughput bucket control group.
1035-
* <p>
1036-
* For more information about throughput bucket please visit
1037-
* <a href="https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/throughput-buckets?tabs=dotnet">Throughput buckets in Azure Cosmos DB</a>
1035+
* Enable the server throughput control group.
1036+
* </br>
1037+
*
1038+
* <!-- src_embed com.azure.cosmos.throughputControl.serverControl -->
1039+
* <!-- end com.azure.cosmos.throughputControl.serverControl -->
10381040
*
10391041
* @param groupConfig the throughput control group config, see {@link ThroughputControlGroupConfig}.
10401042
*/
1043+
@Beta(value = Beta.SinceVersion.V4_74_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
10411044
public void enableServerThroughputControlGroup(ThroughputControlGroupConfig groupConfig) {
10421045
this.asyncContainer.enableServerThroughputControlGroup(groupConfig);
10431046
}

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfig.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
package com.azure.cosmos;
55

66
import com.azure.cosmos.models.PriorityLevel;
7+
import com.azure.cosmos.util.Beta;
78

89
/**
910
* Throughput control group configuration.
@@ -88,6 +89,7 @@ public Double getTargetThroughputThreshold() {
8889
*
8990
* @return the throughput bucket of the throughput control group.
9091
*/
92+
@Beta(value = Beta.SinceVersion.V4_74_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
9193
public Integer getThroughputBucket() {
9294
return this.throughputBucket;
9395
}

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfigBuilder.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -167,8 +167,9 @@ public ThroughputControlGroupConfigBuilder defaultControlGroup(boolean aDefault)
167167
* @param throughputBucket the throughput bucket id.
168168
* @return The {@link ThroughputControlGroupConfigBuilder}.
169169
*/
170+
@Beta(value = Beta.SinceVersion.V4_74_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
170171
public ThroughputControlGroupConfigBuilder throughputBucket(int throughputBucket) {
171-
checkArgument(throughputBucket > 0, "Throughput Bucket should be greater than 0");
172+
checkArgument(throughputBucket >= 0, "Throughput bucket should be no smaller than 0");
172173
this.throughputBucket = throughputBucket;
173174
return this;
174175
}

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6363,22 +6363,22 @@ public void close() {
63636363
}
63646364
}
63656365
@Override
6366-
public synchronized void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono<Integer> throughputQueryMono) {
6366+
public void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono<Integer> throughputQueryMono) {
63676367
checkNotNull(group, "Throughput control group can not be null");
63686368

63696369
this.enableThroughputControlStore();
63706370
this.throughputControlStore.enableSDKThroughputControlGroup(group, throughputQueryMono);
63716371
}
63726372

63736373
@Override
6374-
public synchronized void enableServerThroughputControlGroup(ServerThroughputControlGroup group) {
6374+
public void enableServerThroughputControlGroup(ServerThroughputControlGroup group) {
63756375
checkNotNull(group, "Argument 'group' can not be null");
63766376

63776377
this.enableThroughputControlStore();
63786378
this.throughputControlStore.enableServerThroughputControlGroup(group);
63796379
}
63806380

6381-
private void enableThroughputControlStore() {
6381+
private synchronized void enableThroughputControlStore() {
63826382
if (this.throughputControlEnabled.compareAndSet(false, true)) {
63836383
this.throughputControlStore =
63846384
new ThroughputControlStore(

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlStore.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ public ThroughputControlStore(
3838
this.serverThroughputControlStore = new ServerThroughputControlStore();
3939
}
4040

41-
public void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono<Integer> throughputQueryMono) {
41+
public synchronized void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono<Integer> throughputQueryMono) {
4242
checkNotNull(group, "Throughput control group cannot be null");
4343

4444
if (group.isDefault()) {
@@ -52,7 +52,7 @@ public void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal gr
5252
this.sdkThroughputControlStore.enableThroughputControlGroup(group, throughputQueryMono);
5353
}
5454

55-
public void enableServerThroughputControlGroup(ServerThroughputControlGroup group) {
55+
public synchronized void enableServerThroughputControlGroup(ServerThroughputControlGroup group) {
5656
checkNotNull(group, "Throughput control group cannot be null");
5757

5858
if (group.isDefault()) {

0 commit comments

Comments
 (0)