From 152414b92f93b231e011bc650371828a7988b4c7 Mon Sep 17 00:00:00 2001 From: annie-mac Date: Wed, 16 Jul 2025 10:04:18 -0700 Subject: [PATCH 01/12] add throughput bucket support --- .../directconnectivity/ReflectionUtils.java | 8 +- ...ThroughputControlGroupPropertiesTests.java | 11 +- .../LinkedCancellationTokenSourceTests.java | 2 + .../ThroughputControlTests.java | 2 +- .../ThroughputRequestThrottlerTests.java | 2 + ...lobalThroughputRequestControllerTests.java | 4 +- .../LocalThroughputControllerTests.java | 4 +- ...angesThroughputRequestControllerTests.java | 5 +- .../com/azure/cosmos/CosmosAsyncClient.java | 17 +- .../azure/cosmos/CosmosAsyncContainer.java | 32 +- .../com/azure/cosmos/CosmosContainer.java | 2 +- .../cosmos/ThroughputControlGroupConfig.java | 16 + .../ThroughputControlGroupConfigBuilder.java | 26 +- .../implementation/AsyncDocumentClient.java | 12 +- .../cosmos/implementation/HttpConstants.java | 3 + .../implementation/RxDocumentClientImpl.java | 21 +- .../RxDocumentServiceRequest.java | 6 + .../rntbd/RntbdConstants.java | 3 +- .../rntbd/RntbdRequestHeaders.java | 13 + .../EmptyThroughputContainerController.java | 2 +- .../IThroughputContainerController.java | 12 + .../ThroughputControlGroupFactory.java | 20 +- .../ThroughputControlStore.java | 281 +++------------ .../IThroughputContainerController.java | 13 - ...rSDKThroughputControlGroupProperties.java} | 50 ++- .../{ => sdk}/LinkedCancellationToken.java | 2 +- .../LinkedCancellationTokenSource.java | 2 +- .../sdk/SDKThroughputControlStore.java | 326 ++++++++++++++++++ .../ThroughputControlTrackingUnit.java | 2 +- .../{ => sdk}/ThroughputRequestThrottler.java | 3 +- .../config/GlobalThroughputControlGroup.java | 13 +- .../config/LocalThroughputControlGroup.java | 13 +- .../SDKThroughputControlGroupInternal.java} | 19 +- .../controller/IThroughputController.java | 2 +- .../SDKThroughputContainerController.java} | 37 +- .../ThroughputProvisioningScope.java | 2 +- .../group/ThroughputGroupControllerBase.java | 22 +- .../ThroughputGroupControllerFactory.java | 17 +- .../GlobalThroughputControlClientItem.java | 2 +- .../GlobalThroughputControlConfigItem.java | 2 +- ...lobalThroughputControlGroupController.java | 8 +- .../global/GlobalThroughputControlItem.java | 2 +- .../ThroughputControlContainerManager.java | 5 +- .../group/global/ThroughputUsageSnapshot.java | 2 +- ...LocalThroughputControlGroupController.java | 8 +- .../GlobalThroughputRequestController.java | 4 +- .../request/IThroughputRequestController.java | 4 +- .../PkRangesThroughputRequestController.java | 4 +- ...oughputControlInitializationException.java | 2 +- ...erverThroughputControlGroupProperties.java | 68 ++++ .../server/ServerThroughputControlStore.java | 109 ++++++ .../ServerThroughputControlGroupInternal.java | 66 ++++ .../config/ThroughputBucketControlGroup.java | 45 +++ .../ServerThroughputContainerController.java | 108 ++++++ .../ServerThroughputGroupControllerBase.java | 9 + .../ThroughputBucketGroupController.java | 44 +++ .../src/main/java/module-info.java | 4 +- 57 files changed, 1131 insertions(+), 392 deletions(-) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{controller/container => }/EmptyThroughputContainerController.java (91%) create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/IThroughputContainerController.java rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{config => }/ThroughputControlGroupFactory.java (70%) delete mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/IThroughputContainerController.java rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{ContainerThroughputControlGroupProperties.java => sdk/ContainerSDKThroughputControlGroupProperties.java} (71%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/LinkedCancellationToken.java (96%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/LinkedCancellationTokenSource.java (96%) create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/SDKThroughputControlStore.java rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/ThroughputControlTrackingUnit.java (98%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/ThroughputRequestThrottler.java (99%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/config/GlobalThroughputControlGroup.java (89%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/config/LocalThroughputControlGroup.java (56%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{config/ThroughputControlGroupInternal.java => sdk/config/SDKThroughputControlGroupInternal.java} (92%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/controller/IThroughputController.java (94%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{controller/container/ThroughputContainerController.java => sdk/controller/container/SDKThroughputContainerController.java} (92%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/controller/container/ThroughputProvisioningScope.java (68%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/controller/group/ThroughputGroupControllerBase.java (91%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/controller/group/ThroughputGroupControllerFactory.java (64%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/controller/group/global/GlobalThroughputControlClientItem.java (94%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/controller/group/global/GlobalThroughputControlConfigItem.java (97%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/controller/group/global/GlobalThroughputControlGroupController.java (94%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/controller/group/global/GlobalThroughputControlItem.java (93%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/controller/group/global/ThroughputControlContainerManager.java (98%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/controller/group/global/ThroughputUsageSnapshot.java (92%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/controller/group/local/LocalThroughputControlGroupController.java (79%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/controller/request/GlobalThroughputRequestController.java (89%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/controller/request/IThroughputRequestController.java (80%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/controller/request/PkRangesThroughputRequestController.java (96%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/exceptions/ThroughputControlInitializationException.java (93%) create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ContainerServerThroughputControlGroupProperties.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ServerThroughputControlStore.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ServerThroughputControlGroupInternal.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ThroughputBucketControlGroup.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputContainerController.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupControllerBase.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ThroughputBucketGroupController.java diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/ReflectionUtils.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/ReflectionUtils.java index 2fe1cef5a869..62831a86e9fb 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/ReflectionUtils.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/ReflectionUtils.java @@ -40,10 +40,10 @@ import com.azure.cosmos.implementation.http.HttpRequest; import com.azure.cosmos.implementation.routing.CollectionRoutingMap; import com.azure.cosmos.implementation.routing.LocationCache; -import com.azure.cosmos.implementation.throughputControl.ThroughputControlTrackingUnit; -import com.azure.cosmos.implementation.throughputControl.ThroughputRequestThrottler; -import com.azure.cosmos.implementation.throughputControl.controller.request.GlobalThroughputRequestController; -import com.azure.cosmos.implementation.throughputControl.controller.request.PkRangesThroughputRequestController; +import com.azure.cosmos.implementation.throughputControl.sdk.ThroughputControlTrackingUnit; +import com.azure.cosmos.implementation.throughputControl.sdk.ThroughputRequestThrottler; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.request.GlobalThroughputRequestController; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.request.PkRangesThroughputRequestController; import com.azure.cosmos.models.CosmosClientTelemetryConfig; import io.netty.handler.ssl.SslContext; import org.apache.commons.lang3.reflect.FieldUtils; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ContainerThroughputControlGroupPropertiesTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ContainerThroughputControlGroupPropertiesTests.java index 9c8965c0610d..5fde672ed95d 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ContainerThroughputControlGroupPropertiesTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ContainerThroughputControlGroupPropertiesTests.java @@ -9,7 +9,8 @@ import com.azure.cosmos.implementation.RxDocumentServiceRequest; import com.azure.cosmos.implementation.TestConfigurations; import com.azure.cosmos.implementation.apachecommons.lang.tuple.Pair; -import com.azure.cosmos.implementation.throughputControl.config.LocalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.sdk.ContainerSDKThroughputControlGroupProperties; +import com.azure.cosmos.implementation.throughputControl.sdk.config.LocalThroughputControlGroup; import com.azure.cosmos.models.PriorityLevel; import org.mockito.Mockito; import org.testng.annotations.Test; @@ -30,8 +31,8 @@ public void enableThroughputControlGroup() { .key(TestConfigurations.MASTER_KEY) .buildAsyncClient(); - ContainerThroughputControlGroupProperties throughputControlContainerProperties = - new ContainerThroughputControlGroupProperties("/testDB/testContainer"); + ContainerSDKThroughputControlGroupProperties throughputControlContainerProperties = + new ContainerSDKThroughputControlGroupProperties("/testDB/testContainer"); CosmosAsyncContainer container = testClient.getDatabase("fakeDatabase").getContainer("fakeContainer"); @@ -160,8 +161,8 @@ public void enableThroughputControlGroupWithoutDefault() { .key(TestConfigurations.MASTER_KEY) .buildAsyncClient(); - ContainerThroughputControlGroupProperties throughputControlContainerProperties = - new ContainerThroughputControlGroupProperties("/testDB/testContainer"); + ContainerSDKThroughputControlGroupProperties throughputControlContainerProperties = + new ContainerSDKThroughputControlGroupProperties("/testDB/testContainer"); // Test: Without default group and request not having the group name, allowRequestToContinueOnInitError // should not throw NPE diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationTokenSourceTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationTokenSourceTests.java index e505391a7c18..b0eb8eae81f8 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationTokenSourceTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationTokenSourceTests.java @@ -3,6 +3,8 @@ package com.azure.cosmos.implementation.throughputControl; +import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationToken; +import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationTokenSource; import org.testng.annotations.Test; import static org.assertj.core.api.Assertions.assertThat; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java index be0d563e54fc..e0e568bd2568 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java @@ -21,7 +21,7 @@ import com.azure.cosmos.implementation.TestConfigurations; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.implementation.apachecommons.lang.tuple.Pair; -import com.azure.cosmos.implementation.throughputControl.controller.group.global.GlobalThroughputControlClientItem; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.group.global.GlobalThroughputControlClientItem; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; import com.azure.cosmos.models.CosmosContainerProperties; import com.azure.cosmos.models.CosmosContainerRequestOptions; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputRequestThrottlerTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputRequestThrottlerTests.java index 299801ef74fc..3b026b3bd773 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputRequestThrottlerTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputRequestThrottlerTests.java @@ -13,6 +13,8 @@ import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils; import com.azure.cosmos.implementation.directconnectivity.StoreResponse; +import com.azure.cosmos.implementation.throughputControl.sdk.ThroughputControlTrackingUnit; +import com.azure.cosmos.implementation.throughputControl.sdk.ThroughputRequestThrottler; import org.mockito.Mockito; import org.testng.annotations.Test; import reactor.test.StepVerifier; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/GlobalThroughputRequestControllerTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/GlobalThroughputRequestControllerTests.java index 527ab40f6211..b3214e88626f 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/GlobalThroughputRequestControllerTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/GlobalThroughputRequestControllerTests.java @@ -8,8 +8,8 @@ import com.azure.cosmos.implementation.RxDocumentServiceRequest; import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils; import com.azure.cosmos.implementation.directconnectivity.StoreResponse; -import com.azure.cosmos.implementation.throughputControl.ThroughputRequestThrottler; -import com.azure.cosmos.implementation.throughputControl.controller.request.GlobalThroughputRequestController; +import com.azure.cosmos.implementation.throughputControl.sdk.ThroughputRequestThrottler; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.request.GlobalThroughputRequestController; import org.mockito.Mockito; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/LocalThroughputControllerTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/LocalThroughputControllerTests.java index 79680b6891b1..3bb45052405d 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/LocalThroughputControllerTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/LocalThroughputControllerTests.java @@ -11,8 +11,8 @@ import com.azure.cosmos.implementation.directconnectivity.StoreResponse; import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper; import com.azure.cosmos.implementation.routing.Range; -import com.azure.cosmos.implementation.throughputControl.config.LocalThroughputControlGroup; -import com.azure.cosmos.implementation.throughputControl.controller.group.local.LocalThroughputControlGroupController; +import com.azure.cosmos.implementation.throughputControl.sdk.config.LocalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.group.local.LocalThroughputControlGroupController; import com.azure.cosmos.models.PriorityLevel; import org.mockito.Mockito; import org.testng.annotations.BeforeClass; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/PkRangesThroughputRequestControllerTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/PkRangesThroughputRequestControllerTests.java index ba454eaf7372..8f392e8f1230 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/PkRangesThroughputRequestControllerTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/PkRangesThroughputRequestControllerTests.java @@ -13,8 +13,8 @@ import com.azure.cosmos.implementation.directconnectivity.StoreResponse; import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper; import com.azure.cosmos.implementation.routing.Range; -import com.azure.cosmos.implementation.throughputControl.ThroughputRequestThrottler; -import com.azure.cosmos.implementation.throughputControl.controller.request.PkRangesThroughputRequestController; +import com.azure.cosmos.implementation.throughputControl.sdk.ThroughputRequestThrottler; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.request.PkRangesThroughputRequestController; import org.mockito.Mockito; import org.testng.annotations.BeforeClass; import org.testng.annotations.BeforeMethod; @@ -31,7 +31,6 @@ import java.util.concurrent.ConcurrentHashMap; import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.InstanceOfAssertFactories.OPTIONAL; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index 6736cd186e18..106038f82fc3 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java @@ -29,7 +29,8 @@ import com.azure.cosmos.implementation.clienttelemetry.TagName; import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdMetrics; import com.azure.cosmos.implementation.faultinjection.IFaultInjectorProvider; -import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroupInternal; import com.azure.cosmos.models.CosmosAuthorizationTokenResolver; import com.azure.cosmos.models.CosmosClientTelemetryConfig; import com.azure.cosmos.models.CosmosContainerIdentity; @@ -572,9 +573,19 @@ DiagnosticsProvider getDiagnosticsProvider() { * @param group Throughput control group going to be enabled. * @param throughputQueryMono The throughput query mono. */ - void enableThroughputControlGroup(ThroughputControlGroupInternal group, Mono throughputQueryMono) { + void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono throughputQueryMono) { checkNotNull(group, "Throughput control group cannot be null"); - this.asyncDocumentClient.enableThroughputControlGroup(group, throughputQueryMono); + this.asyncDocumentClient.enableSDKThroughputControlGroup(group, throughputQueryMono); + } + + /*** + * Enable server throughput control group. + * + * @param group the server throughput control group. + */ + void enableServerThroughputControlGroup(ServerThroughputControlGroupInternal group) { + checkNotNull(group, "Argument 'group' can not be null"); + this.asyncDocumentClient.enableServerThroughputControlGroup(group); } /*** diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index dc4aa1433f0b..be75c1965a78 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -37,9 +37,10 @@ import com.azure.cosmos.implementation.feedranges.FeedRangeInternal; import com.azure.cosmos.implementation.routing.PartitionKeyInternal; import com.azure.cosmos.implementation.routing.Range; -import com.azure.cosmos.implementation.throughputControl.config.GlobalThroughputControlGroup; -import com.azure.cosmos.implementation.throughputControl.config.LocalThroughputControlGroup; -import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupFactory; +import com.azure.cosmos.implementation.throughputControl.sdk.config.GlobalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.sdk.config.LocalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.server.config.ThroughputBucketControlGroup; +import com.azure.cosmos.implementation.throughputControl.ThroughputControlGroupFactory; import com.azure.cosmos.models.CosmosBatch; import com.azure.cosmos.models.CosmosBatchOperationResult; import com.azure.cosmos.models.CosmosBatchRequestOptions; @@ -81,7 +82,6 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; -import java.util.UUID; import java.util.concurrent.Callable; import java.util.concurrent.CompletionException; import java.util.concurrent.atomic.AtomicBoolean; @@ -2783,7 +2783,7 @@ void enableLocalThroughputControlGroup( LocalThroughputControlGroup localControlGroup = ThroughputControlGroupFactory.createThroughputLocalControlGroup(groupConfig, this); - this.database.getClient().enableThroughputControlGroup(localControlGroup, throughputQueryMono); + this.database.getClient().enableSDKThroughputControlGroup(localControlGroup, throughputQueryMono); } /** @@ -2808,7 +2808,7 @@ void enableLocalThroughputControlGroup( * * * - * @param groupConfig The throughput control group configuration, see {@link GlobalThroughputControlGroup}. + * @param groupConfig The throughput control group configuration, see {@link ThroughputControlGroupConfig}. * @param globalControlConfig The global throughput control configuration, see {@link GlobalThroughputControlConfig}. */ public void enableGlobalThroughputControlGroup( @@ -2818,10 +2818,11 @@ public void enableGlobalThroughputControlGroup( this.enableGlobalThroughputControlGroup(groupConfig, globalControlConfig, null); } + /*** * Only used internally. *
- * @param groupConfig The throughput control group configuration, see {@link GlobalThroughputControlGroup}. + * @param groupConfig The throughput control group configuration, see {@link ThroughputControlGroupConfig}. * @param globalControlConfig The global throughput control configuration, see {@link GlobalThroughputControlConfig}. * @param throughputQueryMono The throughput query mono. */ @@ -2833,7 +2834,22 @@ void enableGlobalThroughputControlGroup( GlobalThroughputControlGroup globalControlGroup = ThroughputControlGroupFactory.createThroughputGlobalControlGroup(groupConfig, globalControlConfig, this); - this.database.getClient().enableThroughputControlGroup(globalControlGroup, throughputQueryMono); + this.database.getClient().enableSDKThroughputControlGroup(globalControlGroup, throughputQueryMono); + } + + /*** + * Enable the service throughput bucket control group. + *

+ * For more information about throughput bucket please visit + * Throughput buckets in Azure Cosmos DB + * + * @param groupConfig the throughput control group config, see {@link ThroughputControlGroupConfig}. + */ + public void enableThroughputBucketControlGroup(ThroughputControlGroupConfig groupConfig) { + ThroughputBucketControlGroup throughputBucketControlGroup = + ThroughputControlGroupFactory.createThroughputBucketControlGroup(groupConfig, this); + + this.database.getClient().enableServerThroughputControlGroup(throughputBucketControlGroup); } void configureFaultInjectionProvider(IFaultInjectorProvider injectorProvider) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java index f9c18db1af18..2702a24377b0 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java @@ -3,7 +3,7 @@ package com.azure.cosmos; -import com.azure.cosmos.implementation.throughputControl.config.GlobalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.sdk.config.GlobalThroughputControlGroup; import com.azure.cosmos.models.CosmosBatch; import com.azure.cosmos.models.CosmosBatchOperationResult; import com.azure.cosmos.models.CosmosBatchRequestOptions; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfig.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfig.java index 81b8417666b5..de626b24e708 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfig.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfig.java @@ -13,6 +13,7 @@ public final class ThroughputControlGroupConfig { private final Integer targetThroughput; private final Double targetThroughputThreshold; private final PriorityLevel priorityLevel; + private final Integer throughputBucket; private final boolean isDefault; private final boolean continueOnInitError; @@ -21,12 +22,14 @@ public final class ThroughputControlGroupConfig { Integer targetThroughput, Double targetThroughputThreshold, PriorityLevel priorityLevel, + Integer throughputBucket, boolean isDefault, boolean continueOnInitError) { this.groupName = groupName; this.targetThroughput = targetThroughput; this.targetThroughputThreshold = targetThroughputThreshold; this.priorityLevel = priorityLevel; + this.throughputBucket = throughputBucket; this.isDefault = isDefault; this.continueOnInitError = continueOnInitError; } @@ -76,6 +79,19 @@ public Double getTargetThroughputThreshold() { */ public PriorityLevel getPriorityLevel() { return this.priorityLevel; } + + /*** + * Get the throughput bucket. + *

+ * For more information about throughput bucket please visit + * Throughput buckets in Azure Cosmos DB + * + * @return the throughput bucket of the throughput control group. + */ + public Integer getThroughputBucket() { + return this.throughputBucket; + } + /** * Get whether this throughput control group will be used by default. * diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfigBuilder.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfigBuilder.java index bdd8de91f3d6..9df41d5ce8e4 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfigBuilder.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfigBuilder.java @@ -19,6 +19,7 @@ public class ThroughputControlGroupConfigBuilder { private Double targetThroughputThreshold; private boolean isDefault; private PriorityLevel priorityLevel; + private Integer throughputBucket; private boolean continueOnInitError = DEFAULT_CONTINUE_ON_INIT_ERROR; /** @@ -157,6 +158,21 @@ public ThroughputControlGroupConfigBuilder defaultControlGroup(boolean aDefault) return this; } + /** + * Set the throughput bucket of the group. + *

+ * For more information about throughput bucket please visit + * Throughput buckets in Azure Cosmos DB + * + * @param throughputBucket the throughput bucket id. + * @return The {@link ThroughputControlGroupConfigBuilder}. + */ + public ThroughputControlGroupConfigBuilder throughputBucket(int throughputBucket) { + checkArgument(throughputBucket > 0, "Throughput Bucket should be greater than 0"); + this.throughputBucket = throughputBucket; + return this; + } + /** * Set whether allow request to continue on original request flow if throughput control controller failed on initialization. * If set to true, requests will be able to fall back to original request flow if throughput control controller failed on initialization. @@ -178,9 +194,14 @@ public ThroughputControlGroupConfig build() { if (StringUtils.isEmpty(this.groupName)) { throw new IllegalArgumentException("Group name cannot be null nor empty"); } - if (this.targetThroughput == null && this.targetThroughputThreshold == null && this.priorityLevel == null) { - throw new IllegalArgumentException("All targetThroughput, targetThroughputThreshold and priorityLevel cannot be null or empty."); + + if (this.targetThroughput == null + && this.targetThroughputThreshold == null + && this.priorityLevel == null + && this.throughputBucket == null) { // TODO: any exclusiveness between this two? + throw new IllegalArgumentException("All targetThroughput, targetThroughputThreshold, priorityLevel and throughput bucket cannot be null or empty."); } + if (this.targetThroughput == null && this.targetThroughputThreshold == null) { this.targetThroughput = Integer.MAX_VALUE; } @@ -190,6 +211,7 @@ public ThroughputControlGroupConfig build() { this.targetThroughput, this.targetThroughputThreshold, this.priorityLevel, + this.throughputBucket, this.isDefault, this.continueOnInitError); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java index ba66b3b9367c..3892b0b5c314 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java @@ -21,7 +21,9 @@ import com.azure.cosmos.implementation.directconnectivity.AddressSelector; import com.azure.cosmos.implementation.faultinjection.IFaultInjectorProvider; import com.azure.cosmos.implementation.query.PartitionedQueryExecutionInfo; -import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.server.config.ThroughputBucketControlGroup; import com.azure.cosmos.models.CosmosAuthorizationTokenResolver; import com.azure.cosmos.models.CosmosBatchResponse; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; @@ -1632,7 +1634,13 @@ Flux> readAllDocuments( * * @param group the throughput control group. */ - void enableThroughputControlGroup(ThroughputControlGroupInternal group, Mono throughputQueryMono); + void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono throughputQueryMono); + + /*** + * Enable server throughput control group. + * @param group the server throughput control group. + */ + void enableServerThroughputControlGroup(ServerThroughputControlGroupInternal group); /** * Submits open connection tasks and warms up caches for replicas for containers specified by diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/HttpConstants.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/HttpConstants.java index 9aa1e851dc96..8f8a8175b874 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/HttpConstants.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/HttpConstants.java @@ -291,6 +291,9 @@ public static class HttpHeaders { public static final String GLOBAL_DATABASE_ACCOUNT_NAME = "GlobalDatabaseAccountName"; public static final String THINCLIENT_START_EPK = "x-ms-thinclient-range-min"; public static final String THINCLIENT_END_EPK = "x-ms-thinclient-range-max"; + + // Throughput bucket header + public static final String THROUGHPUT_BUCKET = "x-ms-cosmos-throughput-bucket"; } public static class A_IMHeaderValues { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 6f22fb00fc51..657ab2c3b691 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -69,7 +69,8 @@ import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple; import com.azure.cosmos.implementation.spark.OperationListener; import com.azure.cosmos.implementation.throughputControl.ThroughputControlStore; -import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroupInternal; import com.azure.cosmos.models.CosmosAuthorizationTokenResolver; import com.azure.cosmos.models.CosmosBatchResponse; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; @@ -6253,9 +6254,22 @@ public void close() { } } @Override - public synchronized void enableThroughputControlGroup(ThroughputControlGroupInternal group, Mono throughputQueryMono) { + public void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono throughputQueryMono) { checkNotNull(group, "Throughput control group can not be null"); + this.enableThroughputControlStore(); + this.throughputControlStore.enableSDKThroughputControlGroup(group, throughputQueryMono); + } + + @Override + public void enableServerThroughputControlGroup(ServerThroughputControlGroupInternal group) { + checkNotNull(group, "Argument 'group' can not be null"); + + this.enableThroughputControlStore(); + this.throughputControlStore.enableServerThroughputControlGroup(group); + } + + private synchronized void enableThroughputControlStore() { if (this.throughputControlEnabled.compareAndSet(false, true)) { this.throughputControlStore = new ThroughputControlStore( @@ -6269,10 +6283,9 @@ public synchronized void enableThroughputControlGroup(ThroughputControlGroupInte this.gatewayProxy.enableThroughputControl(throughputControlStore); } } - - this.throughputControlStore.enableThroughputControlGroup(group, throughputQueryMono); } + @Override public Flux submitOpenConnectionTasksAndInitCaches(CosmosContainerProactiveInitConfig proactiveContainerInitConfig) { return this.storeModel.submitOpenConnectionTasksAndInitCaches(proactiveContainerInitConfig); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java index 92bb603ad727..a76035cfa3f7 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java @@ -1171,6 +1171,12 @@ public void setPriorityLevel(PriorityLevel priorityLevel) { } } + public void setThroughputBucket(Integer throughputBucket) { + if (throughputBucket != null) { + this.headers.put(HttpConstants.HttpHeaders.THROUGHPUT_BUCKET, throughputBucket.toString()); + } + } + public Duration getResponseTimeout() { return responseTimeout; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java index e332263685cf..0129ac84f10a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java @@ -598,7 +598,8 @@ public enum RntbdRequestHeader implements RntbdHeader { SDKSupportedCapabilities((short) 0x00A2, RntbdTokenType.ULong, false), ChangeFeedWireFormatVersion((short) 0x00B2, RntbdTokenType.String, false), PriorityLevel((short) 0x00BF, RntbdTokenType.Byte, false), - GlobalDatabaseAccountName((short) 0x00CE, RntbdTokenType.String, false); + GlobalDatabaseAccountName((short) 0x00CE, RntbdTokenType.String, false), + ThroughputBucket((short)0x00DB, RntbdTokenType.Byte, false); //TODO: does the throughput bucket supported in thin client public static final List thinClientHeadersInOrderList = Arrays.asList( EffectivePartitionKey, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java index 8f258eed4487..7c05b97387a4 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java @@ -132,6 +132,7 @@ final class RntbdRequestHeaders extends RntbdTokenStream { this.addChangeFeedWireFormatVersion(headers); this.addPriorityLevel(headers); this.addGlobalDatabaseAccountName(headers); + this.addThroughputBucket(headers); // Normal headers (Strings, Ints, Longs, etc.) @@ -291,6 +292,8 @@ private RntbdToken getCorrelatedActivityId() { private RntbdToken getPriorityLevel() { return this.get(RntbdRequestHeader.PriorityLevel); } + private RntbdToken getThroughputBucket() { return this.get(RntbdRequestHeader.ThroughputBucket); } + private RntbdToken getGlobalDatabaseAccountName() { return this.get(RntbdRequestHeader.GlobalDatabaseAccountName); } @@ -792,6 +795,16 @@ private void addPriorityLevel(final Map headers) } } + private void addThroughputBucket(final Map headers) + { + final String value = headers.get(HttpHeaders.THROUGHPUT_BUCKET); + + if (StringUtils.isNotEmpty(value)) { + final int throughputBucket = Integer.valueOf(value); + this.getThroughputBucket().setValue(throughputBucket); + } + } + private void addGlobalDatabaseAccountName(final Map headers) { final String value = headers.get(HttpHeaders.GLOBAL_DATABASE_ACCOUNT_NAME); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/EmptyThroughputContainerController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/EmptyThroughputContainerController.java similarity index 91% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/EmptyThroughputContainerController.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/EmptyThroughputContainerController.java index 987d9f652bfe..72aeb8e479a0 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/EmptyThroughputContainerController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/EmptyThroughputContainerController.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.container; +package com.azure.cosmos.implementation.throughputControl; import com.azure.cosmos.implementation.RxDocumentServiceRequest; import reactor.core.publisher.Mono; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/IThroughputContainerController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/IThroughputContainerController.java new file mode 100644 index 000000000000..89edffd24e3d --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/IThroughputContainerController.java @@ -0,0 +1,12 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl; + +import com.azure.cosmos.implementation.throughputControl.sdk.controller.IThroughputController; + +/** + * Represents a throughput container controller. + */ +public abstract class IThroughputContainerController implements IThroughputController { +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/ThroughputControlGroupFactory.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupFactory.java similarity index 70% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/ThroughputControlGroupFactory.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupFactory.java index c2a95c91593c..f91cc084756f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/ThroughputControlGroupFactory.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupFactory.java @@ -1,12 +1,15 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.config; +package com.azure.cosmos.implementation.throughputControl; import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.ThroughputControlGroupConfig; import com.azure.cosmos.GlobalThroughputControlConfig; +import com.azure.cosmos.implementation.throughputControl.sdk.config.GlobalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.sdk.config.LocalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.server.config.ThroughputBucketControlGroup; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; @@ -48,4 +51,19 @@ public static GlobalThroughputControlGroup createThroughputGlobalControlGroup( globalControlConfig.getControlItemExpireInterval()); } + + public static ThroughputBucketControlGroup createThroughputBucketControlGroup( + ThroughputControlGroupConfig groupConfig, + CosmosAsyncContainer targetContainer) { + + checkNotNull(groupConfig, "Throughput control group config can not be null"); + checkNotNull(targetContainer, "Throughput target container can not be null"); + + return new ThroughputBucketControlGroup( + groupConfig.getGroupName(), + groupConfig.isDefault(), + groupConfig.getPriorityLevel(), + groupConfig.getThroughputBucket(), + targetContainer); + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlStore.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlStore.java index e565a7774bc0..e34f2ec04caf 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlStore.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlStore.java @@ -5,86 +5,26 @@ import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.ConnectionMode; -import com.azure.cosmos.CosmosException; import com.azure.cosmos.implementation.ResourceType; import com.azure.cosmos.implementation.RxDocumentServiceRequest; import com.azure.cosmos.implementation.Utils; -import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; -import com.azure.cosmos.implementation.apachecommons.lang.tuple.Pair; -import com.azure.cosmos.implementation.caches.AsyncCache; import com.azure.cosmos.implementation.caches.RxClientCollectionCache; import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; -import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal; -import com.azure.cosmos.implementation.throughputControl.controller.IThroughputController; -import com.azure.cosmos.implementation.throughputControl.controller.container.EmptyThroughputContainerController; -import com.azure.cosmos.implementation.throughputControl.controller.container.IThroughputContainerController; -import com.azure.cosmos.implementation.throughputControl.controller.container.ThroughputContainerController; -import com.azure.cosmos.implementation.throughputControl.exceptions.ThroughputControlInitializationException; +import com.azure.cosmos.implementation.throughputControl.sdk.SDKThroughputControlStore; +import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.server.ServerThroughputControlStore; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroupInternal; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import reactor.core.Exceptions; import reactor.core.publisher.Mono; -import java.util.concurrent.ConcurrentHashMap; - -import static com.azure.cosmos.implementation.Exceptions.isNameCacheStale; -import static com.azure.cosmos.implementation.Exceptions.isPartitionKeyMismatchException; -import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; -/** - * This is the entrance class for the whole throughput control work flow pipeline. - * The pipeline will consist of controllers which is implementation of {@link IThroughputController} and {@link ThroughputRequestThrottler}. - * - * Following is a high-level diagram of the pipeline: - * - * +-------------------+ - * | Client | - * +-------------------+ - * | - * | - * | - * +---------------------------------------------------------+ - * | ThroughputControlStore | - * +---------------------------------------------------------+ - * / \ - * / \ - * / \ - * +---------------------------------------+ +---------------------------------------+ - * | Container A controller | | Container B controller | - * +---------------------------------------+ ... +---------------------------------------+ - * / \ - * / \ - * / \ - * +--------------------+ +---------------------+ - * | Group 1 controller | | Group 2 controller | - * +--------------------+ ... +---------------------+ - * | \ - * | \-------------\ - * | \ - * +---------------------------+ +-----------------------------+ - * |Global Request controller | OR | PkRanges Request controller | - * +--------------------------+ +------------------------------+ - * | / \ - * | / \ - * | / \ - * +------------------+ +------------------+ +------------------+ - * |Request throttler | |Request throttler | |Request throttler | - * +------------------+ +------------------+ ... +------------------+ - * - * - */ + public class ThroughputControlStore { private static final Logger logger = LoggerFactory.getLogger(ThroughputControlStore.class); - - private final RxClientCollectionCache collectionCache; - private final ConnectionMode connectionMode; - private final AsyncCache containerControllerCache; - private final ConcurrentHashMap containerMap; - private final RxPartitionKeyRangeCache partitionKeyRangeCache; - - private final LinkedCancellationTokenSource cancellationTokenSource; - private final ConcurrentHashMap cancellationTokenMap; + private final SDKThroughputControlStore sdkThroughputControlStore; + private final ServerThroughputControlStore serverThroughputControlStore; public ThroughputControlStore( RxClientCollectionCache collectionCache, @@ -94,41 +34,36 @@ public ThroughputControlStore( checkNotNull(collectionCache,"RxClientCollectionCache can not be null"); checkNotNull(partitionKeyRangeCache, "PartitionKeyRangeCache can not be null"); - this.collectionCache = collectionCache; - this.connectionMode = connectionMode; - this.containerControllerCache = new AsyncCache<>(); - this.containerMap = new ConcurrentHashMap<>(); - this.partitionKeyRangeCache = partitionKeyRangeCache; - - this.cancellationTokenSource = new LinkedCancellationTokenSource(); - this.cancellationTokenMap = new ConcurrentHashMap<>(); + this.sdkThroughputControlStore = new SDKThroughputControlStore(collectionCache, connectionMode, partitionKeyRangeCache); + this.serverThroughputControlStore = new ServerThroughputControlStore(); } - public void enableThroughputControlGroup(ThroughputControlGroupInternal group, Mono throughputQueryMono) { + public void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono throughputQueryMono) { checkNotNull(group, "Throughput control group cannot be null"); - String containerNameLink = Utils.trimBeginningAndEndingSlashes(BridgeInternal.extractContainerSelfLink(group.getTargetContainer())); - this.containerMap.compute(containerNameLink, (key, throughputControlContainerProperties) -> { - if (throughputControlContainerProperties == null) { - throughputControlContainerProperties = new ContainerThroughputControlGroupProperties(containerNameLink); + if (group.isDefault()) { + //verify a default group is not being defined in the server throughput control store + String containerNameLink = Utils.trimBeginningAndEndingSlashes(BridgeInternal.extractContainerSelfLink(group.getTargetContainer())); + if (this.serverThroughputControlStore.hasDefaultGroup(containerNameLink)) { + throw new IllegalArgumentException("A default group already exists"); } + } - int groupSizeBefore = throughputControlContainerProperties.getThroughputControlGroupSet().size(); - Pair stateAfterEnabling = - throughputControlContainerProperties.enableThroughputControlGroup(group, throughputQueryMono); + this.sdkThroughputControlStore.enableThroughputControlGroup(group, throughputQueryMono); + } - int groupSizeAfter = stateAfterEnabling.getLeft(); - boolean wasGroupConfigUpdated = stateAfterEnabling.getRight(); + public void enableServerThroughputControlGroup(ServerThroughputControlGroupInternal group) { + checkNotNull(group, "Throughput control group cannot be null"); - if ((groupSizeAfter > groupSizeBefore && groupSizeAfter == 1) || wasGroupConfigUpdated) { - // This is the first enabled group for the target container or an existing group was modified - // Clean the current cache in case we have built EmptyThroughputContainerController or an existing - // group was modified - this.containerControllerCache.remove(containerNameLink); + if (group.isDefault()) { + //verify a default group is not being defined in the sdk throughput control store + String containerNameLink = Utils.trimBeginningAndEndingSlashes(BridgeInternal.extractContainerSelfLink(group.getTargetContainer())); + if (this.sdkThroughputControlStore.hasDefaultGroup(containerNameLink)) { + throw new IllegalArgumentException("A default group already exists"); } + } - return throughputControlContainerProperties; - }); + this.serverThroughputControlStore.enableThroughputControlGroup(group); } public Mono processRequest(RxDocumentServiceRequest request, Mono originalRequestMono) { @@ -142,162 +77,28 @@ public Mono processRequest(RxDocumentServiceRequest request, Mono orig } String collectionNameLink = Utils.getCollectionName(request.getResourceAddress()); - return this.resolveContainerController(collectionNameLink) - .flatMap(containerController -> { - if (containerController.canHandleRequest(request)) { - return containerController.processRequest(request, originalRequestMono) - .doOnError(throwable -> this.handleException(collectionNameLink, request, throwable)); - } - - // Unable to find container controller to handle the request, - // It is caused by control store out of sync or the request has staled info. - // We will handle the first scenario by creating a new container controller, - // while fall back to original request Mono for the second scenario. - return this.updateControllerAndRetry(collectionNameLink, request, originalRequestMono); - }) - .onErrorResume(throwable -> { - - Exception unwrappedException = Utils.as(Exceptions.unwrap(throwable), Exception.class); - if (unwrappedException instanceof ThroughputControlInitializationException) { - if (this.shouldContinueRequestOnInitError(request, collectionNameLink, unwrappedException)) { - return originalRequestMono; - } - - return Mono.error(unwrappedException.getCause()); - } - - return Mono.error(throwable); - }); - } - - private boolean shouldContinueRequestOnInitError(RxDocumentServiceRequest request, String collectionNameLink, Throwable throwable) { - if (throwable instanceof ThroughputControlInitializationException) { - ContainerThroughputControlGroupProperties throughputControlContainerProperties = this.containerMap.get(collectionNameLink); - - checkNotNull( - throughputControlContainerProperties, - "Throughput control container properties should not be null"); - checkArgument( - throughputControlContainerProperties.getThroughputControlGroupSet().size() > 0, - "There should be more than one throughput control group"); - - return throughputControlContainerProperties.allowRequestToContinueOnInitError(request); + if (this.serverThroughputControlStore.hasGroup(collectionNameLink, request.getThroughputControlGroupName())) { + return this.serverThroughputControlStore.processRequest(request, originalRequestMono); } - return false; - } - - private Mono updateControllerAndRetry( - String containerNameLink, - RxDocumentServiceRequest request, - Mono originalRequestMono) { - - return this.shouldRefreshContainerController(containerNameLink, request) - .flatMap(shouldRefresh -> { - if (shouldRefresh) { - this.cancellationTokenMap.compute(containerNameLink, (key, cancellationToken) -> { - if (cancellationToken != null) { - cancellationToken.cancel(); - } - - return null; - }); - - this.containerControllerCache.refresh(containerNameLink, () -> this.createAndInitContainerController(containerNameLink)); - return this.resolveContainerController(containerNameLink) - .flatMap(updatedContainerController -> { - if (updatedContainerController.canHandleRequest(request)) { - return updatedContainerController.processRequest(request, originalRequestMono) - .doOnError(throwable -> this.handleException(containerNameLink, request, throwable)); - } else { - // still can not handle the request - logger.warn( - "Can not find container controller to process request {} with collectionRid {} ", - request.getActivityId(), - request.requestContext.resolvedCollectionRid); - - return originalRequestMono; - } - }); - } - - return originalRequestMono; - }); - } - - private Mono resolveContainerController(String containerNameLink) { - checkArgument(StringUtils.isNotEmpty(containerNameLink), "Container name link can not be null or empty"); - - return this.containerControllerCache.getAsync( - containerNameLink, - null, - () -> this.createAndInitContainerController(containerNameLink)) - .onErrorResume(throwable -> Mono.error(new ThroughputControlInitializationException(throwable))); - } - - private Mono createAndInitContainerController(String containerNameLink) { - checkArgument(StringUtils.isNotEmpty(containerNameLink), "Container link should not be null or empty"); - - if (this.containerMap.containsKey(containerNameLink)) { - return Mono.just(this.containerMap.get(containerNameLink)) - .flatMap(throughputControlContainerProperties -> { - LinkedCancellationToken parentToken = - this.cancellationTokenMap.compute( - containerNameLink, - (key, cancellationToken) -> this.cancellationTokenSource.getToken()); - - ThroughputContainerController containerController = - new ThroughputContainerController( - this.collectionCache, - this.connectionMode, - throughputControlContainerProperties.getThroughputControlGroupSet(), - this.partitionKeyRangeCache, - parentToken, - throughputControlContainerProperties.getThroughputQueryMono()); - - return containerController.init(); - }); - } else { - return Mono.just(new EmptyThroughputContainerController()) - .flatMap(EmptyThroughputContainerController::init); + if (this.sdkThroughputControlStore.hasGroup(collectionNameLink, request.getThroughputControlGroupName())) { + return this.sdkThroughputControlStore.processRequest(request, originalRequestMono); } - } - - private Mono shouldRefreshContainerController(String containerLink, RxDocumentServiceRequest request) { - // TODO: populate diagnostics - return this.collectionCache.resolveByNameAsync(null, containerLink, null) - .flatMap(documentCollection -> - Mono.just(StringUtils.equals(documentCollection.getResourceId(), request.requestContext.resolvedCollectionRid))); - } - - private void handleException(String containerNameLink, RxDocumentServiceRequest request, Throwable throwable) { - checkArgument(StringUtils.isNotEmpty(containerNameLink), "Container name link can not be null nor empty"); - checkNotNull(request, "Request can not be null"); - checkNotNull(throwable, "Exception can not be null"); - CosmosException cosmosException = Utils.as(Exceptions.unwrap(throwable), CosmosException.class); - - if (cosmosException != null && - (isNameCacheStale(cosmosException) || isPartitionKeyMismatchException(cosmosException))) { - - this.cancellationTokenMap.compute(containerNameLink,(key, cancellationToken) -> { - if (cancellationToken != null) { - cancellationToken.cancel(); - } - return null; - }); - - String containerLink = Utils.getCollectionName(request.getResourceAddress()); + // can not find exact throughput control group mapping, using default group if any + if (this.serverThroughputControlStore.hasDefaultGroup(collectionNameLink)) { + return this.serverThroughputControlStore.processRequest(request, originalRequestMono); + } - this.collectionCache.refresh(null, containerLink, null); - this.containerControllerCache.refresh( - containerLink, - () -> createAndInitContainerController(containerLink) - ); + if (this.sdkThroughputControlStore.hasDefaultGroup(collectionNameLink)) { + return this.serverThroughputControlStore.processRequest(request, originalRequestMono); } + + // neither store can process the request, fallback to just use the original mono + return originalRequestMono; } public void close() { - this.cancellationTokenSource.close(); + this.sdkThroughputControlStore.close(); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/IThroughputContainerController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/IThroughputContainerController.java deleted file mode 100644 index 35a0629f20e9..000000000000 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/IThroughputContainerController.java +++ /dev/null @@ -1,13 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -package com.azure.cosmos.implementation.throughputControl.controller.container; - -import com.azure.cosmos.implementation.throughputControl.controller.IThroughputController; - -/** - * Represents a throughput container controller. - */ -public interface IThroughputContainerController extends IThroughputController { - -} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ContainerThroughputControlGroupProperties.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerSDKThroughputControlGroupProperties.java similarity index 71% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ContainerThroughputControlGroupProperties.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerSDKThroughputControlGroupProperties.java index d67f0e400266..a675a723cd92 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ContainerThroughputControlGroupProperties.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerSDKThroughputControlGroupProperties.java @@ -1,16 +1,18 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl; +package com.azure.cosmos.implementation.throughputControl.sdk; import com.azure.cosmos.implementation.RxDocumentServiceRequest; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.implementation.apachecommons.lang.tuple.Pair; -import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Mono; +import java.util.HashSet; +import java.util.Map; import java.util.Objects; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; @@ -20,21 +22,21 @@ import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; -public class ContainerThroughputControlGroupProperties { - private static Logger logger = LoggerFactory.getLogger(ContainerThroughputControlGroupProperties.class); +public class ContainerSDKThroughputControlGroupProperties { + private static Logger logger = LoggerFactory.getLogger(ContainerSDKThroughputControlGroupProperties.class); - private final AtomicReference defaultGroup; - private final Set throughputControlGroupSet; + private final AtomicReference defaultGroup; + private final Map throughputControlGroups; private final Set supressInitErrorGroupSet; private final AtomicReference> throughputQueryMonoReference; private final String containerNameLink; - public ContainerThroughputControlGroupProperties(String containerNameLink) { + public ContainerSDKThroughputControlGroupProperties(String containerNameLink) { checkArgument(StringUtils.isNotEmpty(containerNameLink), "Argument 'containerNameLink' should not be empty"); - this.containerNameLink = containerNameLink; this.defaultGroup = new AtomicReference<>(); - this.throughputControlGroupSet = ConcurrentHashMap.newKeySet(); + this.containerNameLink = containerNameLink; + this.throughputControlGroups = new ConcurrentHashMap<>(); this.supressInitErrorGroupSet = ConcurrentHashMap.newKeySet(); this.throughputQueryMonoReference = new AtomicReference<>(); } @@ -42,11 +44,11 @@ public ContainerThroughputControlGroupProperties(String containerNameLink) { /*** * Enable a throughput control group. * - * @param group a {@link ThroughputControlGroupInternal}. + * @param group a {@link SDKThroughputControlGroupInternal}. * * @return the total size of distinct throughput control groups enabled on the container. */ - public Pair enableThroughputControlGroup(ThroughputControlGroupInternal group, Mono throughputQueryMono) { + public Pair enableThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono throughputQueryMono) { checkNotNull(group, "Throughput control group should not be null"); if (group.isDefault()) { @@ -65,8 +67,8 @@ public Pair enableThroughputControlGroup(ThroughputControlGrou // Only throw when two different groups are using the same id (databaseId + containerId + groupName) this - .throughputControlGroupSet - .stream() + .throughputControlGroups + .values() .forEach(existingGroup -> { if (Objects.equals(existingGroup.getIdPrefix(), group.getIdPrefix()) && !existingGroup.equals(group)) { @@ -79,7 +81,7 @@ public Pair enableThroughputControlGroup(ThroughputControlGrou "cannot be changed."); } - this.throughputControlGroupSet.remove(existingGroup); + this.throughputControlGroups.remove(existingGroup); logger.info( "Throughput control group with id-prefix {} already exists with different config. " + "Will update config.", @@ -88,16 +90,17 @@ public Pair enableThroughputControlGroup(ThroughputControlGrou } }); - this.throughputControlGroupSet.add(group); + this.throughputControlGroups.put(group.getGroupName(), group); if (!this.throughputQueryMonoReference.compareAndSet(null, throughputQueryMono)) { logger.debug("ThroughputQueryMono has exists for container {}", containerNameLink); } - return Pair.of(this.throughputControlGroupSet.size(), updatedGroupConfig.get()); + return Pair.of(this.throughputControlGroups.size(), updatedGroupConfig.get()); } - public Set getThroughputControlGroupSet() { - return this.throughputControlGroupSet; + // TODO: double check the data structure here + public Map getThroughputControlGroups() { + return this.throughputControlGroups; } public Mono getThroughputQueryMono() { @@ -117,4 +120,15 @@ public boolean allowRequestToContinueOnInitError(RxDocumentServiceRequest reques return this.supressInitErrorGroupSet.contains(requestGroupName); } + + public boolean isNotEmpty() { + return !this.throughputControlGroups.isEmpty(); + } + public boolean hasDefaultGroup() { + return this.defaultGroup.get() != null; + } + + public boolean hasGroup(String groupName) { + return this.throughputControlGroups.containsKey(groupName); + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationToken.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/LinkedCancellationToken.java similarity index 96% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationToken.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/LinkedCancellationToken.java index c53ff1316fd0..90b09d70c4dc 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationToken.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/LinkedCancellationToken.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl; +package com.azure.cosmos.implementation.throughputControl.sdk; import java.util.ArrayList; import java.util.List; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationTokenSource.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/LinkedCancellationTokenSource.java similarity index 96% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationTokenSource.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/LinkedCancellationTokenSource.java index 0bbfd014fb7d..cf65b9ba2b5e 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationTokenSource.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/LinkedCancellationTokenSource.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl; +package com.azure.cosmos.implementation.throughputControl.sdk; import java.io.Closeable; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/SDKThroughputControlStore.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/SDKThroughputControlStore.java new file mode 100644 index 000000000000..38ae9f597ec8 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/SDKThroughputControlStore.java @@ -0,0 +1,326 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl.sdk; + +import com.azure.cosmos.BridgeInternal; +import com.azure.cosmos.ConnectionMode; +import com.azure.cosmos.CosmosException; +import com.azure.cosmos.implementation.ResourceType; +import com.azure.cosmos.implementation.RxDocumentServiceRequest; +import com.azure.cosmos.implementation.Utils; +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.implementation.apachecommons.lang.tuple.Pair; +import com.azure.cosmos.implementation.caches.AsyncCache; +import com.azure.cosmos.implementation.caches.RxClientCollectionCache; +import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; +import com.azure.cosmos.implementation.throughputControl.EmptyThroughputContainerController; +import com.azure.cosmos.implementation.throughputControl.IThroughputContainerController; +import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.IThroughputController; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.container.SDKThroughputContainerController; +import com.azure.cosmos.implementation.throughputControl.sdk.exceptions.ThroughputControlInitializationException; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import reactor.core.Exceptions; +import reactor.core.publisher.Mono; + +import java.util.HashSet; +import java.util.concurrent.ConcurrentHashMap; + +import static com.azure.cosmos.implementation.Exceptions.isNameCacheStale; +import static com.azure.cosmos.implementation.Exceptions.isPartitionKeyMismatchException; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +/** + * This is the entrance class for the whole throughput control work flow pipeline. + * The pipeline will consist of controllers which is implementation of {@link IThroughputController} and {@link ThroughputRequestThrottler}. + * + * Following is a high-level diagram of the pipeline: + * + * +-------------------+ + * | Client | + * +-------------------+ + * | + * | + * | + * +---------------------------------------------------------+ + * | ThroughputControlStore | + * +---------------------------------------------------------+ + * / \ + * / \ + * / \ + * +---------------------------------------+ +---------------------------------------+ + * | Container A controller | | Container B controller | + * +---------------------------------------+ ... +---------------------------------------+ + * / \ + * / \ + * / \ + * +--------------------+ +---------------------+ + * | Group 1 controller | | Group 2 controller | + * +--------------------+ ... +---------------------+ + * | \ + * | \-------------\ + * | \ + * +---------------------------+ +-----------------------------+ + * |Global Request controller | OR | PkRanges Request controller | + * +--------------------------+ +------------------------------+ + * | / \ + * | / \ + * | / \ + * +------------------+ +------------------+ +------------------+ + * |Request throttler | |Request throttler | |Request throttler | + * +------------------+ +------------------+ ... +------------------+ + * + * + */ +public class SDKThroughputControlStore { + private static final Logger logger = LoggerFactory.getLogger(SDKThroughputControlStore.class); + + private final RxClientCollectionCache collectionCache; + private final ConnectionMode connectionMode; + private final AsyncCache containerControllerCache; + private final ConcurrentHashMap containerMap; + private final RxPartitionKeyRangeCache partitionKeyRangeCache; + + private final LinkedCancellationTokenSource cancellationTokenSource; + private final ConcurrentHashMap cancellationTokenMap; + + public SDKThroughputControlStore( + RxClientCollectionCache collectionCache, + ConnectionMode connectionMode, + RxPartitionKeyRangeCache partitionKeyRangeCache) { + + checkNotNull(collectionCache,"RxClientCollectionCache can not be null"); + checkNotNull(partitionKeyRangeCache, "PartitionKeyRangeCache can not be null"); + + this.collectionCache = collectionCache; + this.connectionMode = connectionMode; + this.containerControllerCache = new AsyncCache<>(); + this.containerMap = new ConcurrentHashMap<>(); + this.partitionKeyRangeCache = partitionKeyRangeCache; + + this.cancellationTokenSource = new LinkedCancellationTokenSource(); + this.cancellationTokenMap = new ConcurrentHashMap<>(); + } + + public void enableThroughputControlGroup( + SDKThroughputControlGroupInternal group, + Mono throughputQueryMono) { + checkNotNull(group, "Throughput control group cannot be null"); + + String containerNameLink = Utils.trimBeginningAndEndingSlashes(BridgeInternal.extractContainerSelfLink(group.getTargetContainer())); + this.containerMap.compute(containerNameLink, (key, throughputControlContainerProperties) -> { + if (throughputControlContainerProperties == null) { + throughputControlContainerProperties = new ContainerSDKThroughputControlGroupProperties(containerNameLink); + } + + int groupSizeBefore = throughputControlContainerProperties.getThroughputControlGroups().size(); + Pair stateAfterEnabling = + throughputControlContainerProperties.enableThroughputControlGroup(group, throughputQueryMono); + + int groupSizeAfter = stateAfterEnabling.getLeft(); + boolean wasGroupConfigUpdated = stateAfterEnabling.getRight(); + + if ((groupSizeAfter > groupSizeBefore && groupSizeAfter == 1) || wasGroupConfigUpdated) { + // This is the first enabled group for the target container or an existing group was modified + // Clean the current cache in case we have built EmptyThroughputContainerController or an existing + // group was modified + this.containerControllerCache.remove(containerNameLink); + } + + return throughputControlContainerProperties; + }); + } + + public Mono processRequest(RxDocumentServiceRequest request, Mono originalRequestMono) { + checkNotNull(request, "Request can not be null"); + checkNotNull(originalRequestMono, "originalRequestMono can not be null"); + + // Currently, we will only target two resource types. + // If in the future we find other useful scenarios for throughput control, add more resource type here. + if (request.getResourceType() != ResourceType.Document && request.getResourceType() != ResourceType.StoredProcedure) { + return originalRequestMono; + } + + String collectionNameLink = Utils.getCollectionName(request.getResourceAddress()); + return this.resolveContainerController(collectionNameLink) + .flatMap(containerController -> { + if (containerController.canHandleRequest(request)) { + return containerController.processRequest(request, originalRequestMono) + .doOnError(throwable -> this.handleException(collectionNameLink, request, throwable)); + } + + // Unable to find container controller to handle the request, + // It is caused by control store out of sync or the request has staled info. + // We will handle the first scenario by creating a new container controller, + // while fall back to original request Mono for the second scenario. + return this.updateControllerAndRetry(collectionNameLink, request, originalRequestMono); + }) + .onErrorResume(throwable -> { + + Exception unwrappedException = Utils.as(Exceptions.unwrap(throwable), Exception.class); + if (unwrappedException instanceof ThroughputControlInitializationException) { + if (this.shouldContinueRequestOnInitError(request, collectionNameLink, unwrappedException)) { + return originalRequestMono; + } + + return Mono.error(unwrappedException.getCause()); + } + + return Mono.error(throwable); + }); + } + + private boolean shouldContinueRequestOnInitError(RxDocumentServiceRequest request, String collectionNameLink, Throwable throwable) { + if (throwable instanceof ThroughputControlInitializationException) { + ContainerSDKThroughputControlGroupProperties throughputControlContainerProperties = this.containerMap.get(collectionNameLink); + + checkNotNull( + throughputControlContainerProperties, + "Throughput control container properties should not be null"); + checkArgument( + throughputControlContainerProperties.getThroughputControlGroups().size() > 0, + "There should be more than one throughput control group"); + + return throughputControlContainerProperties.allowRequestToContinueOnInitError(request); + } + + return false; + } + + private Mono updateControllerAndRetry( + String containerNameLink, + RxDocumentServiceRequest request, + Mono originalRequestMono) { + + return this.shouldRefreshContainerController(containerNameLink, request) + .flatMap(shouldRefresh -> { + if (shouldRefresh) { + this.cancellationTokenMap.compute(containerNameLink, (key, cancellationToken) -> { + if (cancellationToken != null) { + cancellationToken.cancel(); + } + + return null; + }); + + this.containerControllerCache.refresh(containerNameLink, () -> this.createAndInitContainerController(containerNameLink)); + return this.resolveContainerController(containerNameLink) + .flatMap(updatedContainerController -> { + if (updatedContainerController.canHandleRequest(request)) { + return updatedContainerController.processRequest(request, originalRequestMono) + .doOnError(throwable -> this.handleException(containerNameLink, request, throwable)); + } else { + // still can not handle the request + logger.warn( + "Can not find container controller to process request {} with collectionRid {} ", + request.getActivityId(), + request.requestContext.resolvedCollectionRid); + + return originalRequestMono; + } + }); + } + + return originalRequestMono; + }); + } + + private Mono resolveContainerController(String containerNameLink) { + checkArgument(StringUtils.isNotEmpty(containerNameLink), "Container name link can not be null or empty"); + + return this.containerControllerCache.getAsync( + containerNameLink, + null, + () -> this.createAndInitContainerController(containerNameLink)) + .onErrorResume(throwable -> Mono.error(new ThroughputControlInitializationException(throwable))); + } + + private Mono createAndInitContainerController(String containerNameLink) { + checkArgument(StringUtils.isNotEmpty(containerNameLink), "Container link should not be null or empty"); + + if (this.containerMap.containsKey(containerNameLink)) { + return Mono.just(this.containerMap.get(containerNameLink)) + .flatMap(throughputControlContainerProperties -> { + LinkedCancellationToken parentToken = + this.cancellationTokenMap.compute( + containerNameLink, + (key, cancellationToken) -> this.cancellationTokenSource.getToken()); + + SDKThroughputContainerController containerController = + new SDKThroughputContainerController( + this.collectionCache, + this.connectionMode, + new HashSet<>(throughputControlContainerProperties.getThroughputControlGroups().values()), + this.partitionKeyRangeCache, + parentToken, + throughputControlContainerProperties.getThroughputQueryMono()); + + return containerController.init(); + }); + } else { + return Mono.just(new EmptyThroughputContainerController()) + .flatMap(EmptyThroughputContainerController::init); + } + } + + private Mono shouldRefreshContainerController(String containerLink, RxDocumentServiceRequest request) { + // TODO: populate diagnostics + return this.collectionCache.resolveByNameAsync(null, containerLink, null) + .flatMap(documentCollection -> + Mono.just(StringUtils.equals(documentCollection.getResourceId(), request.requestContext.resolvedCollectionRid))); + } + + private void handleException(String containerNameLink, RxDocumentServiceRequest request, Throwable throwable) { + checkArgument(StringUtils.isNotEmpty(containerNameLink), "Container name link can not be null nor empty"); + checkNotNull(request, "Request can not be null"); + checkNotNull(throwable, "Exception can not be null"); + + CosmosException cosmosException = Utils.as(Exceptions.unwrap(throwable), CosmosException.class); + + if (cosmosException != null && + (isNameCacheStale(cosmosException) || isPartitionKeyMismatchException(cosmosException))) { + + this.cancellationTokenMap.compute(containerNameLink,(key, cancellationToken) -> { + if (cancellationToken != null) { + cancellationToken.cancel(); + } + return null; + }); + + String containerLink = Utils.getCollectionName(request.getResourceAddress()); + + this.collectionCache.refresh(null, containerLink, null); + this.containerControllerCache.refresh( + containerLink, + () -> createAndInitContainerController(containerLink) + ); + } + } + + public boolean hasDefaultGroup(String containerNameLink) { + if (this.containerMap.containsKey(containerNameLink)) { + return this.containerMap.get(containerNameLink).hasDefaultGroup(); + } + + return false; + } + + public boolean hasGroup(String containerNameLink, String throughputControlGroupName) { + if (StringUtils.isEmpty(throughputControlGroupName)) { + return false; + } + + if (this.containerMap.containsKey(containerNameLink)) { + return this.containerMap.get(containerNameLink).hasGroup(throughputControlGroupName); + } + + return false; + } + + public void close() { + this.cancellationTokenSource.close(); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTrackingUnit.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputControlTrackingUnit.java similarity index 98% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTrackingUnit.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputControlTrackingUnit.java index 020f56c8f172..1d24df68a90f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTrackingUnit.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputControlTrackingUnit.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl; +package com.azure.cosmos.implementation.throughputControl.sdk; import com.azure.cosmos.implementation.OperationType; import org.slf4j.Logger; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputRequestThrottler.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputRequestThrottler.java similarity index 99% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputRequestThrottler.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputRequestThrottler.java index e54632cea1e7..82411afdd66d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputRequestThrottler.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputRequestThrottler.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl; +package com.azure.cosmos.implementation.throughputControl.sdk; import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.CosmosException; @@ -20,7 +20,6 @@ import reactor.core.Exceptions; import reactor.core.publisher.Mono; -import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicReference; import java.util.concurrent.locks.ReentrantReadWriteLock; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/GlobalThroughputControlGroup.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/GlobalThroughputControlGroup.java similarity index 89% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/GlobalThroughputControlGroup.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/GlobalThroughputControlGroup.java index 5ee5b55fd3e5..bd0574e7d354 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/GlobalThroughputControlGroup.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/GlobalThroughputControlGroup.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.config; +package com.azure.cosmos.implementation.throughputControl.sdk.config; import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.models.PriorityLevel; @@ -10,7 +10,7 @@ import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; -public class GlobalThroughputControlGroup extends ThroughputControlGroupInternal { +public class GlobalThroughputControlGroup extends SDKThroughputControlGroupInternal { private static final Duration DEFAULT_CONTROL_ITEM_RENEW_INTERVAL = Duration.ofSeconds(5); private final CosmosAsyncContainer globalControlContainer; @@ -29,7 +29,14 @@ public GlobalThroughputControlGroup( Duration controlItemRenewInterval, Duration controlItemExpireInterval) { - super (groupName, targetContainer, targetThroughput, targetThroughputThreshold, priorityLevel, isDefault, continueOnInitError); + super ( + groupName, + targetContainer, + targetThroughput, + targetThroughputThreshold, + priorityLevel, + isDefault, + continueOnInitError); checkNotNull(globalControlContainer, "Global control container can not be null"); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/LocalThroughputControlGroup.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/LocalThroughputControlGroup.java similarity index 56% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/LocalThroughputControlGroup.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/LocalThroughputControlGroup.java index 15476deb9fb4..6d1a2caa672f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/LocalThroughputControlGroup.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/LocalThroughputControlGroup.java @@ -1,12 +1,12 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.config; +package com.azure.cosmos.implementation.throughputControl.sdk.config; import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.models.PriorityLevel; -public class LocalThroughputControlGroup extends ThroughputControlGroupInternal { +public class LocalThroughputControlGroup extends SDKThroughputControlGroupInternal { public LocalThroughputControlGroup( String groupName, @@ -16,6 +16,13 @@ public LocalThroughputControlGroup( PriorityLevel priorityLevel, boolean isDefault, boolean continueOnInitError) { - super (groupName, targetContainer, targetThroughput, targetThroughputThreshold, priorityLevel, isDefault, continueOnInitError); + super ( + groupName, + targetContainer, + targetThroughput, + targetThroughputThreshold, + priorityLevel, + isDefault, + continueOnInitError); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/ThroughputControlGroupInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/SDKThroughputControlGroupInternal.java similarity index 92% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/ThroughputControlGroupInternal.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/SDKThroughputControlGroupInternal.java index 3ffa06c2933a..3ca028425f17 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/ThroughputControlGroupInternal.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/SDKThroughputControlGroupInternal.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.config; +package com.azure.cosmos.implementation.throughputControl.sdk.config; import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; @@ -12,7 +12,7 @@ import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; -public abstract class ThroughputControlGroupInternal { +public abstract class SDKThroughputControlGroupInternal { private final String groupName; private final String idPrefix; private final String id; @@ -21,9 +21,9 @@ public abstract class ThroughputControlGroupInternal { private final CosmosAsyncContainer targetContainer; private final Integer targetThroughput; private final Double targetThroughputThreshold; - private PriorityLevel priorityLevel; + private final PriorityLevel priorityLevel; - public ThroughputControlGroupInternal( + public SDKThroughputControlGroupInternal( String groupName, CosmosAsyncContainer targetContainer, Integer targetThroughput, @@ -59,7 +59,11 @@ public ThroughputControlGroupInternal( getThroughputIdSuffix(targetThroughput, targetThroughputThreshold, priorityLevel)); } - public static String getThroughputIdSuffix(Integer throughput, Double throughputThreshold, PriorityLevel priorityLevel) { + public static String getThroughputIdSuffix( + Integer throughput, + Double throughputThreshold, + PriorityLevel priorityLevel) { + StringBuilder sb = new StringBuilder(); if (throughput != null) { sb.append("t-"); @@ -186,7 +190,7 @@ public boolean equals(Object other) { return false; } - ThroughputControlGroupInternal that = (ThroughputControlGroupInternal) other; + SDKThroughputControlGroupInternal that = (SDKThroughputControlGroupInternal) other; return Objects.equals(this.id, that.id) && this.isDefault == that.isDefault @@ -196,6 +200,7 @@ public boolean equals(Object other) { && Objects.equals(this.priorityLevel, that.priorityLevel); } + public boolean hasSameIdentity(Object other) { if (this == other) { return true; @@ -205,7 +210,7 @@ public boolean hasSameIdentity(Object other) { return false; } - ThroughputControlGroupInternal that = (ThroughputControlGroupInternal) other; + SDKThroughputControlGroupInternal that = (SDKThroughputControlGroupInternal) other; return Objects.equals(this.idPrefix, that.idPrefix) && this.isDefault == that.isDefault diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/IThroughputController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/IThroughputController.java similarity index 94% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/IThroughputController.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/IThroughputController.java index 203479f8f9e4..355e8a3a0401 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/IThroughputController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/IThroughputController.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller; +package com.azure.cosmos.implementation.throughputControl.sdk.controller; import com.azure.cosmos.implementation.RxDocumentServiceRequest; import reactor.core.publisher.Mono; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/ThroughputContainerController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/container/SDKThroughputContainerController.java similarity index 92% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/ThroughputContainerController.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/container/SDKThroughputContainerController.java index 9f5e938f22db..eac4f7d6fcaa 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/ThroughputContainerController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/container/SDKThroughputContainerController.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.container; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.container; import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.ConnectionMode; @@ -22,12 +22,13 @@ import com.azure.cosmos.implementation.caches.AsyncCache; import com.azure.cosmos.implementation.caches.RxCollectionCache; import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; -import com.azure.cosmos.implementation.throughputControl.LinkedCancellationToken; -import com.azure.cosmos.implementation.throughputControl.LinkedCancellationTokenSource; -import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal; -import com.azure.cosmos.implementation.throughputControl.controller.group.ThroughputGroupControllerBase; -import com.azure.cosmos.implementation.throughputControl.controller.group.ThroughputGroupControllerFactory; -import com.azure.cosmos.implementation.throughputControl.exceptions.ThroughputControlInitializationException; +import com.azure.cosmos.implementation.throughputControl.IThroughputContainerController; +import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationToken; +import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationTokenSource; +import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.group.ThroughputGroupControllerBase; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.group.ThroughputGroupControllerFactory; +import com.azure.cosmos.implementation.throughputControl.sdk.exceptions.ThroughputControlInitializationException; import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.ModelBridgeInternal; import com.azure.cosmos.models.ThroughputProperties; @@ -53,8 +54,8 @@ * 2. Create and initialize throughput group controllers * 3. Start throughput refresh task if necessary */ -public class ThroughputContainerController implements IThroughputContainerController { - private static final Logger logger = LoggerFactory.getLogger(ThroughputContainerController.class); +public class SDKThroughputContainerController implements IThroughputContainerController { + private static final Logger logger = LoggerFactory.getLogger(SDKThroughputContainerController.class); private static final Duration DEFAULT_THROUGHPUT_REFRESH_INTERVAL = Duration.ofMinutes(15); private static final int NO_OFFER_EXCEPTION_STATUS_CODE = HttpConstants.StatusCodes.BADREQUEST; @@ -64,7 +65,7 @@ public class ThroughputContainerController implements IThroughputContainerContro private final RxCollectionCache collectionCache; private final ConnectionMode connectionMode; private final AsyncCache groupControllerCache; - private final Set groups; + private final Set groups; private final AtomicReference maxContainerThroughput; private final RxPartitionKeyRangeCache partitionKeyRangeCache; private final CosmosAsyncContainer targetContainer; @@ -78,10 +79,10 @@ public class ThroughputContainerController implements IThroughputContainerContro private String targetDatabaseRid; private ThroughputProvisioningScope throughputProvisioningScope; - public ThroughputContainerController( + public SDKThroughputContainerController( RxCollectionCache collectionCache, ConnectionMode connectionMode, - Set groups, + Set groups, RxPartitionKeyRangeCache partitionKeyRangeCache, LinkedCancellationToken parentToken, Mono throughputQueryMono) { @@ -108,7 +109,7 @@ public ThroughputContainerController( this.throughputQueryMono = throughputQueryMono == null ? this.resolveContainerMaxThroughputCore() : throughputQueryMono; } - private ThroughputProvisioningScope getThroughputResolveLevel(Set groupConfigs) { + private ThroughputProvisioningScope getThroughputResolveLevel(Set groupConfigs) { if (groupConfigs.stream().anyMatch(groupConfig -> groupConfig.getTargetThroughputThreshold() != null)) { // Throughput can be provisioned on container level or database level, will start from container return ThroughputProvisioningScope.CONTAINER; @@ -177,7 +178,7 @@ private Mono resolveContainerThroughput() { } } - private Mono resolveContainerMaxThroughput() { + private Mono resolveContainerMaxThroughput() { return this.throughputQueryMono .flatMap(maxThroughput -> { this.maxContainerThroughput.set(maxThroughput); @@ -311,7 +312,7 @@ private Mono> getOrCreateThroug return Mono.just(new Utils.ValueHolder<>(this.defaultGroupController)); } - for (ThroughputControlGroupInternal group : this.groups) { + for (SDKThroughputControlGroupInternal group : this.groups) { if (StringUtils.equals(groupName, group.getGroupName())) { return this.resolveThroughputGroupController(group) .map(Utils.ValueHolder::new); @@ -333,13 +334,13 @@ public boolean canHandleRequest(RxDocumentServiceRequest request) { return StringUtils.equals(this.targetContainerRid, request.requestContext.resolvedCollectionRid); } - private Mono createAndInitializeGroupControllers() { + private Mono createAndInitializeGroupControllers() { return Flux.fromIterable(this.groups) .flatMap(this::resolveThroughputGroupController) .then(Mono.just(this)); } - private Mono resolveThroughputGroupController(ThroughputControlGroupInternal group) { + private Mono resolveThroughputGroupController(SDKThroughputControlGroupInternal group) { return this.groupControllerCache.getAsync( group.getGroupName(), null, @@ -347,7 +348,7 @@ private Mono resolveThroughputGroupController(Thr .onErrorResume(throwable -> Mono.error(new ThroughputControlInitializationException(throwable))); } - private Mono createAndInitializeGroupController(ThroughputControlGroupInternal group) { + private Mono createAndInitializeGroupController(SDKThroughputControlGroupInternal group) { LinkedCancellationToken parentToken = this.cancellationTokenMap.compute( group.getGroupName(), diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/ThroughputProvisioningScope.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/container/ThroughputProvisioningScope.java similarity index 68% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/ThroughputProvisioningScope.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/container/ThroughputProvisioningScope.java index a161cb8ed830..c61a06dd9b71 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/ThroughputProvisioningScope.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/container/ThroughputProvisioningScope.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.container; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.container; public enum ThroughputProvisioningScope { NONE, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/ThroughputGroupControllerBase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/ThroughputGroupControllerBase.java similarity index 91% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/ThroughputGroupControllerBase.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/ThroughputGroupControllerBase.java index 54bae311e044..473ad31ee0a4 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/ThroughputGroupControllerBase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/ThroughputGroupControllerBase.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.group; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.group; import com.azure.cosmos.ConnectionMode; import com.azure.cosmos.CosmosException; @@ -11,14 +11,14 @@ import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.implementation.caches.AsyncCache; import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; -import com.azure.cosmos.implementation.throughputControl.LinkedCancellationToken; -import com.azure.cosmos.implementation.throughputControl.LinkedCancellationTokenSource; -import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal; -import com.azure.cosmos.implementation.throughputControl.controller.IThroughputController; -import com.azure.cosmos.implementation.throughputControl.controller.request.GlobalThroughputRequestController; -import com.azure.cosmos.implementation.throughputControl.controller.request.IThroughputRequestController; -import com.azure.cosmos.implementation.throughputControl.controller.request.PkRangesThroughputRequestController; -import com.azure.cosmos.implementation.throughputControl.exceptions.ThroughputControlInitializationException; +import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationToken; +import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationTokenSource; +import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.IThroughputController; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.request.GlobalThroughputRequestController; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.request.IThroughputRequestController; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.request.PkRangesThroughputRequestController; +import com.azure.cosmos.implementation.throughputControl.sdk.exceptions.ThroughputControlInitializationException; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.Exceptions; @@ -44,7 +44,7 @@ public abstract class ThroughputGroupControllerBase implements IThroughputContro private final Duration DEFAULT_THROUGHPUT_USAGE_RESET_DURATION = Duration.ofSeconds(1); private final ConnectionMode connectionMode; - private final ThroughputControlGroupInternal group; + private final SDKThroughputControlGroupInternal group; private final AtomicInteger maxContainerThroughput; private final RxPartitionKeyRangeCache partitionKeyRangeCache; private final AsyncCache requestControllerAsyncCache; @@ -55,7 +55,7 @@ public abstract class ThroughputGroupControllerBase implements IThroughputContro public ThroughputGroupControllerBase( ConnectionMode connectionMode, - ThroughputControlGroupInternal group, + SDKThroughputControlGroupInternal group, Integer maxContainerThroughput, RxPartitionKeyRangeCache partitionKeyRangeCache, String targetContainerRid, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/ThroughputGroupControllerFactory.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/ThroughputGroupControllerFactory.java similarity index 64% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/ThroughputGroupControllerFactory.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/ThroughputGroupControllerFactory.java index 01d1c9dbb040..12646e32933c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/ThroughputGroupControllerFactory.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/ThroughputGroupControllerFactory.java @@ -1,23 +1,22 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.group; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.group; import com.azure.cosmos.ConnectionMode; -import com.azure.cosmos.implementation.GlobalEndpointManager; import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; -import com.azure.cosmos.implementation.throughputControl.LinkedCancellationToken; -import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal; -import com.azure.cosmos.implementation.throughputControl.config.GlobalThroughputControlGroup; -import com.azure.cosmos.implementation.throughputControl.config.LocalThroughputControlGroup; -import com.azure.cosmos.implementation.throughputControl.controller.group.global.GlobalThroughputControlGroupController; -import com.azure.cosmos.implementation.throughputControl.controller.group.local.LocalThroughputControlGroupController; +import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationToken; +import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.sdk.config.GlobalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.sdk.config.LocalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.group.global.GlobalThroughputControlGroupController; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.group.local.LocalThroughputControlGroupController; public class ThroughputGroupControllerFactory { public static ThroughputGroupControllerBase createController( ConnectionMode connectionMode, - ThroughputControlGroupInternal group, + SDKThroughputControlGroupInternal group, Integer maxContainerThroughput, RxPartitionKeyRangeCache partitionKeyRangeCache, String targetCollectionRid, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlClientItem.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlClientItem.java similarity index 94% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlClientItem.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlClientItem.java index a20578e3f4e9..5e220c89261e 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlClientItem.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlClientItem.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.group.global; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.group.global; import com.fasterxml.jackson.annotation.JsonProperty; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlConfigItem.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlConfigItem.java similarity index 97% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlConfigItem.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlConfigItem.java index b1578ea9d435..314a49bd9c41 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlConfigItem.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlConfigItem.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.group.global; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.group.global; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.fasterxml.jackson.annotation.JsonIgnore; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlGroupController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlGroupController.java similarity index 94% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlGroupController.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlGroupController.java index 1916c38de7f3..53be0743edfc 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlGroupController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlGroupController.java @@ -1,15 +1,15 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.group.global; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.group.global; import com.azure.cosmos.ConnectionMode; import com.azure.cosmos.implementation.CosmosSchedulers; import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; import com.azure.cosmos.implementation.guava25.collect.EvictingQueue; -import com.azure.cosmos.implementation.throughputControl.LinkedCancellationToken; -import com.azure.cosmos.implementation.throughputControl.config.GlobalThroughputControlGroup; -import com.azure.cosmos.implementation.throughputControl.controller.group.ThroughputGroupControllerBase; +import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationToken; +import com.azure.cosmos.implementation.throughputControl.sdk.config.GlobalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.group.ThroughputGroupControllerBase; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlItem.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlItem.java similarity index 93% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlItem.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlItem.java index bfaeb5e2ba04..af6d87c239bf 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlItem.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlItem.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.group.global; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.group.global; import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.annotation.JsonProperty; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/ThroughputControlContainerManager.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/ThroughputControlContainerManager.java similarity index 98% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/ThroughputControlContainerManager.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/ThroughputControlContainerManager.java index fd382a10ead9..e5c87df99d43 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/ThroughputControlContainerManager.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/ThroughputControlContainerManager.java @@ -1,14 +1,14 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.group.global; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.group.global; import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.CosmosException; import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.UUIDs; import com.azure.cosmos.implementation.Utils; -import com.azure.cosmos.implementation.throughputControl.config.GlobalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.sdk.config.GlobalThroughputControlGroup; import com.azure.cosmos.models.CosmosItemRequestOptions; import com.azure.cosmos.models.PartitionKey; import com.azure.cosmos.models.SqlParameter; @@ -23,7 +23,6 @@ import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.List; -import java.util.UUID; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/ThroughputUsageSnapshot.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/ThroughputUsageSnapshot.java similarity index 92% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/ThroughputUsageSnapshot.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/ThroughputUsageSnapshot.java index e26c038cd4e9..c2b23da2e5c7 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/ThroughputUsageSnapshot.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/ThroughputUsageSnapshot.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.group.global; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.group.global; import java.time.Duration; import java.time.Instant; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/local/LocalThroughputControlGroupController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/local/LocalThroughputControlGroupController.java similarity index 79% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/local/LocalThroughputControlGroupController.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/local/LocalThroughputControlGroupController.java index 2c5410f5df6e..16def2b5b221 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/local/LocalThroughputControlGroupController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/local/LocalThroughputControlGroupController.java @@ -1,14 +1,14 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.group.local; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.group.local; import com.azure.cosmos.ConnectionMode; import com.azure.cosmos.implementation.CosmosSchedulers; import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; -import com.azure.cosmos.implementation.throughputControl.LinkedCancellationToken; -import com.azure.cosmos.implementation.throughputControl.config.LocalThroughputControlGroup; -import com.azure.cosmos.implementation.throughputControl.controller.group.ThroughputGroupControllerBase; +import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationToken; +import com.azure.cosmos.implementation.throughputControl.sdk.config.LocalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.group.ThroughputGroupControllerBase; import reactor.core.publisher.Mono; public class LocalThroughputControlGroupController extends ThroughputGroupControllerBase { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/request/GlobalThroughputRequestController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/request/GlobalThroughputRequestController.java similarity index 89% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/request/GlobalThroughputRequestController.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/request/GlobalThroughputRequestController.java index 6fdec7d63b4f..e2424a7ee0fb 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/request/GlobalThroughputRequestController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/request/GlobalThroughputRequestController.java @@ -1,11 +1,11 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.request; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.request; import com.azure.cosmos.implementation.RxDocumentServiceRequest; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; -import com.azure.cosmos.implementation.throughputControl.ThroughputRequestThrottler; +import com.azure.cosmos.implementation.throughputControl.sdk.ThroughputRequestThrottler; import reactor.core.publisher.Mono; import java.util.concurrent.atomic.AtomicReference; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/request/IThroughputRequestController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/request/IThroughputRequestController.java similarity index 80% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/request/IThroughputRequestController.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/request/IThroughputRequestController.java index 2e1fcb4437c3..56a66994ab8c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/request/IThroughputRequestController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/request/IThroughputRequestController.java @@ -1,9 +1,9 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.request; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.request; -import com.azure.cosmos.implementation.throughputControl.controller.IThroughputController; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.IThroughputController; /** * Represents a throughput request controller. diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/request/PkRangesThroughputRequestController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/request/PkRangesThroughputRequestController.java similarity index 96% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/request/PkRangesThroughputRequestController.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/request/PkRangesThroughputRequestController.java index 9d57c793e249..fd7f56918e3f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/request/PkRangesThroughputRequestController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/request/PkRangesThroughputRequestController.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.request; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.request; import com.azure.cosmos.implementation.PartitionKeyRange; import com.azure.cosmos.implementation.RxDocumentServiceRequest; @@ -9,7 +9,7 @@ import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper; import com.azure.cosmos.implementation.routing.Range; -import com.azure.cosmos.implementation.throughputControl.ThroughputRequestThrottler; +import com.azure.cosmos.implementation.throughputControl.sdk.ThroughputRequestThrottler; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Mono; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/exceptions/ThroughputControlInitializationException.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/exceptions/ThroughputControlInitializationException.java similarity index 93% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/exceptions/ThroughputControlInitializationException.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/exceptions/ThroughputControlInitializationException.java index ece6c75d1c79..a1345fe0c64b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/exceptions/ThroughputControlInitializationException.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/exceptions/ThroughputControlInitializationException.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.exceptions; +package com.azure.cosmos.implementation.throughputControl.sdk.exceptions; import com.azure.cosmos.implementation.throughputControl.ThroughputControlStore; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ContainerServerThroughputControlGroupProperties.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ContainerServerThroughputControlGroupProperties.java new file mode 100644 index 000000000000..6fd411ee854f --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ContainerServerThroughputControlGroupProperties.java @@ -0,0 +1,68 @@ +package com.azure.cosmos.implementation.throughputControl.server; + +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroupInternal; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicReference; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +public class ContainerServerThroughputControlGroupProperties { + private static Logger logger = LoggerFactory.getLogger(ContainerServerThroughputControlGroupProperties.class); + + private final String containerNameLink; + private final AtomicReference defaultGroup; + private final Map throughputControlGroups; + + public ContainerServerThroughputControlGroupProperties(String containerNameLink) { + checkArgument(StringUtils.isNotEmpty(containerNameLink), "Argument 'containerNameLink' should not be empty"); + + this.containerNameLink = containerNameLink; + this.defaultGroup = new AtomicReference<>(); + this.throughputControlGroups = new ConcurrentHashMap<>(); + } + + /*** + * Enable a server throughput control group. + * + * @param group a {@link ServerThroughputControlGroupInternal}. + * + * @return the total size of distinct server throughput control groups enabled on the container. + */ + public int enableThroughputControlGroup(ServerThroughputControlGroupInternal group) { + checkNotNull(group, "Argument 'group' should not be null'"); + + if (group.isDefault()) { + if (!this.defaultGroup.compareAndSet(null, group)) { + if (!this.defaultGroup.get().equals(group)) { + throw new IllegalArgumentException("A default group already exists"); + } + } + } + + ServerThroughputControlGroupInternal serverThroughputControlGroup = + this.throughputControlGroups.computeIfAbsent(group.getGroupName(), (key) -> group); + if (!serverThroughputControlGroup.equals(group)) { + throw new IllegalArgumentException("A group with same name already exists, name: " + group.getGroupName()); + } + + return this.throughputControlGroups.size(); + } + + public Map getThroughputControlGroups() { + return this.throughputControlGroups; + } + + public boolean hasDefaultGroup() { + return this.defaultGroup.get() != null; + } + + public boolean hasGroup(String groupName) { + return this.throughputControlGroups.containsKey(groupName); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ServerThroughputControlStore.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ServerThroughputControlStore.java new file mode 100644 index 000000000000..64309baa07f5 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ServerThroughputControlStore.java @@ -0,0 +1,109 @@ +package com.azure.cosmos.implementation.throughputControl.server; + +import com.azure.cosmos.BridgeInternal; +import com.azure.cosmos.implementation.ResourceType; +import com.azure.cosmos.implementation.RxDocumentServiceRequest; +import com.azure.cosmos.implementation.Utils; +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.implementation.caches.AsyncCache; +import com.azure.cosmos.implementation.throughputControl.EmptyThroughputContainerController; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.server.controller.ServerThroughputContainerController; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import reactor.core.publisher.Mono; + +import java.util.concurrent.ConcurrentHashMap; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +public class ServerThroughputControlStore { + private static final Logger logger = LoggerFactory.getLogger(ServerThroughputControlStore.class); + + private final AsyncCache containerControllerCache; + private final ConcurrentHashMap containerMap; + + public ServerThroughputControlStore() { + this.containerControllerCache = new AsyncCache<>(); + this.containerMap = new ConcurrentHashMap<>(); + } + + public void enableThroughputControlGroup(ServerThroughputControlGroupInternal group) { + checkNotNull(group, "Throughput control group cannot be null"); + + String containerNameLink = Utils.trimBeginningAndEndingSlashes(BridgeInternal.extractContainerSelfLink(group.getTargetContainer())); + this.containerMap.compute(containerNameLink, (key, throughputControlContainerProperties) -> { + if (throughputControlContainerProperties == null) { + throughputControlContainerProperties = new ContainerServerThroughputControlGroupProperties(containerNameLink); + } + + int groupSizeAfterEnabling = throughputControlContainerProperties.enableThroughputControlGroup(group); + if (groupSizeAfterEnabling == 1) { + // This is the first enabled group for the target container or an existing group was modified + // Clean the current cache in case we have built EmptyThroughputContainerController or an existing + // group was modified + this.containerControllerCache.remove(containerNameLink); + } + + return throughputControlContainerProperties; + }); + } + + public Mono processRequest(RxDocumentServiceRequest request, Mono originalRequestMono) { + checkNotNull(request, "Request can not be null"); + checkNotNull(originalRequestMono, "originalRequestMono can not be null"); + + // Currently, we will only target two resource types. + // If in the future we find other useful scenarios for throughput control, add more resource type here. + if (request.getResourceType() != ResourceType.Document && request.getResourceType() != ResourceType.StoredProcedure) { + return originalRequestMono; + } + + String collectionNameLink = Utils.getCollectionName(request.getResourceAddress()); + return this.resolveContainerController(collectionNameLink) + .flatMap(containerController -> containerController.processRequest(request, originalRequestMono)); + } + + public boolean hasDefaultGroup(String containerNameLink) { + if (this.containerMap.containsKey(containerNameLink)) { + return this.containerMap.get(containerNameLink).hasDefaultGroup(); + } + + return false; + } + + public boolean hasGroup(String containerNameLink, String throughputControlGroupName) { + if (throughputControlGroupName.isEmpty()) { + return false; + } + + if (this.containerMap.containsKey(containerNameLink)) { + return this.containerMap.get(containerNameLink).hasGroup(throughputControlGroupName); + } + + return false; + } + + private Mono resolveContainerController(String containerNameLink) { + return this.containerControllerCache.getAsync( + containerNameLink, + null, + () -> this.createAndInitContainerController(containerNameLink)); + } + + private Mono createAndInitContainerController(String containerNameLink) { + checkArgument(StringUtils.isNotEmpty(containerNameLink), "Container link should not be null or empty"); + + if (this.containerMap.containsKey(containerNameLink)) { + ContainerServerThroughputControlGroupProperties containerProperties = + this.containerMap.get(containerNameLink); + return Mono + .just(new ServerThroughputContainerController(containerProperties.getThroughputControlGroups())) + .flatMap(ServerThroughputContainerController::init); + } else { + return Mono.just(new EmptyThroughputContainerController()) + .flatMap(EmptyThroughputContainerController::init); + } + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ServerThroughputControlGroupInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ServerThroughputControlGroupInternal.java new file mode 100644 index 000000000000..d1337a10f6b4 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ServerThroughputControlGroupInternal.java @@ -0,0 +1,66 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl.server.config; + +import com.azure.cosmos.CosmosAsyncContainer; +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.models.PriorityLevel; + +import java.util.Objects; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +public abstract class ServerThroughputControlGroupInternal { + private final String groupName; + private final boolean isDefault; + private final CosmosAsyncContainer targetContainer; + private final PriorityLevel priorityLevel; + + public ServerThroughputControlGroupInternal( + String groupName, + boolean isDefault, + PriorityLevel priorityLevel, + CosmosAsyncContainer targetContainer) { + + checkArgument(StringUtils.isNotEmpty(groupName), "Argument 'groupName' cannot be null or empty."); + checkNotNull(targetContainer, "Argument 'targetContainer' can not be null"); + + this.groupName = groupName; + this.isDefault = isDefault; + this.targetContainer = targetContainer; + this.priorityLevel = priorityLevel; + } + + public String getGroupName() { + return groupName; + } + + public boolean isDefault() { + return isDefault; + } + + public CosmosAsyncContainer getTargetContainer() { + return targetContainer; + } + + public PriorityLevel getPriorityLevel() { + return priorityLevel; + } + + @Override + public boolean equals(Object o) { + if (o == null || getClass() != o.getClass()) return false; + ServerThroughputControlGroupInternal that = (ServerThroughputControlGroupInternal) o; + return isDefault == that.isDefault + && Objects.equals(groupName, that.groupName) + && Objects.equals(targetContainer.getId(), that.targetContainer.getId()) + && Objects.equals(priorityLevel, that.priorityLevel); + } + + @Override + public int hashCode() { + return Objects.hash(groupName, isDefault, targetContainer, priorityLevel); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ThroughputBucketControlGroup.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ThroughputBucketControlGroup.java new file mode 100644 index 000000000000..ced6a97151e5 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ThroughputBucketControlGroup.java @@ -0,0 +1,45 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl.server.config; + +import com.azure.cosmos.CosmosAsyncContainer; +import com.azure.cosmos.models.PriorityLevel; + +import java.util.Objects; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; + +public class ThroughputBucketControlGroup extends ServerThroughputControlGroupInternal { + private final Integer throughputBucket; + + public ThroughputBucketControlGroup( + String groupName, + boolean isDefault, + PriorityLevel priorityLevel, + Integer throughputBucket, + CosmosAsyncContainer targetContainer) { + + super(groupName, isDefault, priorityLevel, targetContainer); + + checkArgument(throughputBucket != null && throughputBucket > 0, "Target throughput should be greater than 0"); + this.throughputBucket = throughputBucket; + } + + public Integer getThroughputBucket() { + return this.throughputBucket; + } + + @Override + public boolean equals(Object o) { + if (o == null || getClass() != o.getClass()) return false; + if (!super.equals(o)) return false; + ThroughputBucketControlGroup that = (ThroughputBucketControlGroup) o; + return Objects.equals(throughputBucket, that.throughputBucket); + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), throughputBucket); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputContainerController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputContainerController.java new file mode 100644 index 000000000000..368d2a95246c --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputContainerController.java @@ -0,0 +1,108 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl.server.controller; + +import com.azure.cosmos.implementation.RxDocumentServiceRequest; +import com.azure.cosmos.implementation.Utils; +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.implementation.caches.AsyncCache; +import com.azure.cosmos.implementation.throughputControl.IThroughputContainerController; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.IThroughputController; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.server.config.ThroughputBucketControlGroup; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +import java.util.Map; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +public class ServerThroughputContainerController implements IThroughputController { + private static final Logger logger = LoggerFactory.getLogger(ServerThroughputContainerController.class); + + private final AsyncCache groupControllerCache; + private final Map groups; + + private ServerThroughputGroupControllerBase defaultGroupController; + public ServerThroughputContainerController(Map groups) { + + checkArgument(groups != null && !groups.isEmpty(), "Throughput groups can not be null or empty"); + + this.groupControllerCache = new AsyncCache<>(); + this.groups = groups; + } + + @Override + @SuppressWarnings("unchecked") + public Mono init() { + return this.createAndInitializeGroupControllers().thenReturn((T) this); + } + + @Override + public Mono processRequest(RxDocumentServiceRequest request, Mono originalRequestMono) { + checkNotNull(request, "Request can not be null"); + checkNotNull(originalRequestMono, "Original request mono can not be null"); + + return this.getOrCreateThroughputGroupController(request.getThroughputControlGroupName()) + .flatMap(groupController -> { + if (groupController.v != null) { + return groupController.v.processRequest(request, originalRequestMono); + } + + return originalRequestMono; + }); + } + + private Mono> getOrCreateThroughputGroupController(String groupName) { + + // If there is no control group defined, using the default group controller + if (StringUtils.isEmpty(groupName)) { + return Mono.just(new Utils.ValueHolder<>(this.defaultGroupController)); + } + + ServerThroughputControlGroupInternal group = this.groups.get(groupName); + if (group == null) { + // If the request is associated with a group not enabled, will fall back to the default one. + return Mono.just(new Utils.ValueHolder<>(this.defaultGroupController)); + } + + return this.resolveThroughputGroupController(group).map(Utils.ValueHolder::new); + } + + @Override + public boolean canHandleRequest(RxDocumentServiceRequest request) { + return true; + } + + private Mono createAndInitializeGroupControllers() { + return Flux.fromIterable(this.groups.values()) + .flatMap(this::resolveThroughputGroupController) + .then(Mono.just(this)); + } + + private Mono resolveThroughputGroupController(ServerThroughputControlGroupInternal group) { + return this.groupControllerCache.getAsync( + group.getGroupName(), + null, + () -> this.createAndInitializeGroupController(group)); + } + + private Mono createAndInitializeGroupController(ServerThroughputControlGroupInternal group) { + if (group instanceof ThroughputBucketControlGroup) { + // create throughput bucket group controller + ThroughputBucketGroupController throughputBucketGroupController = + new ThroughputBucketGroupController((ThroughputBucketControlGroup) group); + if (throughputBucketGroupController.isDefault()) { + this.defaultGroupController = throughputBucketGroupController; + } + return Mono.just(throughputBucketGroupController); + } + + return Mono.error( + new IllegalArgumentException("Server throughput control group type is not supported: " + group.getClass().getSimpleName())); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupControllerBase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupControllerBase.java new file mode 100644 index 000000000000..ce9948f38bf6 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupControllerBase.java @@ -0,0 +1,9 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl.server.controller; + +import com.azure.cosmos.implementation.throughputControl.sdk.controller.IThroughputController; + +public abstract class ServerThroughputGroupControllerBase implements IThroughputController { +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ThroughputBucketGroupController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ThroughputBucketGroupController.java new file mode 100644 index 000000000000..138a5b1fe949 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ThroughputBucketGroupController.java @@ -0,0 +1,44 @@ +package com.azure.cosmos.implementation.throughputControl.server.controller; + +import com.azure.cosmos.implementation.RxDocumentServiceRequest; +import com.azure.cosmos.implementation.throughputControl.server.config.ThroughputBucketControlGroup; +import reactor.core.publisher.Mono; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +public class ThroughputBucketGroupController extends ServerThroughputGroupControllerBase { + private final ThroughputBucketControlGroup throughputBucketControlGroup; + + public ThroughputBucketGroupController(ThroughputBucketControlGroup throughputBucketControlGroup) { + checkNotNull(throughputBucketControlGroup, "Argument 'throughputBucketControlGroup' cannot be null."); + this.throughputBucketControlGroup = throughputBucketControlGroup; + } + + @Override + public boolean canHandleRequest(RxDocumentServiceRequest request) { + return true; + } + + @Override + @SuppressWarnings("unchecked") + public Mono init() { + return (Mono) Mono.just(this); + } + + @Override + public Mono processRequest(RxDocumentServiceRequest request, Mono originalRequestMono) { + if (this.throughputBucketControlGroup.getPriorityLevel() != null) { + request.setPriorityLevel(this.throughputBucketControlGroup.getPriorityLevel()); + } + + if (this.throughputBucketControlGroup.getThroughputBucket() != null) { + request.setThroughputBucket(this.throughputBucketControlGroup.getThroughputBucket()); + } + + return originalRequestMono; + } + + public boolean isDefault() { + return this.throughputBucketControlGroup.isDefault(); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/module-info.java b/sdk/cosmos/azure-cosmos/src/main/java/module-info.java index 345b4a00e929..82b7f11e5056 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/module-info.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/module-info.java @@ -69,7 +69,7 @@ opens com.azure.cosmos.implementation.clienttelemetry to com.fasterxml.jackson.databind; opens com.azure.cosmos.util to com.fasterxml.jackson.databind; opens com.azure.cosmos.implementation.throughputControl to com.fasterxml.jackson.databind; - opens com.azure.cosmos.implementation.throughputControl.controller.group.global to com.fasterxml.jackson.databind; + opens com.azure.cosmos.implementation.throughputControl.sdk.controller.group.global to com.fasterxml.jackson.databind; opens com.azure.cosmos.implementation.perPartitionCircuitBreaker to com.fasterxml.jackson.databind; opens com.azure.cosmos.implementation.perPartitionAutomaticFailover to com.fasterxml.jackson.databind; @@ -79,6 +79,8 @@ exports com.azure.cosmos.implementation.routing to com.azure.cosmos.test; opens com.azure.cosmos to com.azure.cosmos.test, com.azure.spring.data.cosmos, com.fasterxml.jackson.databind, com.fasterxml.jackson.module.afterburner, java.logging; opens com.azure.cosmos.models to com.azure.cosmos.test, com.azure.spring.data.cosmos, com.fasterxml.jackson.databind, com.fasterxml.jackson.module.afterburner, java.logging; + opens com.azure.cosmos.implementation.throughputControl.sdk to com.fasterxml.jackson.databind; + opens com.azure.cosmos.implementation.throughputControl.sdk.config to com.fasterxml.jackson.databind; uses com.azure.cosmos.implementation.guava25.base.PatternCompiler; uses com.azure.core.util.tracing.Tracer; From 0cbf74dd2123671bc791efa0cf81460789214944 Mon Sep 17 00:00:00 2001 From: annie-mac Date: Wed, 16 Jul 2025 10:18:42 -0700 Subject: [PATCH 02/12] reconstruct test structure for test --- .../src/test/java/com/azure/cosmos/ClosedClientTests.java | 2 +- .../com/azure/cosmos/EndToEndTimeOutWithAvailabilityTest.java | 2 +- .../src/test/java/com/azure/cosmos/ExcludeRegionTests.java | 2 +- .../azure/cosmos/PerPartitionAutomaticFailoverE2ETests.java | 2 +- .../cosmos/SessionConsistencyWithRegionScopingTests.java | 2 +- .../faultinjection/ExcludedRegionWithFaultInjectionTests.java | 2 +- .../FaultInjectionConnectionErrorRuleTests.java | 2 +- .../FaultInjectionMetadataRequestRuleTests.java | 2 +- .../FaultInjectionServerErrorRuleOnDirectTests.java | 2 +- .../FaultInjectionServerErrorRuleOnGatewayTests.java | 2 +- .../azure/cosmos/faultinjection/FaultInjectionTestBase.java | 2 +- .../azure/cosmos/faultinjection/SessionRetryOptionsTests.java | 3 +-- .../directconnectivity/MetadataRequestRetryPolicyTests.java | 2 +- .../ContainerThroughputControlGroupPropertiesTests.java | 3 +-- .../{ => sdk}/LinkedCancellationTokenSourceTests.java | 4 +--- .../implementation/throughputControl/{ => sdk}/TestItem.java | 2 +- .../ThroughputControlGroupConfigConfigurationTests.java | 3 +-- .../throughputControl/{ => sdk}/ThroughputControlTests.java | 2 +- .../{ => sdk}/ThroughputRequestThrottlerTests.java | 4 +--- .../controller/GlobalThroughputRequestControllerTests.java | 2 +- .../{ => sdk}/controller/LocalThroughputControllerTests.java | 2 +- .../controller/PkRangesThroughputRequestControllerTests.java | 2 +- .../java/com/azure/cosmos/rx/ClientRetryPolicyE2ETests.java | 2 +- .../java/com/azure/cosmos/rx/FullFidelityChangeFeedTest.java | 2 +- .../com/azure/cosmos/rx/WebExceptionRetryPolicyE2ETests.java | 2 +- .../sdk/ContainerSDKThroughputControlGroupProperties.java | 1 - 26 files changed, 25 insertions(+), 33 deletions(-) rename sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/ContainerThroughputControlGroupPropertiesTests.java (97%) rename sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/LinkedCancellationTokenSourceTests.java (81%) rename sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/TestItem.java (93%) rename sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/ThroughputControlGroupConfigConfigurationTests.java (97%) rename sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/ThroughputControlTests.java (99%) rename sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/ThroughputRequestThrottlerTests.java (97%) rename sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/controller/GlobalThroughputRequestControllerTests.java (98%) rename sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/controller/LocalThroughputControllerTests.java (98%) rename sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/{ => sdk}/controller/PkRangesThroughputRequestControllerTests.java (98%) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ClosedClientTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ClosedClientTests.java index 3949f6cafcc3..0853bcfa0a1f 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ClosedClientTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ClosedClientTests.java @@ -9,7 +9,7 @@ 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.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosItemRequestOptions; import com.azure.cosmos.models.CosmosPatchOperations; import com.azure.cosmos.models.CosmosQueryRequestOptions; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/EndToEndTimeOutWithAvailabilityTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/EndToEndTimeOutWithAvailabilityTest.java index b73bf1d39759..6219dae1cb3b 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/EndToEndTimeOutWithAvailabilityTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/EndToEndTimeOutWithAvailabilityTest.java @@ -9,7 +9,7 @@ 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.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosItemIdentity; import com.azure.cosmos.models.CosmosItemRequestOptions; import com.azure.cosmos.models.CosmosItemResponse; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ExcludeRegionTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ExcludeRegionTests.java index 03aa0deb98e4..67f255617be0 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ExcludeRegionTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ExcludeRegionTests.java @@ -10,7 +10,7 @@ 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.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosBatch; import com.azure.cosmos.models.CosmosBatchRequestOptions; import com.azure.cosmos.models.CosmosBatchResponse; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PerPartitionAutomaticFailoverE2ETests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PerPartitionAutomaticFailoverE2ETests.java index 4bd5006e53ba..a78618e5ea36 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PerPartitionAutomaticFailoverE2ETests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PerPartitionAutomaticFailoverE2ETests.java @@ -37,7 +37,7 @@ import com.azure.cosmos.implementation.http.HttpResponse; import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper; import com.azure.cosmos.implementation.routing.RegionalRoutingContext; -import com.azure.cosmos.implementation.throughputControl.TestItem; +import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosBatch; import com.azure.cosmos.models.CosmosBatchResponse; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/SessionConsistencyWithRegionScopingTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/SessionConsistencyWithRegionScopingTests.java index 1b913ccf6acd..fc666f46efba 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/SessionConsistencyWithRegionScopingTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/SessionConsistencyWithRegionScopingTests.java @@ -24,7 +24,7 @@ import com.azure.cosmos.implementation.guava25.hash.Funnels; import com.azure.cosmos.implementation.guava25.hash.PrimitiveSink; import com.azure.cosmos.implementation.routing.PartitionKeyInternal; -import com.azure.cosmos.implementation.throughputControl.TestItem; +import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosBatch; import com.azure.cosmos.models.CosmosBatchOperationResult; import com.azure.cosmos.models.CosmosBatchRequestOptions; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/ExcludedRegionWithFaultInjectionTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/ExcludedRegionWithFaultInjectionTests.java index 6ddd63647243..2038baa9d8d0 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/ExcludedRegionWithFaultInjectionTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/ExcludedRegionWithFaultInjectionTests.java @@ -19,7 +19,7 @@ import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.TestConfigurations; import com.azure.cosmos.implementation.Utils; -import com.azure.cosmos.implementation.throughputControl.TestItem; +import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosBatch; import com.azure.cosmos.models.CosmosBatchRequestOptions; import com.azure.cosmos.models.CosmosBatchResponse; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionConnectionErrorRuleTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionConnectionErrorRuleTests.java index c9694190c34f..dc45e3134c68 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionConnectionErrorRuleTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionConnectionErrorRuleTests.java @@ -16,7 +16,7 @@ import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils; import com.azure.cosmos.implementation.directconnectivity.RntbdTransportClient; import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdEndpoint; -import com.azure.cosmos.implementation.throughputControl.TestItem; +import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosContainerIdentity; import com.azure.cosmos.models.FeedRange; import com.azure.cosmos.models.PartitionKey; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionMetadataRequestRuleTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionMetadataRequestRuleTests.java index 38ed45b785bd..39c1828e92f4 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionMetadataRequestRuleTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionMetadataRequestRuleTests.java @@ -18,7 +18,7 @@ import com.azure.cosmos.implementation.MetadataDiagnosticsContext; import com.azure.cosmos.implementation.OperationType; import com.azure.cosmos.implementation.Utils; -import com.azure.cosmos.implementation.throughputControl.TestItem; +import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.FeedRange; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java index f3496fa84d62..899b262dc168 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java @@ -20,7 +20,7 @@ import com.azure.cosmos.implementation.ResourceType; import com.azure.cosmos.implementation.TestConfigurations; import com.azure.cosmos.implementation.Utils; -import com.azure.cosmos.implementation.throughputControl.TestItem; +import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.FeedRange; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnGatewayTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnGatewayTests.java index 8c82c7966592..1ef55aa5fa33 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnGatewayTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnGatewayTests.java @@ -19,7 +19,7 @@ import com.azure.cosmos.implementation.ResourceType; import com.azure.cosmos.implementation.TestConfigurations; import com.azure.cosmos.implementation.Utils; -import com.azure.cosmos.implementation.throughputControl.TestItem; +import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.FeedRange; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionTestBase.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionTestBase.java index 22164aee953f..0996726cedde 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionTestBase.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionTestBase.java @@ -8,7 +8,7 @@ import com.azure.cosmos.CosmosDiagnostics; import com.azure.cosmos.CosmosException; import com.azure.cosmos.implementation.OperationType; -import com.azure.cosmos.implementation.throughputControl.TestItem; +import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; import com.azure.cosmos.models.CosmosPatchOperations; import com.azure.cosmos.models.CosmosQueryRequestOptions; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/SessionRetryOptionsTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/SessionRetryOptionsTests.java index cab6bfa64a90..10d2b930c35f 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/SessionRetryOptionsTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/SessionRetryOptionsTests.java @@ -20,7 +20,7 @@ import com.azure.cosmos.implementation.ImplementationBridgeHelpers; import com.azure.cosmos.implementation.OperationType; import com.azure.cosmos.implementation.TestConfigurations; -import com.azure.cosmos.implementation.throughputControl.TestItem; +import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosItemRequestOptions; import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.CosmosPatchOperations; @@ -55,7 +55,6 @@ import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicReference; -import java.util.stream.Collectors; import static org.assertj.core.api.AssertionsForClassTypes.assertThat; import static org.testng.AssertJUnit.fail; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/MetadataRequestRetryPolicyTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/MetadataRequestRetryPolicyTests.java index 13b8cb738289..679d4444c6cb 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/MetadataRequestRetryPolicyTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/MetadataRequestRetryPolicyTests.java @@ -33,7 +33,7 @@ import com.azure.cosmos.implementation.http.HttpClientConfig; import com.azure.cosmos.implementation.http.HttpTimeoutPolicyControlPlaneHotPath; import com.azure.cosmos.implementation.routing.RegionalRoutingContext; -import com.azure.cosmos.implementation.throughputControl.TestItem; +import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosItemRequestOptions; import com.azure.cosmos.models.CosmosPatchOperations; import com.azure.cosmos.models.CosmosQueryRequestOptions; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ContainerThroughputControlGroupPropertiesTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerThroughputControlGroupPropertiesTests.java similarity index 97% rename from sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ContainerThroughputControlGroupPropertiesTests.java rename to sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerThroughputControlGroupPropertiesTests.java index 5fde672ed95d..592e01afb58e 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ContainerThroughputControlGroupPropertiesTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerThroughputControlGroupPropertiesTests.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl; +package com.azure.cosmos.implementation.throughputControl.sdk; import com.azure.cosmos.CosmosAsyncClient; import com.azure.cosmos.CosmosAsyncContainer; @@ -9,7 +9,6 @@ import com.azure.cosmos.implementation.RxDocumentServiceRequest; import com.azure.cosmos.implementation.TestConfigurations; import com.azure.cosmos.implementation.apachecommons.lang.tuple.Pair; -import com.azure.cosmos.implementation.throughputControl.sdk.ContainerSDKThroughputControlGroupProperties; import com.azure.cosmos.implementation.throughputControl.sdk.config.LocalThroughputControlGroup; import com.azure.cosmos.models.PriorityLevel; import org.mockito.Mockito; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationTokenSourceTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/LinkedCancellationTokenSourceTests.java similarity index 81% rename from sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationTokenSourceTests.java rename to sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/LinkedCancellationTokenSourceTests.java index b0eb8eae81f8..8a1cc2686a0c 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationTokenSourceTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/LinkedCancellationTokenSourceTests.java @@ -1,10 +1,8 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl; +package com.azure.cosmos.implementation.throughputControl.sdk; -import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationToken; -import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationTokenSource; import org.testng.annotations.Test; import static org.assertj.core.api.Assertions.assertThat; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/TestItem.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/TestItem.java similarity index 93% rename from sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/TestItem.java rename to sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/TestItem.java index 8e03339f67c3..cc9b9422e903 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/TestItem.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/TestItem.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl; +package com.azure.cosmos.implementation.throughputControl.sdk; import java.util.UUID; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupConfigConfigurationTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputControlGroupConfigConfigurationTests.java similarity index 97% rename from sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupConfigConfigurationTests.java rename to sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputControlGroupConfigConfigurationTests.java index 597286a9883f..26e0de3144b0 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupConfigConfigurationTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputControlGroupConfigConfigurationTests.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl; +package com.azure.cosmos.implementation.throughputControl.sdk; import com.azure.cosmos.CosmosAsyncClient; import com.azure.cosmos.CosmosAsyncContainer; @@ -19,7 +19,6 @@ import java.time.Duration; import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.assertThatExceptionOfType; import static org.assertj.core.api.Assertions.assertThatThrownBy; public class ThroughputControlGroupConfigConfigurationTests extends TestSuiteBase { diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputControlTests.java similarity index 99% rename from sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java rename to sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputControlTests.java index e0e568bd2568..be31747b7291 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputControlTests.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl; +package com.azure.cosmos.implementation.throughputControl.sdk; import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.ConnectionMode; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputRequestThrottlerTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputRequestThrottlerTests.java similarity index 97% rename from sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputRequestThrottlerTests.java rename to sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputRequestThrottlerTests.java index 3b026b3bd773..21e3a0f06f34 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputRequestThrottlerTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputRequestThrottlerTests.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl; +package com.azure.cosmos.implementation.throughputControl.sdk; import com.azure.cosmos.implementation.DocumentServiceRequestContext; import com.azure.cosmos.implementation.HttpConstants; @@ -13,8 +13,6 @@ import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils; import com.azure.cosmos.implementation.directconnectivity.StoreResponse; -import com.azure.cosmos.implementation.throughputControl.sdk.ThroughputControlTrackingUnit; -import com.azure.cosmos.implementation.throughputControl.sdk.ThroughputRequestThrottler; import org.mockito.Mockito; import org.testng.annotations.Test; import reactor.test.StepVerifier; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/GlobalThroughputRequestControllerTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/GlobalThroughputRequestControllerTests.java similarity index 98% rename from sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/GlobalThroughputRequestControllerTests.java rename to sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/GlobalThroughputRequestControllerTests.java index b3214e88626f..364c83499dd3 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/GlobalThroughputRequestControllerTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/GlobalThroughputRequestControllerTests.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller; +package com.azure.cosmos.implementation.throughputControl.sdk.controller; import com.azure.cosmos.implementation.DocumentServiceRequestContext; import com.azure.cosmos.implementation.OperationType; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/LocalThroughputControllerTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/LocalThroughputControllerTests.java similarity index 98% rename from sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/LocalThroughputControllerTests.java rename to sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/LocalThroughputControllerTests.java index 3bb45052405d..60202d2622a8 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/LocalThroughputControllerTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/LocalThroughputControllerTests.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller; +package com.azure.cosmos.implementation.throughputControl.sdk.controller; import com.azure.cosmos.ConnectionMode; import com.azure.cosmos.CosmosAsyncContainer; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/PkRangesThroughputRequestControllerTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/PkRangesThroughputRequestControllerTests.java similarity index 98% rename from sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/PkRangesThroughputRequestControllerTests.java rename to sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/PkRangesThroughputRequestControllerTests.java index 8f392e8f1230..fcc8f6c6bba2 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/PkRangesThroughputRequestControllerTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/PkRangesThroughputRequestControllerTests.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller; +package com.azure.cosmos.implementation.throughputControl.sdk.controller; import com.azure.cosmos.implementation.DocumentServiceRequestContext; import com.azure.cosmos.implementation.OperationType; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ClientRetryPolicyE2ETests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ClientRetryPolicyE2ETests.java index 75e7fab8a206..486dd3e1d2c2 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ClientRetryPolicyE2ETests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ClientRetryPolicyE2ETests.java @@ -20,7 +20,7 @@ import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.OperationType; import com.azure.cosmos.implementation.Utils; -import com.azure.cosmos.implementation.throughputControl.TestItem; +import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.CosmosPatchOperations; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/FullFidelityChangeFeedTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/FullFidelityChangeFeedTest.java index 2e43dcdfbb26..4b7eb6ea612e 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/FullFidelityChangeFeedTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/FullFidelityChangeFeedTest.java @@ -6,7 +6,7 @@ import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.CosmosAsyncDatabase; import com.azure.cosmos.CosmosClientBuilder; -import com.azure.cosmos.implementation.throughputControl.TestItem; +import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.ChangeFeedPolicy; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; import com.azure.cosmos.models.CosmosContainerProperties; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/WebExceptionRetryPolicyE2ETests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/WebExceptionRetryPolicyE2ETests.java index bc544fbc855d..8ada8f9df8db 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/WebExceptionRetryPolicyE2ETests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/WebExceptionRetryPolicyE2ETests.java @@ -16,7 +16,7 @@ import com.azure.cosmos.implementation.ImplementationBridgeHelpers; import com.azure.cosmos.implementation.OperationType; import com.azure.cosmos.implementation.RequestTimeline; -import com.azure.cosmos.implementation.throughputControl.TestItem; +import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerSDKThroughputControlGroupProperties.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerSDKThroughputControlGroupProperties.java index a675a723cd92..05c3180e5798 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerSDKThroughputControlGroupProperties.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerSDKThroughputControlGroupProperties.java @@ -11,7 +11,6 @@ import org.slf4j.LoggerFactory; import reactor.core.publisher.Mono; -import java.util.HashSet; import java.util.Map; import java.util.Objects; import java.util.Set; From eb72c4316b6095557f6ab1768c11c626897b2910 Mon Sep 17 00:00:00 2001 From: annie-mac Date: Wed, 16 Jul 2025 11:59:10 -0700 Subject: [PATCH 03/12] refactor --- .../com/azure/cosmos/ClosedClientTests.java | 11 ++- .../EndToEndTimeOutWithAvailabilityTest.java | 39 +++++----- .../com/azure/cosmos/ExcludeRegionTests.java | 33 +++++---- ...PerPartitionAutomaticFailoverE2ETests.java | 35 +++++---- ...sionConsistencyWithRegionScopingTests.java | 16 ++--- ...ExcludedRegionWithFaultInjectionTests.java | 52 +++++++------- ...aultInjectionConnectionErrorRuleTests.java | 14 ++-- ...aultInjectionMetadataRequestRuleTests.java | 22 +++--- ...InjectionServerErrorRuleOnDirectTests.java | 42 +++++------ ...njectionServerErrorRuleOnGatewayTests.java | 26 +++---- .../FaultInjectionTestBase.java | 20 +++--- .../SessionRetryOptionsTests.java | 26 +++---- .../MetadataRequestRetryPolicyTests.java | 12 ++-- .../cosmos/rx/ClientRetryPolicyE2ETests.java | 44 ++++++------ .../cosmos/rx/FullFidelityChangeFeedTest.java | 72 ++++++++++--------- .../azure/cosmos/CosmosAsyncContainer.java | 3 +- .../implementation/AsyncDocumentClient.java | 4 +- .../implementation/RxDocumentClientImpl.java | 1 - .../IThroughputContainerController.java | 4 +- .../IThroughputController.java | 2 +- ...erSDKThroughputControlGroupProperties.java | 4 -- .../sdk/SDKThroughputControlStore.java | 29 +------- .../SDKThroughputControlGroupInternal.java | 1 - .../SDKThroughputContainerController.java | 28 ++++---- .../group/ThroughputGroupControllerBase.java | 2 +- .../request/IThroughputRequestController.java | 2 +- ...erverThroughputControlGroupProperties.java | 6 +- .../server/ServerThroughputControlStore.java | 6 +- .../ServerThroughputContainerController.java | 3 +- .../ServerThroughputGroupControllerBase.java | 2 +- 30 files changed, 262 insertions(+), 299 deletions(-) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{sdk/controller => }/IThroughputController.java (94%) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ClosedClientTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ClosedClientTests.java index 0853bcfa0a1f..ade9497e76b1 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ClosedClientTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ClosedClientTests.java @@ -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.sdk.TestItem; import com.azure.cosmos.models.CosmosItemRequestOptions; import com.azure.cosmos.models.CosmosPatchOperations; import com.azure.cosmos.models.CosmosQueryRequestOptions; @@ -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) { @@ -121,7 +120,7 @@ public void failFastWhenUsingClosedClient(OperationType operationType) { private void performDocumentOperation( CosmosAsyncContainer asyncContainer, - TestItem testItem, + TestObject testItem, OperationType operationType) { if (operationType == OperationType.Query) { @@ -131,7 +130,7 @@ private void performDocumentOperation( SqlQuerySpec sqlQuerySpec = new SqlQuerySpec(query); asyncContainer - .queryItems(sqlQuerySpec, queryRequestOptions, TestItem.class) + .queryItems(sqlQuerySpec, queryRequestOptions, TestObject.class) .byPage() .blockFirst(); } @@ -139,7 +138,7 @@ private void performDocumentOperation( if (operationType == OperationType.Read) { asyncContainer - .readItem(testItem.getId(), new PartitionKey(testItem.getId()), TestItem.class) + .readItem(testItem.getId(), new PartitionKey(testItem.getId()), TestObject.class) .block(); } @@ -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(); } } diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/EndToEndTimeOutWithAvailabilityTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/EndToEndTimeOutWithAvailabilityTest.java index 6219dae1cb3b..5d82877afe68 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/EndToEndTimeOutWithAvailabilityTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/EndToEndTimeOutWithAvailabilityTest.java @@ -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.sdk.TestItem; import com.azure.cosmos.models.CosmosItemIdentity; import com.azure.cosmos.models.CosmosItemRequestOptions; import com.azure.cosmos.models.CosmosItemResponse; @@ -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(); @@ -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(); @@ -306,7 +305,7 @@ public void afterClass() { private CosmosDiagnostics performDocumentOperation( CosmosAsyncContainer cosmosAsyncContainer, OperationType operationType, - TestItem createdItem, + TestObject createdItem, CosmosItemRequestOptions cosmosItemRequestOptions, boolean ignoreE2E2LatencyCfgOnRequestOptions, QueryFlavor queryFlavor) { @@ -330,8 +329,8 @@ private CosmosDiagnostics performDocumentOperation( logger.info("Running readAllItems..."); - FeedResponse response = cosmosAsyncContainer - .readAllItems(queryRequestOptions, TestItem.class) + FeedResponse response = cosmosAsyncContainer + .readAllItems(queryRequestOptions, TestObject.class) .byPage() .blockFirst(); @@ -347,11 +346,11 @@ private CosmosDiagnostics performDocumentOperation( logger.info("Running readMany..."); - FeedResponse response = cosmosAsyncContainer + FeedResponse response = cosmosAsyncContainer .readMany( Arrays.asList(new CosmosItemIdentity(new PartitionKey(createdItem.getMypk()), createdItem.getId())), readManyRequestOptions, - TestItem.class) + TestObject.class) .block(); assertThat(response).isNotNull(); @@ -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 itemFeedResponse = - cosmosAsyncContainer.queryItems(query, queryRequestOptions, TestItem.class).byPage().blockFirst(); + FeedResponse itemFeedResponse = + cosmosAsyncContainer.queryItems(query, queryRequestOptions, TestObject.class).byPage().blockFirst(); assertThat(itemFeedResponse).isNotNull(); @@ -378,11 +377,11 @@ private CosmosDiagnostics performDocumentOperation( if (operationType == OperationType.Read) { - CosmosItemResponse response = cosmosAsyncContainer.readItem( + CosmosItemResponse response = cosmosAsyncContainer.readItem( createdItem.getId(), new PartitionKey(createdItem.getMypk()), cosmosItemRequestOptions, - TestItem.class) + TestObject.class) .block(); assertThat(response).isNotNull(); @@ -391,7 +390,7 @@ private CosmosDiagnostics performDocumentOperation( } if (operationType == OperationType.Replace) { - CosmosItemResponse response = cosmosAsyncContainer.replaceItem( + CosmosItemResponse response = cosmosAsyncContainer.replaceItem( createdItem, createdItem.getId(), new PartitionKey(createdItem.getMypk()), @@ -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 response = cosmosAsyncContainer .deleteItem(toBeDeletedItem, cosmosItemRequestOptions) @@ -416,8 +415,8 @@ private CosmosDiagnostics performDocumentOperation( } if (operationType == OperationType.Create) { - CosmosItemResponse response = cosmosAsyncContainer - .createItem(TestItem.createNewItem(), cosmosItemRequestOptions) + CosmosItemResponse response = cosmosAsyncContainer + .createItem(TestObject.create(), cosmosItemRequestOptions) .block(); assertThat(response).isNotNull(); @@ -426,8 +425,8 @@ private CosmosDiagnostics performDocumentOperation( } if (operationType == OperationType.Upsert) { - CosmosItemResponse response = cosmosAsyncContainer - .upsertItem(TestItem.createNewItem(), cosmosItemRequestOptions) + CosmosItemResponse response = cosmosAsyncContainer + .upsertItem(TestObject.create(), cosmosItemRequestOptions) .block(); assertThat(response).isNotNull(); @@ -444,8 +443,8 @@ private CosmosDiagnostics performDocumentOperation( CosmosPatchItemRequestOptions patchItemRequestOptions = new CosmosPatchItemRequestOptions(); patchItemRequestOptions.setNonIdempotentWriteRetryPolicy(true, true); patchItemRequestOptions.setCosmosEndToEndOperationLatencyPolicyConfig(config); - CosmosItemResponse response = cosmosAsyncContainer - .patchItem(createdItem.getId(), new PartitionKey(createdItem.getMypk()), patchOperations, patchItemRequestOptions, TestItem.class) + CosmosItemResponse response = cosmosAsyncContainer + .patchItem(createdItem.getId(), new PartitionKey(createdItem.getMypk()), patchOperations, patchItemRequestOptions, TestObject.class) .block(); assertThat(response).isNotNull(); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ExcludeRegionTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ExcludeRegionTests.java index 67f255617be0..e0e91464d82d 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ExcludeRegionTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ExcludeRegionTests.java @@ -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.sdk.TestItem; import com.azure.cosmos.models.CosmosBatch; import com.azure.cosmos.models.CosmosBatchRequestOptions; import com.azure.cosmos.models.CosmosBatchResponse; @@ -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); @@ -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) @@ -229,7 +228,7 @@ private List getPreferredRegionList(CosmosAsyncClient client) { private CosmosDiagnosticsContext performDocumentOperation( CosmosAsyncContainer cosmosAsyncContainer, OperationType operationType, - TestItem createdItem, + TestObject createdItem, List excludeRegions, CosmosEndToEndOperationLatencyPolicyConfig cosmosEndToEndOperationLatencyPolicyConfig) { @@ -240,8 +239,8 @@ private CosmosDiagnosticsContext performDocumentOperation( String query = String.format("SELECT * from c where c.id = '%s'", createdItem.getId()); queryRequestOptions.setExcludedRegions(excludeRegions); - FeedResponse itemFeedResponse = - cosmosAsyncContainer.queryItems(query, queryRequestOptions, TestItem.class).byPage().blockFirst(); + FeedResponse itemFeedResponse = + cosmosAsyncContainer.queryItems(query, queryRequestOptions, TestObject.class).byPage().blockFirst(); assertThat(itemFeedResponse).isNotNull(); CosmosDiagnostics cosmosDiagnostics = itemFeedResponse.getCosmosDiagnostics(); @@ -271,11 +270,11 @@ private CosmosDiagnosticsContext performDocumentOperation( throw new RuntimeException(e); } - CosmosItemResponse itemResponse = cosmosAsyncContainer.readItem( + CosmosItemResponse itemResponse = cosmosAsyncContainer.readItem( createdItem.getId(), new PartitionKey(createdItem.getMypk()), cosmosItemRequestOptions, - TestItem.class).block(); + TestObject.class).block(); assertThat(itemResponse).isNotNull(); assertThat(itemResponse.getDiagnostics()).isNotNull(); @@ -295,7 +294,7 @@ private CosmosDiagnosticsContext performDocumentOperation( throw new RuntimeException(e); } - CosmosItemResponse itemResponse = cosmosAsyncContainer.replaceItem( + CosmosItemResponse itemResponse = cosmosAsyncContainer.replaceItem( createdItem, createdItem.getId(), new PartitionKey(createdItem.getMypk()), @@ -313,7 +312,7 @@ private CosmosDiagnosticsContext performDocumentOperation( if (operationType == OperationType.Delete) { - TestItem itemToBeDeleted = TestItem.createNewItem(); + TestObject itemToBeDeleted = TestObject.create(); cosmosAsyncContainer.createItem(itemToBeDeleted, cosmosItemRequestOptions).block(); @@ -337,8 +336,8 @@ private CosmosDiagnosticsContext performDocumentOperation( } if (operationType == OperationType.Create) { - CosmosItemResponse itemResponse = cosmosAsyncContainer - .createItem(TestItem.createNewItem(), cosmosItemRequestOptions).block(); + CosmosItemResponse itemResponse = cosmosAsyncContainer + .createItem(TestObject.create(), cosmosItemRequestOptions).block(); assertThat(itemResponse).isNotNull(); assertThat(itemResponse.getDiagnostics()).isNotNull(); @@ -351,8 +350,8 @@ private CosmosDiagnosticsContext performDocumentOperation( } if (operationType == OperationType.Upsert) { - CosmosItemResponse itemResponse - = cosmosAsyncContainer.upsertItem(TestItem.createNewItem(), cosmosItemRequestOptions).block(); + CosmosItemResponse itemResponse + = cosmosAsyncContainer.upsertItem(TestObject.create(), cosmosItemRequestOptions).block(); assertThat(itemResponse).isNotNull(); assertThat(itemResponse.getDiagnostics()).isNotNull(); @@ -375,8 +374,8 @@ private CosmosDiagnosticsContext performDocumentOperation( patchItemRequestOptions.setCosmosEndToEndOperationLatencyPolicyConfig(cosmosEndToEndOperationLatencyPolicyConfig); patchItemRequestOptions.setExcludedRegions(excludeRegions); - CosmosItemResponse itemResponse = cosmosAsyncContainer - .patchItem(createdItem.getId(), new PartitionKey(createdItem.getMypk()), patchOperations, patchItemRequestOptions, TestItem.class) + CosmosItemResponse itemResponse = cosmosAsyncContainer + .patchItem(createdItem.getId(), new PartitionKey(createdItem.getMypk()), patchOperations, patchItemRequestOptions, TestObject.class) .block(); assertThat(itemResponse).isNotNull(); @@ -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); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PerPartitionAutomaticFailoverE2ETests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PerPartitionAutomaticFailoverE2ETests.java index a78618e5ea36..e53658ae0dde 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PerPartitionAutomaticFailoverE2ETests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PerPartitionAutomaticFailoverE2ETests.java @@ -37,7 +37,6 @@ import com.azure.cosmos.implementation.http.HttpResponse; import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper; import com.azure.cosmos.implementation.routing.RegionalRoutingContext; -import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosBatch; import com.azure.cosmos.models.CosmosBatchResponse; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; @@ -1015,7 +1014,7 @@ public void testPpafWithWriteFailoverWithEligibleErrorStatusCodes( regionalRoutingContextWithIssues, cosmosException); - TestItem testItem = TestItem.createNewItem(); + TestObject testItem = TestObject.create(); Function> dataPlaneOperation = resolveDataPlaneOperation(operationType); @@ -1109,7 +1108,7 @@ public void testPpafWithWriteFailoverWithEligibleErrorStatusCodes( shouldThrowReadTimeoutExceptionWhenNetworkError, shouldUseE2ETimeout); - TestItem testItem = TestItem.createNewItem(); + TestObject testItem = TestObject.create(); Function> dataPlaneOperation = resolveDataPlaneOperation(operationType); @@ -1378,7 +1377,7 @@ private Function> resolveDa return (paramsWrapper) -> { CosmosAsyncContainer asyncContainer = paramsWrapper.asyncContainer; - TestItem createdTestObject = paramsWrapper.createdTestItem; + TestObject createdTestObject = paramsWrapper.createdTestItem; CosmosItemRequestOptions itemRequestOptions = paramsWrapper.itemRequestOptions; try { @@ -1405,12 +1404,12 @@ private Function> resolveDa return (paramsWrapper) -> { CosmosAsyncContainer asyncContainer = paramsWrapper.asyncContainer; - TestItem createdTestObject = paramsWrapper.createdTestItem; + TestObject createdTestObject = paramsWrapper.createdTestItem; CosmosItemRequestOptions itemRequestOptions = paramsWrapper.itemRequestOptions; try { - CosmosItemResponse upsertItemResponse = asyncContainer.upsertItem( + CosmosItemResponse upsertItemResponse = asyncContainer.upsertItem( createdTestObject, new PartitionKey(createdTestObject.getId()), itemRequestOptions) @@ -1431,12 +1430,12 @@ private Function> resolveDa return (paramsWrapper) -> { CosmosAsyncContainer asyncContainer = paramsWrapper.asyncContainer; - TestItem createdTestObject = TestItem.createNewItem(); + TestObject createdTestObject = TestObject.create(); CosmosItemRequestOptions itemRequestOptions = paramsWrapper.itemRequestOptions; try { - CosmosItemResponse createItemResponse = asyncContainer.createItem( + CosmosItemResponse createItemResponse = asyncContainer.createItem( createdTestObject, new PartitionKey(createdTestObject.getId()), itemRequestOptions) @@ -1457,7 +1456,7 @@ private Function> resolveDa return (paramsWrapper) -> { CosmosAsyncContainer asyncContainer = paramsWrapper.asyncContainer; - TestItem createdTestObject = paramsWrapper.createdTestItem; + TestObject createdTestObject = paramsWrapper.createdTestItem; CosmosItemRequestOptions itemRequestOptions = paramsWrapper.itemRequestOptions; try { @@ -1483,19 +1482,19 @@ private Function> resolveDa return (paramsWrapper) -> { CosmosAsyncContainer asyncContainer = paramsWrapper.asyncContainer; - TestItem createdTestObject = paramsWrapper.createdTestItem; + TestObject createdTestObject = paramsWrapper.createdTestItem; CosmosPatchItemRequestOptions patchItemRequestOptions = (CosmosPatchItemRequestOptions) paramsWrapper.patchItemRequestOptions; CosmosPatchOperations patchOperations = CosmosPatchOperations.create().add("/number", 555); try { - CosmosItemResponse patchItemResponse = asyncContainer.patchItem( + CosmosItemResponse patchItemResponse = asyncContainer.patchItem( createdTestObject.getId(), new PartitionKey(createdTestObject.getId()), patchOperations, patchItemRequestOptions, - TestItem.class) + TestObject.class) .block(); return new ResponseWrapper<>(patchItemResponse); @@ -1540,12 +1539,12 @@ private Function> resolveDa return (paramsWrapper) -> { CosmosAsyncContainer asyncContainer = paramsWrapper.asyncContainer; - TestItem createdTestObject = paramsWrapper.createdTestItem; + TestObject createdTestObject = paramsWrapper.createdTestItem; CosmosItemRequestOptions itemRequestOptions = paramsWrapper.itemRequestOptions; try { - CosmosItemResponse deleteItemResponse = asyncContainer.replaceItem( + CosmosItemResponse deleteItemResponse = asyncContainer.replaceItem( createdTestObject, createdTestObject.getId(), new PartitionKey(createdTestObject.getId()), @@ -1566,7 +1565,7 @@ private Function> resolveDa case Batch: return (paramsWrapper) -> { - TestItem testObject = TestItem.createNewItem(); + TestObject testObject = TestObject.create(); CosmosBatch batch = CosmosBatch.createCosmosBatch(new PartitionKey(testObject.getId())); CosmosAsyncContainer asyncContainer = paramsWrapper.asyncContainer; @@ -1591,9 +1590,9 @@ private Function> resolveDa try { - FeedResponse feedResponseFromChangeFeed = asyncContainer.queryChangeFeed( + FeedResponse feedResponseFromChangeFeed = asyncContainer.queryChangeFeed( CosmosChangeFeedRequestOptions.createForProcessingFromBeginning(paramsWrapper.feedRangeToDrainForChangeFeed == null ? FeedRange.forFullRange() : paramsWrapper.feedRangeToDrainForChangeFeed), - TestItem.class) + TestObject.class) .byPage() .blockLast(); @@ -1651,7 +1650,7 @@ private static class ResponseWrapper { private static class OperationInvocationParamsWrapper { public CosmosAsyncContainer asyncContainer; - public TestItem createdTestItem; + public TestObject createdTestItem; public CosmosItemRequestOptions itemRequestOptions; public CosmosQueryRequestOptions queryRequestOptions; public CosmosItemRequestOptions patchItemRequestOptions; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/SessionConsistencyWithRegionScopingTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/SessionConsistencyWithRegionScopingTests.java index fc666f46efba..0ac72022717b 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/SessionConsistencyWithRegionScopingTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/SessionConsistencyWithRegionScopingTests.java @@ -24,7 +24,6 @@ import com.azure.cosmos.implementation.guava25.hash.Funnels; import com.azure.cosmos.implementation.guava25.hash.PrimitiveSink; import com.azure.cosmos.implementation.routing.PartitionKeyInternal; -import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosBatch; import com.azure.cosmos.models.CosmosBatchOperationResult; import com.azure.cosmos.models.CosmosBatchRequestOptions; @@ -62,6 +61,7 @@ import reactor.core.publisher.Mono; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; @@ -322,7 +322,7 @@ public Object[][] readYouWriteWithNoExplicitRegionSwitchingTestContext() { Flux createOperationsFlux = Flux.range(0, createOperationCount).map(i -> { String documentId = UUID.randomUUID().toString(); - TestItem testItem = new TestItem(documentId, documentId, documentId); + TestObject testItem = new TestObject(documentId, documentId, Arrays.asList(), documentId); idsToCreate.add(documentId); @@ -353,7 +353,7 @@ public Object[][] readYouWriteWithNoExplicitRegionSwitchingTestContext() { bulkReadResponses.forEach(bulkReadResponse -> { assertThat(bulkReadResponse.getResponse()).isNotNull(); assertThat(bulkReadResponse.getResponse().getStatusCode()).isEqualTo(HttpConstants.StatusCodes.OK); - pksRead.add(bulkReadResponse.getResponse().getItem(TestItem.class).getMypk()); + pksRead.add(bulkReadResponse.getResponse().getItem(TestObject.class).getMypk()); }); return pksRead; @@ -366,7 +366,7 @@ public Object[][] readYouWriteWithNoExplicitRegionSwitchingTestContext() { Flux createOperationsFlux = Flux.range(0, createOperationCount).map(i -> { String documentId = UUID.randomUUID().toString(); - TestItem testItem = new TestItem(documentId, documentId, documentId); + TestObject testItem = new TestObject(documentId, documentId, Arrays.asList(), documentId); idsAddedByBulkCreate.add(documentId); return CosmosBulkOperations.getCreateItemOperation(testItem, new PartitionKey(documentId)); @@ -444,7 +444,7 @@ public Object[][] readYouWriteWithNoExplicitRegionSwitchingTestContext() { Flux createOperationsFlux = Flux.range(0, createOperationCount).map(i -> { String documentId = UUID.randomUUID().toString(); - TestItem testItem = new TestItem(documentId, documentId, documentId); + TestObject testItem = new TestObject(documentId, documentId, Arrays.asList(), documentId); idsAddedByBulkCreate.add(documentId); return CosmosBulkOperations.getCreateItemOperation(testItem, new PartitionKey(documentId)); @@ -1191,7 +1191,7 @@ public Object[][] readYouWriteWithExplicitRegionSwitchingTestContext() { Flux createOperationsFlux = Flux.range(0, createOperationCount).map(i -> { String documentId = UUID.randomUUID().toString(); - TestItem testItem = new TestItem(documentId, documentId, documentId); + TestObject testItem = new TestObject(documentId, documentId, Arrays.asList(), documentId); idsToCreate.add(documentId); @@ -1247,7 +1247,7 @@ public Object[][] readYouWriteWithExplicitRegionSwitchingTestContext() { Flux createOperationsFlux = Flux.range(0, createOperationCount).map(i -> { String documentId = UUID.randomUUID().toString(); - TestItem testItem = new TestItem(documentId, documentId, documentId); + TestObject testItem = new TestObject(documentId, documentId, Arrays.asList(), documentId); idsAddedByCreates.add(documentId); return CosmosBulkOperations.getCreateItemOperation(testItem, new PartitionKey(documentId)); @@ -1320,7 +1320,7 @@ public Object[][] readYouWriteWithExplicitRegionSwitchingTestContext() { Flux createOperationsFlux = Flux.range(0, createOperationCount).map(i -> { String documentId = UUID.randomUUID().toString(); - TestItem testItem = new TestItem(documentId, documentId, documentId); + TestObject testItem = new TestObject(documentId, documentId, Arrays.asList(), documentId); idsAddedByBulkCreate.add(documentId); return CosmosBulkOperations.getCreateItemOperation(testItem, new PartitionKey(documentId)); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/ExcludedRegionWithFaultInjectionTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/ExcludedRegionWithFaultInjectionTests.java index 2038baa9d8d0..79bc5818cc92 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/ExcludedRegionWithFaultInjectionTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/ExcludedRegionWithFaultInjectionTests.java @@ -12,6 +12,7 @@ import com.azure.cosmos.CosmosExcludedRegions; import com.azure.cosmos.CosmosRegionSwitchHint; import com.azure.cosmos.SessionRetryOptionsBuilder; +import com.azure.cosmos.TestObject; import com.azure.cosmos.implementation.AsyncDocumentClient; import com.azure.cosmos.implementation.DatabaseAccount; import com.azure.cosmos.implementation.DatabaseAccountLocation; @@ -19,7 +20,6 @@ import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.TestConfigurations; import com.azure.cosmos.implementation.Utils; -import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosBatch; import com.azure.cosmos.models.CosmosBatchRequestOptions; import com.azure.cosmos.models.CosmosBatchResponse; @@ -150,15 +150,15 @@ public Object[][] regionExclusionReadAfterCreateTestConfigs() { Function> readItemCallback = (params) -> { - TestItem alreadyCreatedItem = params.createdItem; + TestObject alreadyCreatedItem = params.createdItem; try { - CosmosItemResponse response = params.cosmosAsyncContainer + CosmosItemResponse response = params.cosmosAsyncContainer .readItem( alreadyCreatedItem.getId(), new PartitionKey(alreadyCreatedItem.getId()), - TestItem.class) + TestObject.class) .block(); return new OperationExecutionResult<>(response); @@ -481,7 +481,7 @@ public Object[][] regionExclusionQueryAfterCreateTestConfigs() { Function> queryItemCallback = (params) -> { - TestItem alreadyCreatedItem = params.createdItem; + TestObject alreadyCreatedItem = params.createdItem; String query = String.format("SELECT * FROM c WHERE c.id = '%s'", alreadyCreatedItem.getId()); CosmosQueryRequestOptions queryRequestOptions = params.queryRequestOptionsForCallbackAfterMutation != null @@ -490,8 +490,8 @@ public Object[][] regionExclusionQueryAfterCreateTestConfigs() { SqlQuerySpec sqlQuerySpec = new SqlQuerySpec(query); try { - FeedResponse feedResponse = params.cosmosAsyncContainer - .queryItems(sqlQuerySpec, queryRequestOptions, TestItem.class) + FeedResponse feedResponse = params.cosmosAsyncContainer + .queryItems(sqlQuerySpec, queryRequestOptions, TestObject.class) .byPage() .blockFirst(); @@ -767,9 +767,9 @@ public Object[][] regionExclusionWriteAfterCreateTestConfigs() { itemRequestOptions.setNonIdempotentWriteRetryPolicy(params.nonIdempotentWriteRetriesEnabled, true); - CosmosItemResponse response = params.cosmosAsyncContainer + CosmosItemResponse response = params.cosmosAsyncContainer .createItem( - new TestItem(newDocumentId, newDocumentId, newDocumentId), + new TestObject(newDocumentId, newDocumentId, Arrays.asList(), newDocumentId), new PartitionKey(newDocumentId), itemRequestOptions) .block(); @@ -792,8 +792,8 @@ public Object[][] regionExclusionWriteAfterCreateTestConfigs() { Function> replaceItemCallback = (params) -> { - TestItem alreadyCreatedItem = params.createdItem; - alreadyCreatedItem.setProp(UUID.randomUUID().toString()); + TestObject alreadyCreatedItem = params.createdItem; + alreadyCreatedItem.setStringProp(UUID.randomUUID().toString()); CosmosItemRequestOptions itemRequestOptions = params.itemRequestOptionsForCallbackAfterMutation != null ? params.itemRequestOptionsForCallbackAfterMutation : new CosmosItemRequestOptions(); @@ -802,7 +802,7 @@ public Object[][] regionExclusionWriteAfterCreateTestConfigs() { try { - CosmosItemResponse response = params.cosmosAsyncContainer + CosmosItemResponse response = params.cosmosAsyncContainer .replaceItem( alreadyCreatedItem, alreadyCreatedItem.getId(), @@ -827,7 +827,7 @@ public Object[][] regionExclusionWriteAfterCreateTestConfigs() { Function> deleteItemCallback = (params) -> { - TestItem alreadyCreatedItem = params.createdItem; + TestObject alreadyCreatedItem = params.createdItem; CosmosItemRequestOptions itemRequestOptions = params.itemRequestOptionsForCallbackAfterMutation != null ? params.itemRequestOptionsForCallbackAfterMutation : new CosmosItemRequestOptions(); @@ -860,8 +860,8 @@ public Object[][] regionExclusionWriteAfterCreateTestConfigs() { Function> upsertExistingItemCallback = (params) -> { - TestItem alreadyCreatedItem = params.createdItem; - alreadyCreatedItem.setProp(UUID.randomUUID().toString()); + TestObject alreadyCreatedItem = params.createdItem; + alreadyCreatedItem.setStringProp(UUID.randomUUID().toString()); CosmosItemRequestOptions itemRequestOptions = params.itemRequestOptionsForCallbackAfterMutation != null ? params.itemRequestOptionsForCallbackAfterMutation : new CosmosItemRequestOptions(); @@ -870,7 +870,7 @@ public Object[][] regionExclusionWriteAfterCreateTestConfigs() { try { - CosmosItemResponse response = params.cosmosAsyncContainer.upsertItem( + CosmosItemResponse response = params.cosmosAsyncContainer.upsertItem( alreadyCreatedItem, new PartitionKey(alreadyCreatedItem.getId()), itemRequestOptions).block(); @@ -893,7 +893,7 @@ public Object[][] regionExclusionWriteAfterCreateTestConfigs() { Function> upsertNonExistingItemCallback = (params) -> { - TestItem newItem = TestItem.createNewItem(); + TestObject newItem = TestObject.create(); CosmosItemRequestOptions itemRequestOptions = params.itemRequestOptionsForCallbackAfterMutation != null ? params.itemRequestOptionsForCallbackAfterMutation : new CosmosItemRequestOptions(); @@ -902,7 +902,7 @@ public Object[][] regionExclusionWriteAfterCreateTestConfigs() { try { - CosmosItemResponse response = params.cosmosAsyncContainer.upsertItem( + CosmosItemResponse response = params.cosmosAsyncContainer.upsertItem( newItem, new PartitionKey(newItem.getId()), itemRequestOptions).block(); @@ -933,12 +933,12 @@ public Object[][] regionExclusionWriteAfterCreateTestConfigs() { patchItemRequestOptions.setNonIdempotentWriteRetryPolicy(params.nonIdempotentWriteRetriesEnabled, true); try { - CosmosItemResponse response = params.cosmosAsyncContainer.patchItem( + CosmosItemResponse response = params.cosmosAsyncContainer.patchItem( params.createdItem.getId(), new PartitionKey(params.createdItem.getId()), patchOperations, (CosmosPatchItemRequestOptions) patchItemRequestOptions, - TestItem.class + TestObject.class ).block(); return new OperationExecutionResult<>(response); @@ -1558,7 +1558,7 @@ public Object[][] regionExclusionBatchTestConfigs() { Function> batchCreateAndReadCallback = (params) -> { String documentId = UUID.randomUUID().toString(); - TestItem testItem = new TestItem(documentId, documentId, documentId); + TestObject testItem = new TestObject(documentId, documentId, Arrays.asList(), documentId); CosmosBatch batch = CosmosBatch.createCosmosBatch(new PartitionKey(documentId)); @@ -1933,7 +1933,7 @@ public Object[][] regionExclusionBulkTestConfigs() { Flux createOperationsFlux = Flux.range(0, totalCreates).map(i -> { String documentId = UUID.randomUUID().toString(); - TestItem testItem = new TestItem(documentId, documentId, documentId); + TestObject testItem = new TestObject(documentId, documentId, Arrays.asList(), documentId); return CosmosBulkOperations.getCreateItemOperation(testItem, new PartitionKey(documentId)); @@ -2364,7 +2364,7 @@ private void execute(MutationTestConfig mutationTestConfig, boolean shouldInject .getDatabase(this.cosmosAsyncContainer.getDatabase().getId()) .getContainer(this.cosmosAsyncContainer.getId()); - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); Thread.sleep(5_000); @@ -2675,11 +2675,11 @@ private static AccountLevelLocationContext getAccountLevelLocationContext(Databa } private static String getRegionResolvedForDefaultEndpoint(CosmosAsyncContainer container, List preferredRegions) { - TestItem testItem = TestItem.createNewItem(); + TestObject testItem = TestObject.create(); CosmosItemRequestOptions itemRequestOptions = new CosmosItemRequestOptions(); itemRequestOptions.setExcludedRegions(preferredRegions); - CosmosItemResponse response = container.createItem(testItem, itemRequestOptions).block(); + CosmosItemResponse response = container.createItem(testItem, itemRequestOptions).block(); Set contactedRegion = response.getDiagnostics().getContactedRegionNames(); @@ -2902,7 +2902,7 @@ public MutationTestConfig withPerRegionDuplicateCount(int perRegionDuplicateCoun private static class ItemOperationInvocationParameters { public boolean nonIdempotentWriteRetriesEnabled; public CosmosAsyncContainer cosmosAsyncContainer; - public TestItem createdItem; + public TestObject createdItem; public CosmosItemRequestOptions itemRequestOptionsForCallbackAfterMutation; public CosmosItemRequestOptions patchItemRequestOptionsForCallbackAfterMutation; public CosmosQueryRequestOptions queryRequestOptionsForCallbackAfterMutation; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionConnectionErrorRuleTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionConnectionErrorRuleTests.java index dc45e3134c68..b3eedd79ff17 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionConnectionErrorRuleTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionConnectionErrorRuleTests.java @@ -8,6 +8,7 @@ import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.CosmosClientBuilder; import com.azure.cosmos.CosmosContainerProactiveInitConfigBuilder; +import com.azure.cosmos.TestObject; import com.azure.cosmos.implementation.AsyncDocumentClient; import com.azure.cosmos.implementation.DatabaseAccount; import com.azure.cosmos.implementation.DatabaseAccountLocation; @@ -16,7 +17,6 @@ import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils; import com.azure.cosmos.implementation.directconnectivity.RntbdTransportClient; import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdEndpoint; -import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosContainerIdentity; import com.azure.cosmos.models.FeedRange; import com.azure.cosmos.models.PartitionKey; @@ -103,7 +103,7 @@ public void faultInjectionConnectionErrorRuleTestWithNoConnectionWarmup(FaultInj CosmosAsyncContainer singlePartitionContainer = getSharedSinglePartitionCosmosContainer(client); // validate one channel exists - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); singlePartitionContainer.createItem(createdItem).block(); RntbdTransportClient rntbdTransportClient = (RntbdTransportClient) ReflectionUtils.getTransportClient(this.client); @@ -146,7 +146,7 @@ public void faultInjectionConnectionErrorRuleTestWithNoConnectionWarmup(FaultInj assertThat(connectionErrorRule.getHitCountDetails()).isNull(); // do another request to open a new connection - singlePartitionContainer.createItem(TestItem.createNewItem()).block(); + singlePartitionContainer.createItem(TestObject.create()).block(); Thread.sleep(Duration.ofSeconds(2).toMillis()); // the configured connection rule should have disabled after 2s, so the connection will remain open @@ -200,7 +200,7 @@ public void faultInjectionConnectionErrorRuleTestWithConnectionWarmup(FaultInjec singlePartitionContainer = connectionWarmupClient.getDatabase(databaseId).getContainer(containerId); // validate one channel exists - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); singlePartitionContainer.createItem(createdItem).block(); RntbdTransportClient rntbdTransportClient = (RntbdTransportClient) ReflectionUtils.getTransportClient(connectionWarmupClient); @@ -269,7 +269,7 @@ public void faultInjectionConnectionErrorRuleTestWithConnectionWarmup(FaultInjec assertThat(connectionErrorRule.getHitCountDetails()).isNull(); // do another request to open a new connection - singlePartitionContainer.createItem(TestItem.createNewItem()).block(); + singlePartitionContainer.createItem(TestObject.create()).block(); Thread.sleep(Duration.ofSeconds(2).toMillis()); // the configured connection rule should have disabled after 2s, so the connection will remain open @@ -300,7 +300,7 @@ public void connectionCloseError_NoEndpoint_NoWarmup() { CosmosAsyncContainer singlePartitionContainer = getSharedSinglePartitionCosmosContainer(client); // validate one channel exists - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); singlePartitionContainer.createItem(createdItem).block(); RntbdTransportClient rntbdTransportClient = (RntbdTransportClient) ReflectionUtils.getTransportClient(this.client); @@ -339,7 +339,7 @@ public void connectionCloseError_NoEndpoint_NoWarmup() { assertThat(connectionErrorRule.getHitCountDetails()).isNull(); // do another request to open a new connection - singlePartitionContainer.createItem(TestItem.createNewItem()).block(); + singlePartitionContainer.createItem(TestObject.create()).block(); Thread.sleep(Duration.ofSeconds(2).toMillis()); // the configured connection rule should have disabled after 2s, so the connection will remain open diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionMetadataRequestRuleTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionMetadataRequestRuleTests.java index 39c1828e92f4..eb6c81686644 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionMetadataRequestRuleTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionMetadataRequestRuleTests.java @@ -10,6 +10,7 @@ import com.azure.cosmos.CosmosDiagnostics; import com.azure.cosmos.CosmosEndToEndOperationLatencyPolicyConfigBuilder; import com.azure.cosmos.CosmosException; +import com.azure.cosmos.TestObject; import com.azure.cosmos.implementation.AsyncDocumentClient; import com.azure.cosmos.implementation.DatabaseAccount; import com.azure.cosmos.implementation.DatabaseAccountLocation; @@ -18,7 +19,6 @@ import com.azure.cosmos.implementation.MetadataDiagnosticsContext; import com.azure.cosmos.implementation.OperationType; import com.azure.cosmos.implementation.Utils; -import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.FeedRange; @@ -190,7 +190,7 @@ public void faultInjectionServerErrorRuleTests_AddressRefresh_ConnectionDelay(bo .duration(Duration.ofMinutes(5)) .build(); try { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); container.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules( @@ -295,7 +295,7 @@ public void faultInjectionServerErrorRuleTests_AddressRefresh_ResponseDelay( .build(); try { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); container.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules( @@ -358,7 +358,7 @@ public void faultInjectionServerErrorRuleTests_AddressRefresh_byPartition(boolea // first create few documents for (int i = 0; i < 10; i++) { - container.createItem(TestItem.createNewItem()).block(); + container.createItem(TestObject.create()).block(); } List feedRanges = container.getFeedRanges().block(); @@ -367,14 +367,14 @@ public void faultInjectionServerErrorRuleTests_AddressRefresh_byPartition(boolea CosmosQueryRequestOptions cosmosQueryRequestOptions = new CosmosQueryRequestOptions(); cosmosQueryRequestOptions.setFeedRange(feedRanges.get(0)); String query = "select * from c"; - TestItem itemOnFeedRange1 = container.queryItems(query, cosmosQueryRequestOptions, TestItem.class) + TestObject itemOnFeedRange1 = container.queryItems(query, cosmosQueryRequestOptions, TestObject.class) .byPage(1) .blockFirst() .getResults() .get(0); cosmosQueryRequestOptions.setFeedRange(feedRanges.get(1)); - TestItem itemOnFeedRange2 = container.queryItems(query, cosmosQueryRequestOptions, TestItem.class) + TestObject itemOnFeedRange2 = container.queryItems(query, cosmosQueryRequestOptions, TestObject.class) .byPage(1) .blockFirst() .getResults() @@ -509,7 +509,7 @@ public void faultInjectionServerErrorRuleTests_AddressRefresh_TooManyRequest(boo .duration(Duration.ofMinutes(5)) .build(); try { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); container.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules( @@ -589,7 +589,7 @@ public void faultInjectionServerErrorRuleTests_PartitionKeyRanges_DelayError( try { // create few items to first make sure the collection cache, pkRanges cache is being populated for (int i = 0; i < 10; i++) { - container.createItem(TestItem.createNewItem()).block(); + container.createItem(TestObject.create()).block(); } CosmosFaultInjectionHelper.configureFaultInjectionRules( @@ -598,7 +598,7 @@ public void faultInjectionServerErrorRuleTests_PartitionKeyRanges_DelayError( .block(); try { - CosmosDiagnostics cosmosDiagnostics = container.createItem(TestItem.createNewItem()).block().getDiagnostics(); + CosmosDiagnostics cosmosDiagnostics = container.createItem(TestObject.create()).block().getDiagnostics(); // The PkRanges requests may have retried in another region, // but the create request will only be retried locally for PARTITION_IS_SPLITTING assertThat(cosmosDiagnostics.getContactedRegionNames().size()).isEqualTo(1); @@ -695,7 +695,7 @@ public void faultInjectionServerErrorRuleTests_CollectionRead_ConnectionDelay(bo // issue few requests to first populate all the necessary caches // as connection delay will impact all other operations, and in this test, we want to limit the scope to only collection read for (int i = 0; i < 10; i++) { - container.createItem(TestItem.createNewItem()).block(); + container.createItem(TestObject.create()).block(); } CosmosFaultInjectionHelper.configureFaultInjectionRules( @@ -704,7 +704,7 @@ public void faultInjectionServerErrorRuleTests_CollectionRead_ConnectionDelay(bo .block(); try { - CosmosDiagnostics cosmosDiagnostics = container.createItem(TestItem.createNewItem()).block().getDiagnostics(); + CosmosDiagnostics cosmosDiagnostics = container.createItem(TestObject.create()).block().getDiagnostics(); // validate CONTAINER_LOOK_UP ObjectNode diagnosticsNode = (ObjectNode) Utils.getSimpleObjectMapper().readTree(cosmosDiagnostics.toString()); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java index 899b262dc168..7c4e5c898a03 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java @@ -10,6 +10,7 @@ import com.azure.cosmos.CosmosClientBuilder; import com.azure.cosmos.CosmosDiagnostics; import com.azure.cosmos.DirectConnectionConfig; +import com.azure.cosmos.TestObject; import com.azure.cosmos.implementation.AsyncDocumentClient; import com.azure.cosmos.implementation.DatabaseAccount; import com.azure.cosmos.implementation.DatabaseAccountLocation; @@ -20,7 +21,6 @@ import com.azure.cosmos.implementation.ResourceType; import com.azure.cosmos.implementation.TestConfigurations; import com.azure.cosmos.implementation.Utils; -import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.FeedRange; @@ -206,7 +206,7 @@ public void faultInjectionServerErrorRuleTests_OperationType(OperationType opera .build(); try { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); cosmosAsyncContainer.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules(cosmosAsyncContainer, Arrays.asList(serverGoneErrorRule)).block(); @@ -254,7 +254,7 @@ public void faultInjectionServerErrorRuleTests_OperationType(OperationType opera @Test(groups = {"multi-region"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT) public void faultInjectionServerErrorRuleTests_OperationTypeImpactAddresses(OperationType operationType) throws JsonProcessingException { // Test the operation type can impact which region or replica the rule will be applicable - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); this.cosmosAsyncContainer.createItem(createdItem).block(); String writeRegionServerGoneRuleId = "serverErrorRule-writeRegionOnly-" + UUID.randomUUID(); @@ -413,7 +413,7 @@ public void faultInjectionServerErrorRuleTests_Region(boolean shouldInjectPrefer .getDatabase(this.cosmosAsyncContainer.getDatabase().getId()) .getContainer(this.cosmosAsyncContainer.getId()); - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); container.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules( @@ -457,7 +457,7 @@ public void faultInjectionServerErrorRuleTests_Region(boolean shouldInjectPrefer @Test(groups = {"multi-region", "long"}, timeOut = TIMEOUT) public void faultInjectionServerErrorRuleTests_Partition() throws JsonProcessingException { for (int i = 0; i < 10; i++) { - cosmosAsyncContainer.createItem(TestItem.createNewItem()).block(); + cosmosAsyncContainer.createItem(TestObject.create()).block(); } // getting one item from each feedRange @@ -467,10 +467,10 @@ public void faultInjectionServerErrorRuleTests_Partition() throws JsonProcessing String query = "select * from c"; CosmosQueryRequestOptions cosmosQueryRequestOptions = new CosmosQueryRequestOptions(); cosmosQueryRequestOptions.setFeedRange(feedRanges.get(0)); - TestItem itemOnFeedRange0 = cosmosAsyncContainer.queryItems(query, cosmosQueryRequestOptions, TestItem.class).blockFirst(); + TestObject itemOnFeedRange0 = cosmosAsyncContainer.queryItems(query, cosmosQueryRequestOptions, TestObject.class).blockFirst(); cosmosQueryRequestOptions.setFeedRange(feedRanges.get(1)); - TestItem itemOnFeedRange1 = cosmosAsyncContainer.queryItems(query, cosmosQueryRequestOptions, TestItem.class).blockFirst(); + TestObject itemOnFeedRange1 = cosmosAsyncContainer.queryItems(query, cosmosQueryRequestOptions, TestObject.class).blockFirst(); // set rule by feed range String feedRangeRuleId = "ServerErrorRule-FeedRange-" + UUID.randomUUID(); @@ -568,12 +568,12 @@ public void faultInjectionServerErrorRuleTests_ServerResponseDelay() throws Json .getContainer(cosmosAsyncContainer.getId()); // create a new item to be used by read operations - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); container.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules(container, Arrays.asList(timeoutRule)).block(); - CosmosItemResponse itemResponse = - container.readItem(createdItem.getId(), new PartitionKey(createdItem.getId()), TestItem.class).block(); + CosmosItemResponse itemResponse = + container.readItem(createdItem.getId(), new PartitionKey(createdItem.getId()), TestObject.class).block(); assertThat(timeoutRule.getHitCount()).isEqualTo(1); this.validateHitCount(timeoutRule, 1, OperationType.Read, ResourceType.Document); @@ -633,7 +633,7 @@ public void faultInjectionServerErrorRuleTests_ServerConnectionTimeout() throws .getContainer(cosmosAsyncContainer.getId()); CosmosFaultInjectionHelper.configureFaultInjectionRules(container, Arrays.asList(serverConnectionDelayRule)).block(); - CosmosItemResponse itemResponse = container.createItem(TestItem.createNewItem()).block(); + CosmosItemResponse itemResponse = container.createItem(TestObject.create()).block(); // Due to the replica validation, there could be an extra open connection call flow, while the rule will also be applied on. assertThat(serverConnectionDelayRule.getHitCount()).isBetween(1l, 2l); @@ -688,7 +688,7 @@ public void faultInjectionServerErrorRuleTests_ServerConnectionDelay() throws Js .getContainer(cosmosAsyncContainer.getId()); CosmosFaultInjectionHelper.configureFaultInjectionRules(container, Arrays.asList(serverConnectionDelayRule)).block(); - CosmosDiagnostics cosmosDiagnostics = container.createItem(TestItem.createNewItem()).block().getDiagnostics(); + CosmosDiagnostics cosmosDiagnostics = container.createItem(TestObject.create()).block().getDiagnostics(); // verify the request succeeded and the rule has applied List diagnosticsNode = new ArrayList<>(); @@ -840,7 +840,7 @@ public void faultInjectionServerErrorRuleTests_ServerErrorResponse( } // simulate high channel acquisition/connectionTimeout for read/query - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); cosmosAsyncContainer.createItem(createdItem).block(); String ruleId = "serverErrorRule-" + serverErrorType + "-" + UUID.randomUUID(); @@ -882,7 +882,7 @@ public void faultInjectionServerErrorRuleTests_ServerErrorResponse( @Test(groups = {"multi-region", "long"}, timeOut = TIMEOUT) public void faultInjectionServerErrorRuleTests_HitLimit() throws JsonProcessingException { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); cosmosAsyncContainer.createItem(createdItem).block(); // set rule by feed range @@ -941,7 +941,7 @@ public void afterClass() { @Test(groups = {"long"}, timeOut = TIMEOUT) public void faultInjectionServerErrorRuleTests_includePrimary() throws JsonProcessingException { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); CosmosAsyncContainer singlePartitionContainer = getSharedSinglePartitionCosmosContainer(clientWithoutPreferredRegions); List feedRanges = singlePartitionContainer.getFeedRanges().block(); @@ -1015,7 +1015,7 @@ public void faultInjectionServerErrorRuleTests_StaledAddressesServerGone() throw .build(); try { - TestItem testItem = this.cosmosAsyncContainer.createItem(TestItem.createNewItem()).block().getItem(); + TestObject testItem = this.cosmosAsyncContainer.createItem(TestObject.create()).block().getItem(); CosmosFaultInjectionHelper.configureFaultInjectionRules(this.cosmosAsyncContainer, Arrays.asList(staledAddressesServerGoneRule)).block(); @@ -1078,7 +1078,7 @@ public void connectionAcquisitionTimeoutAlignConnectionTimeout() throws JsonProc // validate none of the channelAcquisition stage take more than 2s List results = new ArrayList<>(); Flux.range(1, 6) - .flatMap(t -> containerWithSinglePartition.createItem(TestItem.createNewItem())) + .flatMap(t -> containerWithSinglePartition.createItem(TestObject.create())) .doOnNext(response -> results.add(response.getDiagnostics())) .blockLast(); @@ -1117,7 +1117,7 @@ public void faultInjectionServerErrorRuleTests_InjectionRate50Percent() throws J .build(); try { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); cosmosAsyncContainer.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules( @@ -1159,7 +1159,7 @@ public void faultInjectionServerErrorRuleTests_InjectionRate100Percent() throws .build(); try { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); cosmosAsyncContainer.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules( @@ -1226,7 +1226,7 @@ public void faultInjectionServerErrorRuleTests_InjectionRate25Percent() throws J .build(); try { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); cosmosAsyncContainer.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules( @@ -1268,7 +1268,7 @@ public void faultInjectionServerErrorRuleTests_InjectionRate75Percent() throws J .build(); try { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); cosmosAsyncContainer.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules( diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnGatewayTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnGatewayTests.java index 1ef55aa5fa33..5467613ba409 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnGatewayTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnGatewayTests.java @@ -10,6 +10,7 @@ import com.azure.cosmos.CosmosDiagnostics; import com.azure.cosmos.CosmosException; import com.azure.cosmos.DirectConnectionConfig; +import com.azure.cosmos.TestObject; import com.azure.cosmos.implementation.AsyncDocumentClient; import com.azure.cosmos.implementation.DatabaseAccount; import com.azure.cosmos.implementation.DatabaseAccountLocation; @@ -19,7 +20,6 @@ import com.azure.cosmos.implementation.ResourceType; import com.azure.cosmos.implementation.TestConfigurations; import com.azure.cosmos.implementation.Utils; -import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.FeedRange; @@ -194,7 +194,7 @@ public void faultInjectionServerErrorRuleTests_Region(boolean shouldUsePreferred .getDatabase(this.cosmosAsyncContainer.getDatabase().getId()) .getContainer(this.cosmosAsyncContainer.getId()); - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); container.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules( @@ -251,7 +251,7 @@ public void faultInjectionServerErrorRuleTests_Partition() throws JsonProcessing .getContainer(cosmosAsyncContainer.getId()); for (int i = 0; i < 10; i++) { - testContainer.createItem(TestItem.createNewItem()).block(); + testContainer.createItem(TestObject.create()).block(); } // getting one item from each feedRange @@ -261,10 +261,10 @@ public void faultInjectionServerErrorRuleTests_Partition() throws JsonProcessing String query = "select * from c"; CosmosQueryRequestOptions cosmosQueryRequestOptions = new CosmosQueryRequestOptions(); cosmosQueryRequestOptions.setFeedRange(feedRanges.get(0)); - TestItem itemOnFeedRange0 = testContainer.queryItems(query, cosmosQueryRequestOptions, TestItem.class).blockFirst(); + TestObject itemOnFeedRange0 = testContainer.queryItems(query, cosmosQueryRequestOptions, TestObject.class).blockFirst(); cosmosQueryRequestOptions.setFeedRange(feedRanges.get(1)); - TestItem itemOnFeedRange1 = testContainer.queryItems(query, cosmosQueryRequestOptions, TestItem.class).blockFirst(); + TestObject itemOnFeedRange1 = testContainer.queryItems(query, cosmosQueryRequestOptions, TestObject.class).blockFirst(); // set rule by feed range String feedRangeRuleId = "ServerErrorRule-FeedRange-" + UUID.randomUUID(); @@ -350,12 +350,12 @@ public void faultInjectionServerErrorRuleTests_ServerResponseDelay() throws Json directConnectionConfig.setConnectTimeout(Duration.ofSeconds(1)); // create a new item to be used by read operations - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); this.cosmosAsyncContainer.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules(this.cosmosAsyncContainer, Arrays.asList(timeoutRule)).block(); - CosmosItemResponse itemResponse = - this.cosmosAsyncContainer.readItem(createdItem.getId(), new PartitionKey(createdItem.getId()), TestItem.class).block(); + CosmosItemResponse itemResponse = + this.cosmosAsyncContainer.readItem(createdItem.getId(), new PartitionKey(createdItem.getId()), TestObject.class).block(); assertThat(timeoutRule.getHitCount()).isEqualTo(1); this.validateHitCount(timeoutRule, 1, OperationType.Read, ResourceType.Document); @@ -398,7 +398,7 @@ public void faultInjectionServerErrorRuleTests_ServerConnectionDelay() throws Js try { CosmosFaultInjectionHelper.configureFaultInjectionRules(this.cosmosAsyncContainer, Arrays.asList(serverConnectionDelayRule)).block(); - CosmosItemResponse itemResponse = this.cosmosAsyncContainer.createItem(TestItem.createNewItem()).block(); + CosmosItemResponse itemResponse = this.cosmosAsyncContainer.createItem(TestObject.create()).block(); assertThat(serverConnectionDelayRule.getHitCount()).isEqualTo(1l); this.validateFaultInjectionRuleApplied( @@ -449,7 +449,7 @@ public void faultInjectionServerErrorRuleTests_ServerErrorResponse( .build(); try { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); cosmosAsyncContainer.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules(cosmosAsyncContainer, Arrays.asList(serverErrorRule)).block(); @@ -459,7 +459,7 @@ public void faultInjectionServerErrorRuleTests_ServerErrorResponse( try { cosmosDiagnostics = cosmosAsyncContainer - .readItem(createdItem.getId(), new PartitionKey(createdItem.getId()), TestItem.class) + .readItem(createdItem.getId(), new PartitionKey(createdItem.getId()), TestObject.class) .block() .getDiagnostics(); } catch (Exception exception) { @@ -469,7 +469,7 @@ public void faultInjectionServerErrorRuleTests_ServerErrorResponse( try { cosmosDiagnostics = cosmosAsyncContainer - .readItem(createdItem.getId(), new PartitionKey(createdItem.getId()), TestItem.class) + .readItem(createdItem.getId(), new PartitionKey(createdItem.getId()), TestObject.class) .block() .getDiagnostics(); fail("Request should fail, but succeeded"); @@ -499,7 +499,7 @@ public void faultInjectionServerErrorRuleTests_HitLimit( OperationType operationType, FaultInjectionOperationType faultInjectionOperationType) throws JsonProcessingException { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); cosmosAsyncContainer.createItem(createdItem).block(); // set rule by feed range diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionTestBase.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionTestBase.java index 0996726cedde..934cb8bd3193 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionTestBase.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionTestBase.java @@ -7,8 +7,8 @@ import com.azure.cosmos.CosmosClientBuilder; import com.azure.cosmos.CosmosDiagnostics; import com.azure.cosmos.CosmosException; +import com.azure.cosmos.TestObject; import com.azure.cosmos.implementation.OperationType; -import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; import com.azure.cosmos.models.CosmosPatchOperations; import com.azure.cosmos.models.CosmosQueryRequestOptions; @@ -27,13 +27,13 @@ public FaultInjectionTestBase(CosmosClientBuilder cosmosClientBuilder) { protected CosmosDiagnostics performDocumentOperation( CosmosAsyncContainer cosmosAsyncContainer, OperationType operationType, - TestItem createdItem) { + TestObject createdItem) { try { if (operationType == OperationType.Query) { CosmosQueryRequestOptions queryRequestOptions = new CosmosQueryRequestOptions(); String query = String.format("SELECT * from c where c.id = '%s'", createdItem.getId()); - FeedResponse itemFeedResponse = - cosmosAsyncContainer.queryItems(query, queryRequestOptions, TestItem.class).byPage().blockFirst(); + FeedResponse itemFeedResponse = + cosmosAsyncContainer.queryItems(query, queryRequestOptions, TestObject.class).byPage().blockFirst(); return itemFeedResponse.getCosmosDiagnostics(); } @@ -49,7 +49,7 @@ protected CosmosDiagnostics performDocumentOperation( return cosmosAsyncContainer.readItem( createdItem.getId(), new PartitionKey(createdItem.getId()), - TestItem.class).block().getDiagnostics(); + TestObject.class).block().getDiagnostics(); } if (operationType == OperationType.Replace) { @@ -64,11 +64,11 @@ protected CosmosDiagnostics performDocumentOperation( } if (operationType == OperationType.Create) { - return cosmosAsyncContainer.createItem(TestItem.createNewItem()).block().getDiagnostics(); + return cosmosAsyncContainer.createItem(TestObject.create()).block().getDiagnostics(); } if (operationType == OperationType.Upsert) { - return cosmosAsyncContainer.upsertItem(TestItem.createNewItem()).block().getDiagnostics(); + return cosmosAsyncContainer.upsertItem(TestObject.create()).block().getDiagnostics(); } if (operationType == OperationType.Patch) { @@ -77,7 +77,7 @@ protected CosmosDiagnostics performDocumentOperation( .create() .add("newPath", "newPath"); return cosmosAsyncContainer - .patchItem(createdItem.getId(), new PartitionKey(createdItem.getId()), patchOperations, TestItem.class) + .patchItem(createdItem.getId(), new PartitionKey(createdItem.getId()), patchOperations, TestObject.class) .block().getDiagnostics(); } } @@ -87,8 +87,8 @@ protected CosmosDiagnostics performDocumentOperation( CosmosChangeFeedRequestOptions changeFeedRequestOptions = CosmosChangeFeedRequestOptions.createForProcessingFromBeginning(feedRanges.get(0)); - FeedResponse firstPage = cosmosAsyncContainer - .queryChangeFeed(changeFeedRequestOptions, TestItem.class) + FeedResponse firstPage = cosmosAsyncContainer + .queryChangeFeed(changeFeedRequestOptions, TestObject.class) .byPage() .blockFirst(); return firstPage.getCosmosDiagnostics(); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/SessionRetryOptionsTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/SessionRetryOptionsTests.java index 10d2b930c35f..261ae9ee0a27 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/SessionRetryOptionsTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/SessionRetryOptionsTests.java @@ -11,6 +11,7 @@ import com.azure.cosmos.CosmosRegionSwitchHint; import com.azure.cosmos.SessionRetryOptions; import com.azure.cosmos.SessionRetryOptionsBuilder; +import com.azure.cosmos.TestObject; import com.azure.cosmos.implementation.AsyncDocumentClient; import com.azure.cosmos.implementation.Configs; import com.azure.cosmos.implementation.DatabaseAccount; @@ -20,7 +21,6 @@ import com.azure.cosmos.implementation.ImplementationBridgeHelpers; import com.azure.cosmos.implementation.OperationType; import com.azure.cosmos.implementation.TestConfigurations; -import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosItemRequestOptions; import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.CosmosPatchOperations; @@ -237,7 +237,7 @@ public void nonWriteOperation_WithReadSessionUnavailable_test( .getDatabase(this.cosmosAsyncContainer.getDatabase().getId()) .getContainer(this.cosmosAsyncContainer.getId()); - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); containerForClientWithPreferredRegions.createItem(createdItem).block(); FaultInjectionRuleBuilder badSessionTokenRuleBuilder = new FaultInjectionRuleBuilder("serverErrorRule-bad-session-token-" + UUID.randomUUID()); @@ -333,7 +333,7 @@ public void writeOperation_withReadSessionUnavailable_test( .configureFaultInjectionRules(asyncContainerFromClientWithPreferredRegions, Arrays.asList(badSessionTokenRule)) .block(); - TestItem testItem = TestItem.createNewItem(); + TestObject testItem = TestObject.create(); validateOperationExecutionResult( performDocumentOperation( @@ -378,7 +378,7 @@ private Map getRegionMap(DatabaseAccount databaseAccount, boolea private OperationExecutionResult performDocumentOperation( CosmosAsyncContainer faultInjectedContainer, - TestItem testItem, + TestObject testItem, OperationType operationType) { AtomicReference operationStart = new AtomicReference<>(Instant.now()); @@ -389,8 +389,8 @@ private OperationExecutionResult performDocumentOperation( CosmosQueryRequestOptions queryRequestOptions = new CosmosQueryRequestOptions(); SqlQuerySpec sqlQuerySpec = new SqlQuerySpec(query); - FeedResponse feedResponse = faultInjectedContainer - .queryItems(sqlQuerySpec, queryRequestOptions, TestItem.class) + FeedResponse feedResponse = faultInjectedContainer + .queryItems(sqlQuerySpec, queryRequestOptions, TestObject.class) .byPage() .doOnSubscribe(ignore -> operationStart.set(Instant.now())) .doOnCancel(() -> operationEnd.set(Instant.now())) @@ -409,8 +409,8 @@ private OperationExecutionResult performDocumentOperation( if (operationType == OperationType.Read) { - CosmosItemResponse itemResponse = faultInjectedContainer - .readItem(testItem.getId(), new PartitionKey(testItem.getId()), TestItem.class) + CosmosItemResponse itemResponse = faultInjectedContainer + .readItem(testItem.getId(), new PartitionKey(testItem.getId()), TestObject.class) .doOnSubscribe(ignore -> operationStart.set(Instant.now())) .doOnSuccess(ignore -> operationEnd.set(Instant.now())) .block(); @@ -426,7 +426,7 @@ private OperationExecutionResult performDocumentOperation( CosmosItemRequestOptions itemRequestOptions = new CosmosItemRequestOptions(); - CosmosItemResponse itemResponse = faultInjectedContainer + CosmosItemResponse itemResponse = faultInjectedContainer .createItem(testItem, new PartitionKey(testItem.getId()), itemRequestOptions) .doOnSubscribe(ignore -> operationStart.set(Instant.now())) .doOnSuccess(ignore -> operationEnd.set(Instant.now())) @@ -447,7 +447,7 @@ private OperationExecutionResult performDocumentOperation( .createItem(testItem, new PartitionKey(testItem.getId()), itemRequestOptions) .block(); - CosmosItemResponse itemResponse = faultInjectedContainer + CosmosItemResponse itemResponse = faultInjectedContainer .replaceItem(testItem, testItem.getId(), new PartitionKey(testItem.getId()), itemRequestOptions) .doOnSubscribe(ignore -> operationStart.set(Instant.now())) .doOnSuccess(ignore -> operationEnd.set(Instant.now())) @@ -483,7 +483,7 @@ private OperationExecutionResult performDocumentOperation( if (operationType == OperationType.Upsert) { CosmosItemRequestOptions itemRequestOptions = new CosmosItemRequestOptions(); - CosmosItemResponse itemResponse = faultInjectedContainer + CosmosItemResponse itemResponse = faultInjectedContainer .upsertItem(testItem, new PartitionKey(testItem.getId()), itemRequestOptions) .doOnSubscribe(ignore -> operationStart.set(Instant.now())) .doOnSuccess(ignore -> operationEnd.set(Instant.now())) @@ -504,8 +504,8 @@ private OperationExecutionResult performDocumentOperation( .createItem(testItem, new PartitionKey(testItem.getId()), itemRequestOptions) .block(); - CosmosItemResponse itemResponse = faultInjectedContainer - .patchItem(testItem.getId(), new PartitionKey(testItem.getId()), patchOperations, TestItem.class) + CosmosItemResponse itemResponse = faultInjectedContainer + .patchItem(testItem.getId(), new PartitionKey(testItem.getId()), patchOperations, TestObject.class) .doOnSubscribe(ignore -> operationStart.set(Instant.now())) .doOnSuccess(ignore -> operationEnd.set(Instant.now())) .block(); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/MetadataRequestRetryPolicyTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/MetadataRequestRetryPolicyTests.java index 679d4444c6cb..210112149959 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/MetadataRequestRetryPolicyTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/MetadataRequestRetryPolicyTests.java @@ -13,6 +13,7 @@ import com.azure.cosmos.CosmosEndToEndOperationLatencyPolicyConfigBuilder; import com.azure.cosmos.CosmosException; import com.azure.cosmos.DirectConnectionConfig; +import com.azure.cosmos.TestObject; import com.azure.cosmos.implementation.AsyncDocumentClient; import com.azure.cosmos.implementation.Configs; import com.azure.cosmos.implementation.DatabaseAccount; @@ -33,7 +34,6 @@ import com.azure.cosmos.implementation.http.HttpClientConfig; import com.azure.cosmos.implementation.http.HttpTimeoutPolicyControlPlaneHotPath; import com.azure.cosmos.implementation.routing.RegionalRoutingContext; -import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosItemRequestOptions; import com.azure.cosmos.models.CosmosPatchOperations; import com.azure.cosmos.models.CosmosQueryRequestOptions; @@ -291,7 +291,7 @@ public void forceBackgroundAddressRefresh_onConnectionTimeoutAndRequestCancellat CosmosEndToEndOperationLatencyPolicyConfig cosmosEndToEndOperationLatencyPolicyConfigForFaultyOperation = new CosmosEndToEndOperationLatencyPolicyConfigBuilder(Duration.ofMillis(500)).build(); - TestItem testItem = new TestItem(UUID.randomUUID().toString(), UUID.randomUUID().toString(), UUID.randomUUID().toString()); + TestObject testItem = new TestObject(UUID.randomUUID().toString(), UUID.randomUUID().toString(), Arrays.asList(), UUID.randomUUID().toString()); performDocumentOperation( container, @@ -378,7 +378,7 @@ public void metadataRetryPolicyTest( private void performDocumentOperation( CosmosAsyncContainer faultInjectedContainer, - TestItem testItem, + TestObject testItem, OperationType faultInjectedOperationType, FaultInjectionRule connectionDelayFault, HttpClientUnderTestWrapper httpClientUnderTestWrapper, @@ -409,7 +409,7 @@ private void performDocumentOperation( SqlQuerySpec sqlQuerySpec = new SqlQuerySpec(query); faultInjectedContainer - .queryItems(sqlQuerySpec, queryRequestOptions, TestItem.class) + .queryItems(sqlQuerySpec, queryRequestOptions, TestObject.class) .byPage() .subscribe(); } else if (faultInjectedOperationType == OperationType.Read) { @@ -427,7 +427,7 @@ private void performDocumentOperation( .setCosmosEndToEndOperationLatencyPolicyConfig(endToEndOperationLatencyPolicyConfigForFaultyOperation); faultInjectedContainer - .readItem(testItem.getId(), new PartitionKey(testItem.getMypk()), requestOptionsForRead, TestItem.class) + .readItem(testItem.getId(), new PartitionKey(testItem.getMypk()), requestOptionsForRead, TestObject.class) .subscribe(); } else if (faultInjectedOperationType == OperationType.Create) { @@ -496,7 +496,7 @@ private void performDocumentOperation( Thread.sleep(idleTimeInMillis); faultInjectedContainer - .patchItem(testItem.getId(), new PartitionKey(testItem.getMypk()), patchOperations, TestItem.class) + .patchItem(testItem.getId(), new PartitionKey(testItem.getMypk()), patchOperations, TestObject.class) .subscribe(); } } diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ClientRetryPolicyE2ETests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ClientRetryPolicyE2ETests.java index 486dd3e1d2c2..7a409299eaac 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ClientRetryPolicyE2ETests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ClientRetryPolicyE2ETests.java @@ -13,6 +13,7 @@ import com.azure.cosmos.CosmosException; import com.azure.cosmos.DirectConnectionConfig; import com.azure.cosmos.FlakyTestRetryAnalyzer; +import com.azure.cosmos.TestObject; import com.azure.cosmos.implementation.AsyncDocumentClient; import com.azure.cosmos.implementation.DatabaseAccount; import com.azure.cosmos.implementation.DatabaseAccountLocation; @@ -20,7 +21,6 @@ import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.OperationType; import com.azure.cosmos.implementation.Utils; -import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.CosmosPatchOperations; @@ -151,7 +151,7 @@ public static Object[] preferredRegionsConfigProvider() { @Test(groups = { "multi-master" }, dataProvider = "preferredRegionsConfigProvider", timeOut = TIMEOUT) public void queryPlanHttpTimeoutWillNotMarkRegionUnavailable(boolean shouldUsePreferredRegionsOnClient) { - TestItem newItem = TestItem.createNewItem(); + TestObject newItem = TestObject.create(); CosmosAsyncContainer resultantCosmosAsyncContainer; CosmosAsyncClient resultantCosmosAsyncClient; @@ -194,7 +194,7 @@ public void queryPlanHttpTimeoutWillNotMarkRegionUnavailable(boolean shouldUsePr try { // validate the query plan will be retried in a different region and the final requests will be succeeded // TODO: Also capture all retries for metadata requests in the diagnostics - FeedResponse firstPage = cosmosAsyncContainerFromClientWithPreferredRegions.queryItems(query, queryRequestOptions, TestItem.class) + FeedResponse firstPage = cosmosAsyncContainerFromClientWithPreferredRegions.queryItems(query, queryRequestOptions, TestObject.class) .byPage() .blockFirst(); @@ -229,7 +229,7 @@ public void addressRefreshHttpTimeoutWillDoCrossRegionRetryForReads(boolean shou throw new SkipException("queryPlanHttpTimeoutWillNotMarkRegionUnavailable() is only meant for DIRECT mode"); } - TestItem newItem = TestItem.createNewItem(); + TestObject newItem = TestObject.create(); resultantCosmosAsyncContainer.createItem(newItem).block(); // create fault injection rules for address refresh @@ -264,8 +264,8 @@ public void addressRefreshHttpTimeoutWillDoCrossRegionRetryForReads(boolean shou resultantCosmosAsyncContainer, Arrays.asList(addressRefreshDelayRule, serverGoneRule)).block(); try { - CosmosItemResponse itemResponse = resultantCosmosAsyncContainer - .readItem(newItem.getId(), new PartitionKey(newItem.getId()), TestItem.class) + CosmosItemResponse itemResponse = resultantCosmosAsyncContainer + .readItem(newItem.getId(), new PartitionKey(newItem.getId()), TestObject.class) .block(); assertThat(itemResponse).isNotNull(); @@ -334,7 +334,7 @@ public void addressRefreshHttpTimeoutWillNotDoCrossRegionRetryForWrites(boolean resultantCosmosAsyncContainer, Arrays.asList(addressRefreshDelayRule, serverGoneRule)).block(); try { - TestItem newItem = TestItem.createNewItem(); + TestObject newItem = TestObject.create(); resultantCosmosAsyncContainer.createItem(newItem).block(); } catch (CosmosException e) { assertThat(e.getDiagnostics().getContactedRegionNames().size()).isEqualTo(1); @@ -372,7 +372,7 @@ public void dataPlaneRequestHttpTimeout( throw new SkipException("queryPlanHttpTimeoutWillNotMarkRegionUnavailable() is only meant for GATEWAY mode"); } - TestItem newItem = TestItem.createNewItem(); + TestObject newItem = TestObject.create(); resultantCosmosAsyncContainer.createItem(newItem).block(); FaultInjectionRule requestHttpTimeoutRule = new FaultInjectionRuleBuilder("requestHttpTimeoutRule" + UUID.randomUUID()) .condition( @@ -398,7 +398,7 @@ public void dataPlaneRequestHttpTimeout( resultantCosmosAsyncContainer, operationType, newItem, - (testItem) -> new PartitionKey(testItem.getId()) + (TestObject) -> new PartitionKey(TestObject.getId()) ).block(); assertThat(cosmosDiagnostics.getContactedRegionNames().size()).isEqualTo(this.preferredRegions.size()); @@ -412,7 +412,7 @@ public void dataPlaneRequestHttpTimeout( resultantCosmosAsyncContainer, operationType, newItem, - (testItem) -> new PartitionKey(testItem.getId()) + (TestObject) -> new PartitionKey(TestObject.getId()) ).block(); fail("dataPlaneRequestHttpTimeout() should have failed for operationType " + operationType); } catch (CosmosException e) { @@ -475,7 +475,7 @@ public void channelAcquisitionExceptionOnWrites( CosmosFaultInjectionHelper.configureFaultInjectionRules(testContainer, Arrays.asList(channelAcquisitionExceptionRule)).block(); try { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); testContainer.createItem(createdItem).block(); // using a higher concurrency to force channelAcquisitionException to happen @@ -486,7 +486,7 @@ public void channelAcquisitionExceptionOnWrites( testContainer, operationType, createdItem, - (testItem) -> new PartitionKey(testItem.getMypk()))) + (TestObject) -> new PartitionKey(TestObject.getMypk()))) .doOnNext(diagnostics -> { // since we have only injected connection delay error in one region, so we should only see 2 regions being contacted eventually assertThat(diagnostics.getContactedRegionNames().size()).isEqualTo(2); @@ -538,13 +538,13 @@ private boolean isChannelAcquisitionExceptionTriggeredRegionRetryExists(String c private Mono performDocumentOperation( CosmosAsyncContainer cosmosAsyncContainer, OperationType operationType, - TestItem createdItem, - Function extractPartitionKeyFunc) { + TestObject createdItem, + Function extractPartitionKeyFunc) { if (operationType == OperationType.Query) { CosmosQueryRequestOptions queryRequestOptions = new CosmosQueryRequestOptions(); String query = String.format("SELECT * from c where c.id = '%s'", createdItem.getId()); - FeedResponse itemFeedResponse = - cosmosAsyncContainer.queryItems(query, queryRequestOptions, TestItem.class).byPage().blockFirst(); + FeedResponse itemFeedResponse = + cosmosAsyncContainer.queryItems(query, queryRequestOptions, TestObject.class).byPage().blockFirst(); return Mono.just(itemFeedResponse.getCosmosDiagnostics()); } @@ -560,7 +560,7 @@ private Mono performDocumentOperation( .readItem( createdItem.getId(), extractPartitionKeyFunc.apply(createdItem), - TestItem.class + TestObject.class ) .map(itemResponse -> itemResponse.getDiagnostics()); } @@ -579,11 +579,11 @@ private Mono performDocumentOperation( } if (operationType == OperationType.Create) { - return cosmosAsyncContainer.createItem(TestItem.createNewItem()).map(itemResponse -> itemResponse.getDiagnostics()); + return cosmosAsyncContainer.createItem(TestObject.create()).map(itemResponse -> itemResponse.getDiagnostics()); } if (operationType == OperationType.Upsert) { - return cosmosAsyncContainer.upsertItem(TestItem.createNewItem()).map(itemResponse -> itemResponse.getDiagnostics()); + return cosmosAsyncContainer.upsertItem(TestObject.create()).map(itemResponse -> itemResponse.getDiagnostics()); } if (operationType == OperationType.Patch) { @@ -596,7 +596,7 @@ private Mono performDocumentOperation( createdItem.getId(), extractPartitionKeyFunc.apply(createdItem), patchOperations, - TestItem.class) + TestObject.class) .map(itemResponse -> itemResponse.getDiagnostics()); } } @@ -606,8 +606,8 @@ private Mono performDocumentOperation( CosmosChangeFeedRequestOptions changeFeedRequestOptions = CosmosChangeFeedRequestOptions.createForProcessingFromBeginning(feedRanges.get(0)); - FeedResponse firstPage = cosmosAsyncContainer - .queryChangeFeed(changeFeedRequestOptions, TestItem.class) + FeedResponse firstPage = cosmosAsyncContainer + .queryChangeFeed(changeFeedRequestOptions, TestObject.class) .byPage() .blockFirst(); return Mono.just(firstPage.getCosmosDiagnostics()); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/FullFidelityChangeFeedTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/FullFidelityChangeFeedTest.java index 4b7eb6ea612e..005e09bb0968 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/FullFidelityChangeFeedTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/FullFidelityChangeFeedTest.java @@ -6,7 +6,7 @@ import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.CosmosAsyncDatabase; import com.azure.cosmos.CosmosClientBuilder; -import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; +import com.azure.cosmos.TestObject; import com.azure.cosmos.models.ChangeFeedPolicy; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; import com.azure.cosmos.models.CosmosContainerProperties; @@ -21,12 +21,14 @@ import org.testng.annotations.Test; import java.time.Duration; +import java.util.Arrays; import java.util.Iterator; import java.util.List; import java.util.UUID; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.fail; +import static org.assertj.core.api.InstanceOfAssertFactories.ARRAY; public class FullFidelityChangeFeedTest extends TestSuiteBase { @@ -60,22 +62,22 @@ public void fullFidelityChangeFeed_FromNowForLogicalPartition() throws Exception continuationToken1 = response.getContinuationToken(); } - TestItem item1 = new TestItem( + TestObject item1 = new TestObject( UUID.randomUUID().toString(), - "mypk-1", "Johnson"); - TestItem item2 = new TestItem( + "mypk-1", Arrays.asList(), "Johnson"); + TestObject item2 = new TestObject( UUID.randomUUID().toString(), - "mypk-1", "Smith"); - TestItem item3 = new TestItem( + "mypk-1", Arrays.asList(), "Smith"); + TestObject item3 = new TestObject( UUID.randomUUID().toString(), - "mypk-2", "John"); + "mypk-2", Arrays.asList(), "John"); cosmosContainer.createItem(item1).block(); cosmosContainer.createItem(item2).block(); - String originalLastNameItem1 = item1.getProp(); - item1.setProp("Gates"); + String originalLastNameItem1 = item1.getStringProp(); + item1.setStringProp("Gates"); cosmosContainer.upsertItem(item1).block(); - String originalLastNameItem2 = item2.getProp(); - item2.setProp("Doe"); + String originalLastNameItem2 = item2.getStringProp(); + item2.setStringProp("Doe"); cosmosContainer.upsertItem(item2).block(); cosmosContainer.deleteItem(item1, new CosmosItemRequestOptions()).block(); @@ -102,14 +104,14 @@ public void fullFidelityChangeFeed_FromNowForLogicalPartition() throws Exception assertThat(itemChanges.get(1).get("metadata").get("operationType").asText()).isEqualTo("create"); // Assert replace of item1 assertThat(itemChanges.get(2).get("current").get("id").asText()).isEqualTo(item1.getId()); - assertThat(itemChanges.get(2).get("current").get("prop").asText()).isEqualTo(item1.getProp()); + assertThat(itemChanges.get(2).get("current").get("prop").asText()).isEqualTo(item1.getStringProp()); assertThat(itemChanges.get(2).get("metadata").get("operationType").asText()).isEqualTo("replace"); if (itemChanges.get(2).get("previous") != null) { assertThat(itemChanges.get(2).get("previous")).isEqualTo(itemChanges.get(0).get("current")); } // Assert replace of item2 assertThat(itemChanges.get(3).get("current").get("id").asText()).isEqualTo(item2.getId()); - assertThat(itemChanges.get(3).get("current").get("prop").asText()).isEqualTo(item2.getProp()); + assertThat(itemChanges.get(3).get("current").get("prop").asText()).isEqualTo(item2.getStringProp()); assertThat(itemChanges.get(3).get("metadata").get("operationType").asText()).isEqualTo("replace"); if (itemChanges.get(3).get("previous") != null) { assertThat(itemChanges.get(3).get("previous")).isEqualTo(itemChanges.get(1).get("current")); @@ -141,8 +143,8 @@ public void fullFidelityChangeFeed_FromNowForLogicalPartition() throws Exception } cosmosContainer.createItem(item3).block(); - String originalLastNameItem3 = item3.getProp(); - item3.setProp("Potter"); + String originalLastNameItem3 = item3.getStringProp(); + item3.setStringProp("Potter"); cosmosContainer.upsertItem(item3).block(); cosmosContainer.deleteItem(item3, new CosmosItemRequestOptions()).block(); @@ -166,7 +168,7 @@ public void fullFidelityChangeFeed_FromNowForLogicalPartition() throws Exception assertThat(itemChanges.get(0).get("metadata").get("operationType").asText()).isEqualTo("create"); // Assert replace of item3 assertThat(itemChanges.get(1).get("current").get("id").asText()).isEqualTo(item3.getId()); - assertThat(itemChanges.get(1).get("current").get("prop").asText()).isEqualTo(item3.getProp()); + assertThat(itemChanges.get(1).get("current").get("prop").asText()).isEqualTo(item3.getStringProp()); assertThat(itemChanges.get(1).get("metadata").get("operationType").asText()).isEqualTo("replace"); if (itemChanges.get(1).get("previous") != null) { assertThat(itemChanges.get(1).get("previous")).isEqualTo(itemChanges.get(0).get("current")); @@ -209,16 +211,16 @@ public void fullFidelityChangeFeed_FromContinuationToken() throws Exception { .createForProcessingFromContinuation(continuationToken); options.allVersionsAndDeletes(); - TestItem item1 = new TestItem( + TestObject item1 = new TestObject( UUID.randomUUID().toString(), - "mypk", "Johnson"); - TestItem item2 = new TestItem( + "mypk", Arrays.asList(), "Johnson"); + TestObject item2 = new TestObject( UUID.randomUUID().toString(), - "mypk", "Smith"); + "mypk", Arrays.asList(), "Smith"); cosmosContainer.upsertItem(item1).block(); cosmosContainer.upsertItem(item2).block(); - String originalLastName = item1.getProp(); - item1.setProp("Gates"); + String originalLastName = item1.getStringProp(); + item1.setStringProp("Gates"); cosmosContainer.upsertItem(item1).block(); cosmosContainer.deleteItem(item1, new CosmosItemRequestOptions()).block(); @@ -237,11 +239,11 @@ public void fullFidelityChangeFeed_FromContinuationToken() throws Exception { assertThat(itemChanges.get(0).get("current").get("prop").asText()).isEqualTo(originalLastName); assertThat(itemChanges.get(0).get("metadata").get("operationType").asText()).isEqualTo("create"); assertThat(itemChanges.get(1).get("current").get("id").asText()).isEqualTo(item2.getId()); - assertThat(itemChanges.get(1).get("current").get("prop").asText()).isEqualTo(item2.getProp()); + assertThat(itemChanges.get(1).get("current").get("prop").asText()).isEqualTo(item2.getStringProp()); assertThat(itemChanges.get(1).get("metadata").get("operationType").asText()).isEqualTo("create"); // Assert replace of item1 assertThat(itemChanges.get(2).get("current").get("id").asText()).isEqualTo(item1.getId()); - assertThat(itemChanges.get(2).get("current").get("prop").asText()).isEqualTo(item1.getProp()); + assertThat(itemChanges.get(2).get("current").get("prop").asText()).isEqualTo(item1.getStringProp()); assertThat(itemChanges.get(2).get("metadata").get("operationType").asText()).isEqualTo("replace"); if (itemChanges.get(2).get("previous") != null) { assertThat(itemChanges.get(2).get("previous")).isEqualTo(itemChanges.get(0).get("current")); @@ -286,9 +288,9 @@ public void fullFidelityChangeFeed_FromContinuationTokenOperationsOrder() throws // Create, replace, and delete 50 objects for 150 total operations for (int i = 0; i < 50; i++) { - TestItem currentItem = new TestItem("item"+ i, "mypk", "Smith"); + TestObject currentItem = new TestObject("item"+ i, "mypk", Arrays.asList(), "Smith"); cosmosContainer.upsertItem(currentItem).block(); - currentItem.setProp("Jefferson"); + currentItem.setStringProp("Jefferson"); cosmosContainer.upsertItem(currentItem).block(); cosmosContainer.deleteItem(currentItem, new CosmosItemRequestOptions()).block(); } @@ -342,18 +344,18 @@ public void fullFidelityChangeFeed_VerifyPreviousPresentOnReplace() throws Excep .createForProcessingFromContinuation(continuationToken); options.allVersionsAndDeletes(); - TestItem item1 = new TestItem( + TestObject item1 = new TestObject( UUID.randomUUID().toString(), - "mypk", "Johnson"); + "mypk", Arrays.asList(), "Johnson"); cosmosContainer.upsertItem(item1).block(); - String originalLastName = item1.getProp(); - item1.setProp("Gates"); + String originalLastName = item1.getStringProp(); + item1.setStringProp("Gates"); cosmosContainer.upsertItem(item1).block(); - String secondLastName = item1.getProp(); - item1.setProp("DiCaprio"); - String thirdLastName = item1.getProp(); + String secondLastName = item1.getStringProp(); + item1.setStringProp("DiCaprio"); + String thirdLastName = item1.getStringProp(); cosmosContainer.upsertItem(item1).block(); - item1.setProp(originalLastName); + item1.setStringProp(originalLastName); cosmosContainer.upsertItem(item1).block(); cosmosContainer.deleteItem(item1, new CosmosItemRequestOptions()).block(); @@ -387,7 +389,7 @@ public void fullFidelityChangeFeed_VerifyPreviousPresentOnReplace() throws Excep assertThat(itemChanges.get(2).get("previous")).isEqualTo(itemChanges.get(1).get("current")); assertThat(itemChanges.get(3).get("previous").get("id").asText()).isEqualTo(item1.getId()); - assertThat(itemChanges.get(3).get("current").get("prop").asText()).isEqualTo(item1.getProp()); + assertThat(itemChanges.get(3).get("current").get("prop").asText()).isEqualTo(item1.getStringProp()); assertThat(itemChanges.get(3).get("metadata").get("operationType").asText()).isEqualTo("replace"); assertThat(itemChanges.get(3).get("metadata").get("previousImageLSN").asText() ).isEqualTo(itemChanges.get(2).get("metadata").get("lsn").asText()); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index be75c1965a78..f34ac20355e3 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -2818,7 +2818,6 @@ public void enableGlobalThroughputControlGroup( this.enableGlobalThroughputControlGroup(groupConfig, globalControlConfig, null); } - /*** * Only used internally. *
@@ -2838,7 +2837,7 @@ void enableGlobalThroughputControlGroup( } /*** - * Enable the service throughput bucket control group. + * Enable the server throughput bucket control group. *

* For more information about throughput bucket please visit * Throughput buckets in Azure Cosmos DB diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java index 3892b0b5c314..00726d6f86e4 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java @@ -23,7 +23,6 @@ import com.azure.cosmos.implementation.query.PartitionedQueryExecutionInfo; import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroupInternal; -import com.azure.cosmos.implementation.throughputControl.server.config.ThroughputBucketControlGroup; import com.azure.cosmos.models.CosmosAuthorizationTokenResolver; import com.azure.cosmos.models.CosmosBatchResponse; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; @@ -1630,7 +1629,7 @@ Flux> readAllDocuments( CosmosItemSerializer getEffectiveItemSerializer(CosmosItemSerializer requestOptionsItemSerializer); /** - * Enable throughput control group. + * Enable sdk throughput control group. * * @param group the throughput control group. */ @@ -1638,6 +1637,7 @@ Flux> readAllDocuments( /*** * Enable server throughput control group. + * * @param group the server throughput control group. */ void enableServerThroughputControlGroup(ServerThroughputControlGroupInternal group); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 657ab2c3b691..ebe624934a81 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -6285,7 +6285,6 @@ private synchronized void enableThroughputControlStore() { } } - @Override public Flux submitOpenConnectionTasksAndInitCaches(CosmosContainerProactiveInitConfig proactiveContainerInitConfig) { return this.storeModel.submitOpenConnectionTasksAndInitCaches(proactiveContainerInitConfig); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/IThroughputContainerController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/IThroughputContainerController.java index 89edffd24e3d..7933ac3e040c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/IThroughputContainerController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/IThroughputContainerController.java @@ -3,10 +3,8 @@ package com.azure.cosmos.implementation.throughputControl; -import com.azure.cosmos.implementation.throughputControl.sdk.controller.IThroughputController; - /** * Represents a throughput container controller. */ -public abstract class IThroughputContainerController implements IThroughputController { +public interface IThroughputContainerController extends IThroughputController { } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/IThroughputController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/IThroughputController.java similarity index 94% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/IThroughputController.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/IThroughputController.java index 355e8a3a0401..c82398dcaa7f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/IThroughputController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/IThroughputController.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.sdk.controller; +package com.azure.cosmos.implementation.throughputControl; import com.azure.cosmos.implementation.RxDocumentServiceRequest; import reactor.core.publisher.Mono; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerSDKThroughputControlGroupProperties.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerSDKThroughputControlGroupProperties.java index 05c3180e5798..7975c2a8b9de 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerSDKThroughputControlGroupProperties.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerSDKThroughputControlGroupProperties.java @@ -97,7 +97,6 @@ public Pair enableThroughputControlGroup(SDKThroughputControlG return Pair.of(this.throughputControlGroups.size(), updatedGroupConfig.get()); } - // TODO: double check the data structure here public Map getThroughputControlGroups() { return this.throughputControlGroups; } @@ -120,9 +119,6 @@ public boolean allowRequestToContinueOnInitError(RxDocumentServiceRequest reques return this.supressInitErrorGroupSet.contains(requestGroupName); } - public boolean isNotEmpty() { - return !this.throughputControlGroups.isEmpty(); - } public boolean hasDefaultGroup() { return this.defaultGroup.get() != null; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/SDKThroughputControlStore.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/SDKThroughputControlStore.java index 38ae9f597ec8..4d4c61e00e90 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/SDKThroughputControlStore.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/SDKThroughputControlStore.java @@ -17,7 +17,7 @@ import com.azure.cosmos.implementation.throughputControl.EmptyThroughputContainerController; import com.azure.cosmos.implementation.throughputControl.IThroughputContainerController; import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; -import com.azure.cosmos.implementation.throughputControl.sdk.controller.IThroughputController; +import com.azure.cosmos.implementation.throughputControl.IThroughputController; import com.azure.cosmos.implementation.throughputControl.sdk.controller.container.SDKThroughputContainerController; import com.azure.cosmos.implementation.throughputControl.sdk.exceptions.ThroughputControlInitializationException; import org.slf4j.Logger; @@ -25,7 +25,6 @@ import reactor.core.Exceptions; import reactor.core.publisher.Mono; -import java.util.HashSet; import java.util.concurrent.ConcurrentHashMap; import static com.azure.cosmos.implementation.Exceptions.isNameCacheStale; @@ -105,9 +104,7 @@ public SDKThroughputControlStore( this.cancellationTokenMap = new ConcurrentHashMap<>(); } - public void enableThroughputControlGroup( - SDKThroughputControlGroupInternal group, - Mono throughputQueryMono) { + public void enableThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono throughputQueryMono) { checkNotNull(group, "Throughput control group cannot be null"); String containerNameLink = Utils.trimBeginningAndEndingSlashes(BridgeInternal.extractContainerSelfLink(group.getTargetContainer())); @@ -253,7 +250,7 @@ private Mono createAndInitContainerController(St new SDKThroughputContainerController( this.collectionCache, this.connectionMode, - new HashSet<>(throughputControlContainerProperties.getThroughputControlGroups().values()), + throughputControlContainerProperties.getThroughputControlGroups(), this.partitionKeyRangeCache, parentToken, throughputControlContainerProperties.getThroughputQueryMono()); @@ -300,26 +297,6 @@ private void handleException(String containerNameLink, RxDocumentServiceRequest } } - public boolean hasDefaultGroup(String containerNameLink) { - if (this.containerMap.containsKey(containerNameLink)) { - return this.containerMap.get(containerNameLink).hasDefaultGroup(); - } - - return false; - } - - public boolean hasGroup(String containerNameLink, String throughputControlGroupName) { - if (StringUtils.isEmpty(throughputControlGroupName)) { - return false; - } - - if (this.containerMap.containsKey(containerNameLink)) { - return this.containerMap.get(containerNameLink).hasGroup(throughputControlGroupName); - } - - return false; - } - public void close() { this.cancellationTokenSource.close(); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/SDKThroughputControlGroupInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/SDKThroughputControlGroupInternal.java index 3ca028425f17..6d393866eab2 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/SDKThroughputControlGroupInternal.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/SDKThroughputControlGroupInternal.java @@ -200,7 +200,6 @@ public boolean equals(Object other) { && Objects.equals(this.priorityLevel, that.priorityLevel); } - public boolean hasSameIdentity(Object other) { if (this == other) { return true; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/container/SDKThroughputContainerController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/container/SDKThroughputContainerController.java index eac4f7d6fcaa..497ec493cd1b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/container/SDKThroughputContainerController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/container/SDKThroughputContainerController.java @@ -41,7 +41,7 @@ import reactor.util.retry.RetrySpec; import java.time.Duration; -import java.util.Set; +import java.util.Map; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicReference; @@ -65,7 +65,7 @@ public class SDKThroughputContainerController implements IThroughputContainerCon private final RxCollectionCache collectionCache; private final ConnectionMode connectionMode; private final AsyncCache groupControllerCache; - private final Set groups; + private final Map groups; private final AtomicReference maxContainerThroughput; private final RxPartitionKeyRangeCache partitionKeyRangeCache; private final CosmosAsyncContainer targetContainer; @@ -82,7 +82,7 @@ public class SDKThroughputContainerController implements IThroughputContainerCon public SDKThroughputContainerController( RxCollectionCache collectionCache, ConnectionMode connectionMode, - Set groups, + Map groups, RxPartitionKeyRangeCache partitionKeyRangeCache, LinkedCancellationToken parentToken, Mono throughputQueryMono) { @@ -99,7 +99,7 @@ public SDKThroughputContainerController( this.maxContainerThroughput = new AtomicReference<>(); this.partitionKeyRangeCache = partitionKeyRangeCache; - this.targetContainer = groups.iterator().next().getTargetContainer(); + this.targetContainer = groups.values().iterator().next().getTargetContainer(); this.client = CosmosBridgeInternal.getContextClient(this.targetContainer); this.throughputProvisioningScope = this.getThroughputResolveLevel(groups); @@ -109,8 +109,8 @@ public SDKThroughputContainerController( this.throughputQueryMono = throughputQueryMono == null ? this.resolveContainerMaxThroughputCore() : throughputQueryMono; } - private ThroughputProvisioningScope getThroughputResolveLevel(Set groupConfigs) { - if (groupConfigs.stream().anyMatch(groupConfig -> groupConfig.getTargetThroughputThreshold() != null)) { + private ThroughputProvisioningScope getThroughputResolveLevel(Map groupConfigs) { + if (groupConfigs.values().stream().anyMatch(groupConfig -> groupConfig.getTargetThroughputThreshold() != null)) { // Throughput can be provisioned on container level or database level, will start from container return ThroughputProvisioningScope.CONTAINER; } else { @@ -312,15 +312,13 @@ private Mono> getOrCreateThroug return Mono.just(new Utils.ValueHolder<>(this.defaultGroupController)); } - for (SDKThroughputControlGroupInternal group : this.groups) { - if (StringUtils.equals(groupName, group.getGroupName())) { - return this.resolveThroughputGroupController(group) - .map(Utils.ValueHolder::new); - } + SDKThroughputControlGroupInternal group = this.groups.get(groupName); + if (group == null) { + // If the request is associated with a group not enabled, will fall back to the default one. + return Mono.just(new Utils.ValueHolder<>(this.defaultGroupController)); } - // If the request is associated with a group not enabled, will fall back to the default one. - return Mono.just(new Utils.ValueHolder<>(this.defaultGroupController)); + return this.resolveThroughputGroupController(group).map(Utils.ValueHolder::new); } public String getTargetContainerRid() { @@ -335,7 +333,7 @@ public boolean canHandleRequest(RxDocumentServiceRequest request) { } private Mono createAndInitializeGroupControllers() { - return Flux.fromIterable(this.groups) + return Flux.fromIterable(this.groups.values()) .flatMap(this::resolveThroughputGroupController) .then(Mono.just(this)); } @@ -388,7 +386,7 @@ private Flux refreshContainerMaxThroughputTask(LinkedCancellationToken can return this.resolveContainerMaxThroughput(); } }) - .flatMapIterable(controller -> this.groups) + .flatMapIterable(controller -> this.groups.values()) .flatMap(this::resolveThroughputGroupController) .doOnNext(groupController -> groupController.onContainerMaxThroughputRefresh(this.maxContainerThroughput.get())) .onErrorResume(throwable -> { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/ThroughputGroupControllerBase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/ThroughputGroupControllerBase.java index 473ad31ee0a4..31886ee017a0 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/ThroughputGroupControllerBase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/ThroughputGroupControllerBase.java @@ -14,7 +14,7 @@ import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationToken; import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationTokenSource; import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; -import com.azure.cosmos.implementation.throughputControl.sdk.controller.IThroughputController; +import com.azure.cosmos.implementation.throughputControl.IThroughputController; import com.azure.cosmos.implementation.throughputControl.sdk.controller.request.GlobalThroughputRequestController; import com.azure.cosmos.implementation.throughputControl.sdk.controller.request.IThroughputRequestController; import com.azure.cosmos.implementation.throughputControl.sdk.controller.request.PkRangesThroughputRequestController; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/request/IThroughputRequestController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/request/IThroughputRequestController.java index 56a66994ab8c..e2c0e541ceed 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/request/IThroughputRequestController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/request/IThroughputRequestController.java @@ -3,7 +3,7 @@ package com.azure.cosmos.implementation.throughputControl.sdk.controller.request; -import com.azure.cosmos.implementation.throughputControl.sdk.controller.IThroughputController; +import com.azure.cosmos.implementation.throughputControl.IThroughputController; /** * Represents a throughput request controller. diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ContainerServerThroughputControlGroupProperties.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ContainerServerThroughputControlGroupProperties.java index 6fd411ee854f..79d409f99670 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ContainerServerThroughputControlGroupProperties.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ContainerServerThroughputControlGroupProperties.java @@ -1,9 +1,10 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + package com.azure.cosmos.implementation.throughputControl.server; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroupInternal; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; @@ -13,7 +14,6 @@ import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; public class ContainerServerThroughputControlGroupProperties { - private static Logger logger = LoggerFactory.getLogger(ContainerServerThroughputControlGroupProperties.class); private final String containerNameLink; private final AtomicReference defaultGroup; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ServerThroughputControlStore.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ServerThroughputControlStore.java index 64309baa07f5..13fb0482ad13 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ServerThroughputControlStore.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ServerThroughputControlStore.java @@ -1,3 +1,6 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + package com.azure.cosmos.implementation.throughputControl.server; import com.azure.cosmos.BridgeInternal; @@ -9,8 +12,6 @@ import com.azure.cosmos.implementation.throughputControl.EmptyThroughputContainerController; import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroupInternal; import com.azure.cosmos.implementation.throughputControl.server.controller.ServerThroughputContainerController; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; import reactor.core.publisher.Mono; import java.util.concurrent.ConcurrentHashMap; @@ -19,7 +20,6 @@ import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; public class ServerThroughputControlStore { - private static final Logger logger = LoggerFactory.getLogger(ServerThroughputControlStore.class); private final AsyncCache containerControllerCache; private final ConcurrentHashMap containerMap; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputContainerController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputContainerController.java index 368d2a95246c..8ce7df52fa28 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputContainerController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputContainerController.java @@ -8,7 +8,6 @@ import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.implementation.caches.AsyncCache; import com.azure.cosmos.implementation.throughputControl.IThroughputContainerController; -import com.azure.cosmos.implementation.throughputControl.sdk.controller.IThroughputController; import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroupInternal; import com.azure.cosmos.implementation.throughputControl.server.config.ThroughputBucketControlGroup; import org.slf4j.Logger; @@ -21,7 +20,7 @@ import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; -public class ServerThroughputContainerController implements IThroughputController { +public class ServerThroughputContainerController implements IThroughputContainerController { private static final Logger logger = LoggerFactory.getLogger(ServerThroughputContainerController.class); private final AsyncCache groupControllerCache; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupControllerBase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupControllerBase.java index ce9948f38bf6..289d33622b9c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupControllerBase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupControllerBase.java @@ -3,7 +3,7 @@ package com.azure.cosmos.implementation.throughputControl.server.controller; -import com.azure.cosmos.implementation.throughputControl.sdk.controller.IThroughputController; +import com.azure.cosmos.implementation.throughputControl.IThroughputController; public abstract class ServerThroughputGroupControllerBase implements IThroughputController { } From 39d69bdeafa51f83350396d95fb6855dad3472e7 Mon Sep 17 00:00:00 2001 From: annie-mac Date: Wed, 16 Jul 2025 13:21:12 -0700 Subject: [PATCH 04/12] refactor --- .../azure/cosmos/CosmosAsyncContainer.java | 2 +- .../sdk/SDKThroughputControlStore.java | 20 +++++++++++++++++++ .../SDKThroughputContainerController.java | 16 +++++++-------- ... => SDKThroughputGroupControllerBase.java} | 6 +++--- .../ThroughputGroupControllerFactory.java | 2 +- ...lobalThroughputControlGroupController.java | 4 ++-- ...LocalThroughputControlGroupController.java | 4 ++-- 7 files changed, 37 insertions(+), 17 deletions(-) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/{ThroughputGroupControllerBase.java => SDKThroughputGroupControllerBase.java} (98%) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index f34ac20355e3..5f16bb9961c8 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -2844,7 +2844,7 @@ void enableGlobalThroughputControlGroup( * * @param groupConfig the throughput control group config, see {@link ThroughputControlGroupConfig}. */ - public void enableThroughputBucketControlGroup(ThroughputControlGroupConfig groupConfig) { + public void enableServerThroughputControlGroup(ThroughputControlGroupConfig groupConfig) { ThroughputBucketControlGroup throughputBucketControlGroup = ThroughputControlGroupFactory.createThroughputBucketControlGroup(groupConfig, this); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/SDKThroughputControlStore.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/SDKThroughputControlStore.java index 4d4c61e00e90..412413fe547d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/SDKThroughputControlStore.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/SDKThroughputControlStore.java @@ -297,6 +297,26 @@ private void handleException(String containerNameLink, RxDocumentServiceRequest } } + public boolean hasDefaultGroup(String containerNameLink) { + if (this.containerMap.containsKey(containerNameLink)) { + return this.containerMap.get(containerNameLink).hasDefaultGroup(); + } + + return false; + } + + public boolean hasGroup(String containerNameLink, String throughputControlGroupName) { + if (throughputControlGroupName.isEmpty()) { + return false; + } + + if (this.containerMap.containsKey(containerNameLink)) { + return this.containerMap.get(containerNameLink).hasGroup(throughputControlGroupName); + } + + return false; + } + public void close() { this.cancellationTokenSource.close(); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/container/SDKThroughputContainerController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/container/SDKThroughputContainerController.java index 497ec493cd1b..eb688dd297ec 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/container/SDKThroughputContainerController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/container/SDKThroughputContainerController.java @@ -26,7 +26,7 @@ import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationToken; import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationTokenSource; import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; -import com.azure.cosmos.implementation.throughputControl.sdk.controller.group.ThroughputGroupControllerBase; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.group.SDKThroughputGroupControllerBase; import com.azure.cosmos.implementation.throughputControl.sdk.controller.group.ThroughputGroupControllerFactory; import com.azure.cosmos.implementation.throughputControl.sdk.exceptions.ThroughputControlInitializationException; import com.azure.cosmos.models.CosmosQueryRequestOptions; @@ -64,7 +64,7 @@ public class SDKThroughputContainerController implements IThroughputContainerCon private final AsyncDocumentClient client; private final RxCollectionCache collectionCache; private final ConnectionMode connectionMode; - private final AsyncCache groupControllerCache; + private final AsyncCache groupControllerCache; private final Map groups; private final AtomicReference maxContainerThroughput; private final RxPartitionKeyRangeCache partitionKeyRangeCache; @@ -74,7 +74,7 @@ public class SDKThroughputContainerController implements IThroughputContainerCon private final ConcurrentHashMap cancellationTokenMap; private final Mono throughputQueryMono; - private ThroughputGroupControllerBase defaultGroupController; + private SDKThroughputGroupControllerBase defaultGroupController; private String targetContainerRid; private String targetDatabaseRid; private ThroughputProvisioningScope throughputProvisioningScope; @@ -305,7 +305,7 @@ public Mono processRequest(RxDocumentServiceRequest request, Mono orig } // TODO: a better way to handle throughput control group enabled after the container initialization - private Mono> getOrCreateThroughputGroupController(String groupName) { + private Mono> getOrCreateThroughputGroupController(String groupName) { // If there is no control group defined, using the default group controller if (StringUtils.isEmpty(groupName)) { @@ -338,7 +338,7 @@ private Mono createAndInitializeGroupControlle .then(Mono.just(this)); } - private Mono resolveThroughputGroupController(SDKThroughputControlGroupInternal group) { + private Mono resolveThroughputGroupController(SDKThroughputControlGroupInternal group) { return this.groupControllerCache.getAsync( group.getGroupName(), null, @@ -346,13 +346,13 @@ private Mono resolveThroughputGroupController(SDK .onErrorResume(throwable -> Mono.error(new ThroughputControlInitializationException(throwable))); } - private Mono createAndInitializeGroupController(SDKThroughputControlGroupInternal group) { + private Mono createAndInitializeGroupController(SDKThroughputControlGroupInternal group) { LinkedCancellationToken parentToken = this.cancellationTokenMap.compute( group.getGroupName(), (key, cancellationToken) -> this.cancellationTokenSource.getToken()); - ThroughputGroupControllerBase groupController = ThroughputGroupControllerFactory.createController( + SDKThroughputGroupControllerBase groupController = ThroughputGroupControllerFactory.createController( this.connectionMode, group, this.maxContainerThroughput.get(), @@ -362,7 +362,7 @@ private Mono createAndInitializeGroupController(S return groupController .init() - .cast(ThroughputGroupControllerBase.class) + .cast(SDKThroughputGroupControllerBase.class) .doOnSuccess(controller -> { if (controller.isDefault()) { this.defaultGroupController = controller; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/ThroughputGroupControllerBase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/SDKThroughputGroupControllerBase.java similarity index 98% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/ThroughputGroupControllerBase.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/SDKThroughputGroupControllerBase.java index 31886ee017a0..5bdc9aaed368 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/ThroughputGroupControllerBase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/SDKThroughputGroupControllerBase.java @@ -39,8 +39,8 @@ * 1. Create and initialize request controller based on connection mode * 2. Schedule reset throughput usage every 1s. */ -public abstract class ThroughputGroupControllerBase implements IThroughputController { - private final static Logger logger = LoggerFactory.getLogger(ThroughputGroupControllerBase.class); +public abstract class SDKThroughputGroupControllerBase implements IThroughputController { + private final static Logger logger = LoggerFactory.getLogger(SDKThroughputGroupControllerBase.class); private final Duration DEFAULT_THROUGHPUT_USAGE_RESET_DURATION = Duration.ofSeconds(1); private final ConnectionMode connectionMode; @@ -53,7 +53,7 @@ public abstract class ThroughputGroupControllerBase implements IThroughputContro protected final AtomicReference groupThroughput; protected final LinkedCancellationTokenSource cancellationTokenSource; - public ThroughputGroupControllerBase( + public SDKThroughputGroupControllerBase( ConnectionMode connectionMode, SDKThroughputControlGroupInternal group, Integer maxContainerThroughput, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/ThroughputGroupControllerFactory.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/ThroughputGroupControllerFactory.java index 12646e32933c..7a26cf4ef508 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/ThroughputGroupControllerFactory.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/ThroughputGroupControllerFactory.java @@ -14,7 +14,7 @@ public class ThroughputGroupControllerFactory { - public static ThroughputGroupControllerBase createController( + public static SDKThroughputGroupControllerBase createController( ConnectionMode connectionMode, SDKThroughputControlGroupInternal group, Integer maxContainerThroughput, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlGroupController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlGroupController.java index 53be0743edfc..0beebbcd8105 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlGroupController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlGroupController.java @@ -9,7 +9,7 @@ import com.azure.cosmos.implementation.guava25.collect.EvictingQueue; import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationToken; import com.azure.cosmos.implementation.throughputControl.sdk.config.GlobalThroughputControlGroup; -import com.azure.cosmos.implementation.throughputControl.sdk.controller.group.ThroughputGroupControllerBase; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.group.SDKThroughputGroupControllerBase; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; @@ -19,7 +19,7 @@ import java.time.Instant; import java.util.concurrent.atomic.AtomicReference; -public class GlobalThroughputControlGroupController extends ThroughputGroupControllerBase { +public class GlobalThroughputControlGroupController extends SDKThroughputGroupControllerBase { private static final Logger logger = LoggerFactory.getLogger(GlobalThroughputControlGroupController.class); private static final double INITIAL_CLIENT_THROUGHPUT_RU_SHARE = 1.0; private static final double INITIAL_THROUGHPUT_USAGE = 1.0; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/local/LocalThroughputControlGroupController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/local/LocalThroughputControlGroupController.java index 16def2b5b221..5e9498fe5924 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/local/LocalThroughputControlGroupController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/local/LocalThroughputControlGroupController.java @@ -8,10 +8,10 @@ import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationToken; import com.azure.cosmos.implementation.throughputControl.sdk.config.LocalThroughputControlGroup; -import com.azure.cosmos.implementation.throughputControl.sdk.controller.group.ThroughputGroupControllerBase; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.group.SDKThroughputGroupControllerBase; import reactor.core.publisher.Mono; -public class LocalThroughputControlGroupController extends ThroughputGroupControllerBase { +public class LocalThroughputControlGroupController extends SDKThroughputGroupControllerBase { public LocalThroughputControlGroupController( ConnectionMode connectionMode, From 6e3025dc88fe2f2338793cfb4f97402adca3fbe3 Mon Sep 17 00:00:00 2001 From: annie-mac Date: Thu, 17 Jul 2025 10:45:46 -0700 Subject: [PATCH 05/12] fix tests --- .../sdk/ThroughputControlGroupConfigConfigurationTests.java | 2 +- .../directconnectivity/rntbd/RntbdRequestHeaders.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputControlGroupConfigConfigurationTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputControlGroupConfigConfigurationTests.java index 26e0de3144b0..61488d7f9d66 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputControlGroupConfigConfigurationTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputControlGroupConfigConfigurationTests.java @@ -96,7 +96,7 @@ public void validatePriorityThroughputControlGroupWithThreshold() { public void validateThroughputControlGroupMalformed() { assertThatThrownBy(() -> new ThroughputControlGroupConfigBuilder().groupName("test").build()) .isInstanceOf(IllegalArgumentException.class) - .hasMessage("All targetThroughput, targetThroughputThreshold and priorityLevel cannot be null or empty."); + .hasMessage("All targetThroughput, targetThroughputThreshold, priorityLevel and throughput bucket cannot be null or empty."); } @BeforeClass(groups = { "emulator" }, timeOut = 4 * SETUP_TIMEOUT) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java index 7c05b97387a4..0c3f9615312f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java @@ -801,7 +801,7 @@ private void addThroughputBucket(final Map headers) if (StringUtils.isNotEmpty(value)) { final int throughputBucket = Integer.valueOf(value); - this.getThroughputBucket().setValue(throughputBucket); + this.getThroughputBucket().setValue((byte)throughputBucket); } } From 708dce197c33d3c4a3b7fefe4103ffcbb1a03935 Mon Sep 17 00:00:00 2001 From: annie-mac Date: Sat, 19 Jul 2025 21:20:33 -0700 Subject: [PATCH 06/12] refactor --- .../throughputControl/{sdk => }/TestItem.java | 2 +- .../{sdk => }/ThroughputControlTests.java | 105 ++++++++++++++++- ...hroughputControlGroupPropertiesTests.java} | 2 +- ...ThroughputControlGroupPropertiesTests.java | 106 ++++++++++++++++++ .../rx/WebExceptionRetryPolicyE2ETests.java | 2 +- .../com/azure/cosmos/CosmosAsyncClient.java | 4 +- .../azure/cosmos/CosmosAsyncContainer.java | 12 +- .../com/azure/cosmos/CosmosContainer.java | 12 ++ .../implementation/AsyncDocumentClient.java | 4 +- .../implementation/RxDocumentClientImpl.java | 4 +- .../ThroughputControlGroupFactory.java | 6 +- .../ThroughputControlStore.java | 9 +- .../sdk/SDKThroughputControlStore.java | 2 +- ...erverThroughputControlGroupProperties.java | 14 +-- .../server/ServerThroughputControlStore.java | 6 +- ...java => ServerThroughputControlGroup.java} | 29 +++-- .../config/ThroughputBucketControlGroup.java | 45 -------- .../ServerThroughputContainerController.java | 25 ++--- ...a => ServerThroughputGroupController.java} | 23 ++-- .../ServerThroughputGroupControllerBase.java | 9 -- 20 files changed, 299 insertions(+), 122 deletions(-) rename sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/{sdk => }/TestItem.java (93%) rename sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/{sdk => }/ThroughputControlTests.java (89%) rename sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/{ContainerThroughputControlGroupPropertiesTests.java => ContainerSDKThroughputControlGroupPropertiesTests.java} (99%) create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/server/ContainerServerThroughputControlGroupPropertiesTests.java rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/{ServerThroughputControlGroupInternal.java => ServerThroughputControlGroup.java} (69%) delete mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ThroughputBucketControlGroup.java rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/{ThroughputBucketGroupController.java => ServerThroughputGroupController.java} (51%) delete mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupControllerBase.java diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/TestItem.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/TestItem.java similarity index 93% rename from sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/TestItem.java rename to sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/TestItem.java index cc9b9422e903..8e03339f67c3 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/TestItem.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/TestItem.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.sdk; +package com.azure.cosmos.implementation.throughputControl; import java.util.UUID; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputControlTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java similarity index 89% rename from sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputControlTests.java rename to sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java index be31747b7291..838210d19656 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputControlTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.sdk; +package com.azure.cosmos.implementation.throughputControl; import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.ConnectionMode; @@ -52,6 +52,7 @@ import java.util.concurrent.atomic.AtomicInteger; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.AssertionsForClassTypes.assertThatThrownBy; import static org.testng.Assert.fail; public class ThroughputControlTests extends TestSuiteBase { @@ -670,8 +671,106 @@ public void enableSameGroupMultipleTimes() { } } + @Test(groups = {"emulator"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT) + public void serverThroughputControl(OperationType operationType) { + // TODO: currently there is no easy way to do e2e testing, so the testing here is to just verify that + // server throughput can be enabled on the container + this.ensureContainer(); + + ThroughputControlGroupConfig serverThroughputControlGroup = + new ThroughputControlGroupConfigBuilder() + .groupName("serverThroughputControl") + .throughputBucket(2) + .build(); + container.enableServerThroughputControlGroup(serverThroughputControlGroup); + + CosmosItemRequestOptions requestOptions = new CosmosItemRequestOptions(); + requestOptions.setContentResponseOnWriteEnabled(true); + requestOptions.setThroughputControlGroupName(serverThroughputControlGroup.getGroupName()); + + CosmosItemResponse createItemResponse = container.createItem(getDocumentDefinition(), requestOptions).block(); + TestItem createdItem = createItemResponse.getItem(); + this.validateRequestNotThrottled( + createItemResponse.getDiagnostics().toString(), + BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + + performDocumentOperation( + this.container, + operationType, + createdItem, + serverThroughputControlGroup.getGroupName()); + } + + @Test(groups = {"emulator"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT) + public void throughputControl_LocalAndServer_requestOptions(OperationType operationType) { + this.ensureContainer(); + + // This test is verify that SDK throughput control group and server throughput control group can be enabled at the same time + + // The create document in this test usually takes around 6.29RU, pick a RU here relatively close, so to test throttled scenario + ThroughputControlGroupConfig groupConfig = + new ThroughputControlGroupConfigBuilder() + .groupName("group-sdk" + UUID.randomUUID()) + .targetThroughput(6) + .build(); + container.enableLocalThroughputControlGroup(groupConfig); + + ThroughputControlGroupConfig serverGroupConfig = + new ThroughputControlGroupConfigBuilder() + .groupName("group-server" + UUID.randomUUID()) + .throughputBucket(3) + .build(); + container.enableServerThroughputControlGroup(serverGroupConfig); + + CosmosItemRequestOptions requestOptions = new CosmosItemRequestOptions(); + requestOptions.setContentResponseOnWriteEnabled(true); + requestOptions.setThroughputControlGroupName(groupConfig.getGroupName()); + + CosmosItemResponse createItemResponse = container.createItem(getDocumentDefinition(), requestOptions).block(); + TestItem createdItem = createItemResponse.getItem(); + this.validateRequestNotThrottled( + createItemResponse.getDiagnostics().toString(), + BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + + // second request to group-1. which will get throttled + CosmosDiagnostics cosmosDiagnostics = performDocumentOperation(this.container, operationType, createdItem, groupConfig.getGroupName()); + this.validateRequestThrottled( + cosmosDiagnostics.toString(), + BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + + // third request to server group, which will not get throttled + requestOptions.setThroughputControlGroupName(serverGroupConfig.getGroupName()); + performDocumentOperation(this.container, operationType, createdItem, groupConfig.getGroupName()); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void throughputControlDefaultGroup_LocalAndServer_requestOptions() { + this.ensureContainer(); + + // This test is verify that only one default throughput control group can be defined across sdk and server control group + ThroughputControlGroupConfig groupConfig = + new ThroughputControlGroupConfigBuilder() + .groupName("group-sdk" + UUID.randomUUID()) + .targetThroughput(6) + .defaultControlGroup(true) + .build(); + container.enableLocalThroughputControlGroup(groupConfig); + + ThroughputControlGroupConfig serverGroupConfig = + new ThroughputControlGroupConfigBuilder() + .groupName("group-server" + UUID.randomUUID()) + .throughputBucket(3) + .defaultControlGroup(true) + .build(); + + assertThatThrownBy( + () ->container.enableServerThroughputControlGroup(serverGroupConfig)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("A default group already exists"); + } + @BeforeClass(groups = { "emulator" }, timeOut = 4 * SETUP_TIMEOUT) - public void before_ThroughputBudgetControllerTest() { + public void before_ThroughputControllerTest() { this.ensureContainer(); } @@ -698,7 +797,7 @@ private void ensureContainer() { } @AfterClass(groups = {"emulator"}, timeOut = TIMEOUT, alwaysRun = true) - public void after_ThroughputBudgetControllerTest() { + public void after_ThroughputControllerTest() { safeClose(this.client); } diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerThroughputControlGroupPropertiesTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerSDKThroughputControlGroupPropertiesTests.java similarity index 99% rename from sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerThroughputControlGroupPropertiesTests.java rename to sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerSDKThroughputControlGroupPropertiesTests.java index 592e01afb58e..a4a3e94fd25e 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerThroughputControlGroupPropertiesTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerSDKThroughputControlGroupPropertiesTests.java @@ -19,7 +19,7 @@ import static org.assertj.core.api.AssertionsForClassTypes.assertThat; import static org.assertj.core.api.AssertionsForClassTypes.assertThatThrownBy; -public class ContainerThroughputControlGroupPropertiesTests { +public class ContainerSDKThroughputControlGroupPropertiesTests { @Test(groups = "emulator") public void enableThroughputControlGroup() { diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/server/ContainerServerThroughputControlGroupPropertiesTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/server/ContainerServerThroughputControlGroupPropertiesTests.java new file mode 100644 index 000000000000..e52f18d16598 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/server/ContainerServerThroughputControlGroupPropertiesTests.java @@ -0,0 +1,106 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl.server; + +import com.azure.cosmos.CosmosAsyncClient; +import com.azure.cosmos.CosmosAsyncContainer; +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.implementation.TestConfigurations; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroup; +import com.azure.cosmos.models.PriorityLevel; +import org.testng.annotations.Test; + +import java.util.UUID; + +import static org.assertj.core.api.AssertionsForClassTypes.assertThat; +import static org.assertj.core.api.AssertionsForClassTypes.assertThatThrownBy; + +public class ContainerServerThroughputControlGroupPropertiesTests { + @Test(groups = "emulator") + public void enableThroughputControlGroup() { + CosmosAsyncClient testClient = null; + try { + testClient = new CosmosClientBuilder() + .endpoint(TestConfigurations.HOST) + .key(TestConfigurations.MASTER_KEY) + .buildAsyncClient(); + + ContainerServerThroughputControlGroupProperties throughputControlContainerProperties = + new ContainerServerThroughputControlGroupProperties("/testDB/testContainer"); + + CosmosAsyncContainer container = testClient.getDatabase("fakeDatabase").getContainer("fakeContainer"); + + // Test 1: add default throughput control group successfully + ServerThroughputControlGroup throughputControlDefaultGroup = new ServerThroughputControlGroup( + "test-" + UUID.randomUUID(), + true, + PriorityLevel.HIGH, + 1, + container); + + Integer currentGroupSize = + throughputControlContainerProperties.enableThroughputControlGroup(throughputControlDefaultGroup); + assertThat(currentGroupSize).isEqualTo(1); + assertThat(throughputControlContainerProperties.hasDefaultGroup()).isTrue(); + assertThat(throughputControlContainerProperties.hasGroup(throughputControlDefaultGroup.getGroupName())).isTrue(); + + // Test 2: add another default throughput control group + ServerThroughputControlGroup throughputControlDefaultGroup2 = new ServerThroughputControlGroup( + "test-" + UUID.randomUUID(), + true, + PriorityLevel.HIGH, + 2, + container); + + assertThatThrownBy( + () -> throughputControlContainerProperties + .enableThroughputControlGroup(throughputControlDefaultGroup2)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("A default group already exists"); + assertThat(throughputControlContainerProperties.hasGroup(throughputControlDefaultGroup2.getGroupName())).isFalse(); + + // Test 3: add a new non-default group + ServerThroughputControlGroup throughputControlGroup1 = new ServerThroughputControlGroup( + "test-" + UUID.randomUUID(), + false, + PriorityLevel.HIGH, + 1, + container); + + currentGroupSize = throughputControlContainerProperties.enableThroughputControlGroup(throughputControlGroup1); + assertThat(currentGroupSize).isEqualTo(2); + assertThat(throughputControlContainerProperties.hasGroup(throughputControlGroup1.getGroupName())).isFalse(); + + // Test 4: add a throughput control group with name as throughputControlGroup1 but with different config + ServerThroughputControlGroup throughputControlGroup2 = new ServerThroughputControlGroup( + throughputControlGroup1.getGroupName(), + false, + PriorityLevel.HIGH, + 2, + container); + assertThatThrownBy( + () -> throughputControlContainerProperties + .enableThroughputControlGroup(throughputControlGroup2)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("A group with same name already exists, name: " + throughputControlGroup1.getGroupName()); + assertThat(throughputControlContainerProperties.hasGroup(throughputControlGroup2.getGroupName())).isFalse(); + + // Test 5: add a throughput control group with same config as throughputControlGroup1, to verify no errors will be thrown + ServerThroughputControlGroup throughputControlGroup3 = new ServerThroughputControlGroup( + throughputControlGroup1.getGroupName(), + throughputControlGroup1.isDefault(), + throughputControlGroup1.getPriorityLevel(), + throughputControlGroup1.getThroughputBucket(), + container); + + currentGroupSize = throughputControlContainerProperties.enableThroughputControlGroup(throughputControlGroup3); + assertThat(currentGroupSize).isEqualTo(2); + + } finally { + if (testClient != null) { + testClient.close(); + } + } + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/WebExceptionRetryPolicyE2ETests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/WebExceptionRetryPolicyE2ETests.java index 8ada8f9df8db..bc544fbc855d 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/WebExceptionRetryPolicyE2ETests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/WebExceptionRetryPolicyE2ETests.java @@ -16,7 +16,7 @@ import com.azure.cosmos.implementation.ImplementationBridgeHelpers; import com.azure.cosmos.implementation.OperationType; import com.azure.cosmos.implementation.RequestTimeline; -import com.azure.cosmos.implementation.throughputControl.sdk.TestItem; +import com.azure.cosmos.implementation.throughputControl.TestItem; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index 106038f82fc3..680b98a6a915 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java @@ -30,7 +30,7 @@ import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdMetrics; import com.azure.cosmos.implementation.faultinjection.IFaultInjectorProvider; import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; -import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroup; import com.azure.cosmos.models.CosmosAuthorizationTokenResolver; import com.azure.cosmos.models.CosmosClientTelemetryConfig; import com.azure.cosmos.models.CosmosContainerIdentity; @@ -583,7 +583,7 @@ void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal group, Mo * * @param group the server throughput control group. */ - void enableServerThroughputControlGroup(ServerThroughputControlGroupInternal group) { + void enableServerThroughputControlGroup(ServerThroughputControlGroup group) { checkNotNull(group, "Argument 'group' can not be null"); this.asyncDocumentClient.enableServerThroughputControlGroup(group); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 5f16bb9961c8..3d40ad1b81a7 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -39,7 +39,7 @@ import com.azure.cosmos.implementation.routing.Range; import com.azure.cosmos.implementation.throughputControl.sdk.config.GlobalThroughputControlGroup; import com.azure.cosmos.implementation.throughputControl.sdk.config.LocalThroughputControlGroup; -import com.azure.cosmos.implementation.throughputControl.server.config.ThroughputBucketControlGroup; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroup; import com.azure.cosmos.implementation.throughputControl.ThroughputControlGroupFactory; import com.azure.cosmos.models.CosmosBatch; import com.azure.cosmos.models.CosmosBatchOperationResult; @@ -2845,10 +2845,14 @@ void enableGlobalThroughputControlGroup( * @param groupConfig the throughput control group config, see {@link ThroughputControlGroupConfig}. */ public void enableServerThroughputControlGroup(ThroughputControlGroupConfig groupConfig) { - ThroughputBucketControlGroup throughputBucketControlGroup = - ThroughputControlGroupFactory.createThroughputBucketControlGroup(groupConfig, this); + if (groupConfig.getPriorityLevel() == null && groupConfig.getThroughputBucket() == null) { + throw new IllegalArgumentException("Config 'priorityLevel' and 'throughputBucket' can not both are null."); + } + + ServerThroughputControlGroup serverThroughputControlGroup = + ThroughputControlGroupFactory.createServerThroughputControlGroup(groupConfig, this); - this.database.getClient().enableServerThroughputControlGroup(throughputBucketControlGroup); + this.database.getClient().enableServerThroughputControlGroup(serverThroughputControlGroup); } void configureFaultInjectionProvider(IFaultInjectorProvider injectorProvider) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java index 2702a24377b0..1b7935aaceeb 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java @@ -1030,6 +1030,18 @@ public void enableGlobalThroughputControlGroup(ThroughputControlGroupConfig grou this.asyncContainer.enableGlobalThroughputControlGroup(groupConfig, globalControlConfig); } + /*** + * Enable the server throughput bucket control group. + *

+ * For more information about throughput bucket please visit + * Throughput buckets in Azure Cosmos DB + * + * @param groupConfig the throughput control group config, see {@link ThroughputControlGroupConfig}. + */ + public void enableServerThroughputControlGroup(ThroughputControlGroupConfig groupConfig) { + this.asyncContainer.enableServerThroughputControlGroup(groupConfig); + } + /** * Initializes the container by warming up the caches and connections for the current read region. * diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java index 00726d6f86e4..7761a588b4a6 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java @@ -22,7 +22,7 @@ import com.azure.cosmos.implementation.faultinjection.IFaultInjectorProvider; import com.azure.cosmos.implementation.query.PartitionedQueryExecutionInfo; import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; -import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroup; import com.azure.cosmos.models.CosmosAuthorizationTokenResolver; import com.azure.cosmos.models.CosmosBatchResponse; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; @@ -1640,7 +1640,7 @@ Flux> readAllDocuments( * * @param group the server throughput control group. */ - void enableServerThroughputControlGroup(ServerThroughputControlGroupInternal group); + void enableServerThroughputControlGroup(ServerThroughputControlGroup group); /** * Submits open connection tasks and warms up caches for replicas for containers specified by diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index ebe624934a81..0aa24f5b1086 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -70,7 +70,7 @@ import com.azure.cosmos.implementation.spark.OperationListener; import com.azure.cosmos.implementation.throughputControl.ThroughputControlStore; import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; -import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroup; import com.azure.cosmos.models.CosmosAuthorizationTokenResolver; import com.azure.cosmos.models.CosmosBatchResponse; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; @@ -6262,7 +6262,7 @@ public void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal gr } @Override - public void enableServerThroughputControlGroup(ServerThroughputControlGroupInternal group) { + public void enableServerThroughputControlGroup(ServerThroughputControlGroup group) { checkNotNull(group, "Argument 'group' can not be null"); this.enableThroughputControlStore(); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupFactory.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupFactory.java index f91cc084756f..0a52d45b831b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupFactory.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupFactory.java @@ -9,7 +9,7 @@ import com.azure.cosmos.GlobalThroughputControlConfig; import com.azure.cosmos.implementation.throughputControl.sdk.config.GlobalThroughputControlGroup; import com.azure.cosmos.implementation.throughputControl.sdk.config.LocalThroughputControlGroup; -import com.azure.cosmos.implementation.throughputControl.server.config.ThroughputBucketControlGroup; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroup; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; @@ -52,14 +52,14 @@ public static GlobalThroughputControlGroup createThroughputGlobalControlGroup( } - public static ThroughputBucketControlGroup createThroughputBucketControlGroup( + public static ServerThroughputControlGroup createServerThroughputControlGroup( ThroughputControlGroupConfig groupConfig, CosmosAsyncContainer targetContainer) { checkNotNull(groupConfig, "Throughput control group config can not be null"); checkNotNull(targetContainer, "Throughput target container can not be null"); - return new ThroughputBucketControlGroup( + return new ServerThroughputControlGroup( groupConfig.getGroupName(), groupConfig.isDefault(), groupConfig.getPriorityLevel(), diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlStore.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlStore.java index e34f2ec04caf..250dbd466ae7 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlStore.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlStore.java @@ -13,7 +13,7 @@ import com.azure.cosmos.implementation.throughputControl.sdk.SDKThroughputControlStore; import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; import com.azure.cosmos.implementation.throughputControl.server.ServerThroughputControlStore; -import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroup; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Mono; @@ -52,12 +52,13 @@ public void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal gr this.sdkThroughputControlStore.enableThroughputControlGroup(group, throughputQueryMono); } - public void enableServerThroughputControlGroup(ServerThroughputControlGroupInternal group) { + public void enableServerThroughputControlGroup(ServerThroughputControlGroup group) { checkNotNull(group, "Throughput control group cannot be null"); if (group.isDefault()) { //verify a default group is not being defined in the sdk throughput control store - String containerNameLink = Utils.trimBeginningAndEndingSlashes(BridgeInternal.extractContainerSelfLink(group.getTargetContainer())); + String containerNameLink = + Utils.trimBeginningAndEndingSlashes(BridgeInternal.extractContainerSelfLink(group.getTargetContainer())); if (this.sdkThroughputControlStore.hasDefaultGroup(containerNameLink)) { throw new IllegalArgumentException("A default group already exists"); } @@ -91,7 +92,7 @@ public Mono processRequest(RxDocumentServiceRequest request, Mono orig } if (this.sdkThroughputControlStore.hasDefaultGroup(collectionNameLink)) { - return this.serverThroughputControlStore.processRequest(request, originalRequestMono); + return this.sdkThroughputControlStore.processRequest(request, originalRequestMono); } // neither store can process the request, fallback to just use the original mono diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/SDKThroughputControlStore.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/SDKThroughputControlStore.java index 412413fe547d..d48cf3262d60 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/SDKThroughputControlStore.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/SDKThroughputControlStore.java @@ -306,7 +306,7 @@ public boolean hasDefaultGroup(String containerNameLink) { } public boolean hasGroup(String containerNameLink, String throughputControlGroupName) { - if (throughputControlGroupName.isEmpty()) { + if (StringUtils.isEmpty(throughputControlGroupName)) { return false; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ContainerServerThroughputControlGroupProperties.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ContainerServerThroughputControlGroupProperties.java index 79d409f99670..472302f6e405 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ContainerServerThroughputControlGroupProperties.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ContainerServerThroughputControlGroupProperties.java @@ -4,7 +4,7 @@ package com.azure.cosmos.implementation.throughputControl.server; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; -import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroup; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; @@ -16,8 +16,8 @@ public class ContainerServerThroughputControlGroupProperties { private final String containerNameLink; - private final AtomicReference defaultGroup; - private final Map throughputControlGroups; + private final AtomicReference defaultGroup; + private final Map throughputControlGroups; public ContainerServerThroughputControlGroupProperties(String containerNameLink) { checkArgument(StringUtils.isNotEmpty(containerNameLink), "Argument 'containerNameLink' should not be empty"); @@ -30,11 +30,11 @@ public ContainerServerThroughputControlGroupProperties(String containerNameLink) /*** * Enable a server throughput control group. * - * @param group a {@link ServerThroughputControlGroupInternal}. + * @param group a {@link ServerThroughputControlGroup}. * * @return the total size of distinct server throughput control groups enabled on the container. */ - public int enableThroughputControlGroup(ServerThroughputControlGroupInternal group) { + public int enableThroughputControlGroup(ServerThroughputControlGroup group) { checkNotNull(group, "Argument 'group' should not be null'"); if (group.isDefault()) { @@ -45,7 +45,7 @@ public int enableThroughputControlGroup(ServerThroughputControlGroupInternal gro } } - ServerThroughputControlGroupInternal serverThroughputControlGroup = + ServerThroughputControlGroup serverThroughputControlGroup = this.throughputControlGroups.computeIfAbsent(group.getGroupName(), (key) -> group); if (!serverThroughputControlGroup.equals(group)) { throw new IllegalArgumentException("A group with same name already exists, name: " + group.getGroupName()); @@ -54,7 +54,7 @@ public int enableThroughputControlGroup(ServerThroughputControlGroupInternal gro return this.throughputControlGroups.size(); } - public Map getThroughputControlGroups() { + public Map getThroughputControlGroups() { return this.throughputControlGroups; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ServerThroughputControlStore.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ServerThroughputControlStore.java index 13fb0482ad13..1146e2e72bf9 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ServerThroughputControlStore.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ServerThroughputControlStore.java @@ -10,7 +10,7 @@ import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.implementation.caches.AsyncCache; import com.azure.cosmos.implementation.throughputControl.EmptyThroughputContainerController; -import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroup; import com.azure.cosmos.implementation.throughputControl.server.controller.ServerThroughputContainerController; import reactor.core.publisher.Mono; @@ -29,7 +29,7 @@ public ServerThroughputControlStore() { this.containerMap = new ConcurrentHashMap<>(); } - public void enableThroughputControlGroup(ServerThroughputControlGroupInternal group) { + public void enableThroughputControlGroup(ServerThroughputControlGroup group) { checkNotNull(group, "Throughput control group cannot be null"); String containerNameLink = Utils.trimBeginningAndEndingSlashes(BridgeInternal.extractContainerSelfLink(group.getTargetContainer())); @@ -74,7 +74,7 @@ public boolean hasDefaultGroup(String containerNameLink) { } public boolean hasGroup(String containerNameLink, String throughputControlGroupName) { - if (throughputControlGroupName.isEmpty()) { + if (StringUtils.isEmpty(throughputControlGroupName)) { return false; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ServerThroughputControlGroupInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ServerThroughputControlGroup.java similarity index 69% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ServerThroughputControlGroupInternal.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ServerThroughputControlGroup.java index d1337a10f6b4..f2dac5d227ad 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ServerThroughputControlGroupInternal.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ServerThroughputControlGroup.java @@ -12,55 +12,64 @@ import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; -public abstract class ServerThroughputControlGroupInternal { +public class ServerThroughputControlGroup { private final String groupName; private final boolean isDefault; private final CosmosAsyncContainer targetContainer; private final PriorityLevel priorityLevel; + private final Integer throughputBucket; - public ServerThroughputControlGroupInternal( + public ServerThroughputControlGroup( String groupName, boolean isDefault, PriorityLevel priorityLevel, + Integer throughputBucket, CosmosAsyncContainer targetContainer) { checkArgument(StringUtils.isNotEmpty(groupName), "Argument 'groupName' cannot be null or empty."); checkNotNull(targetContainer, "Argument 'targetContainer' can not be null"); + checkArgument(throughputBucket != null && throughputBucket > 0, "Target throughput should be greater than 0"); this.groupName = groupName; this.isDefault = isDefault; this.targetContainer = targetContainer; this.priorityLevel = priorityLevel; + this.throughputBucket = throughputBucket; } public String getGroupName() { - return groupName; + return this.groupName; } public boolean isDefault() { - return isDefault; + return this.isDefault; } public CosmosAsyncContainer getTargetContainer() { - return targetContainer; + return this.targetContainer; } public PriorityLevel getPriorityLevel() { - return priorityLevel; + return this.priorityLevel; + } + + public Integer getThroughputBucket() { + return this.throughputBucket; } @Override public boolean equals(Object o) { if (o == null || getClass() != o.getClass()) return false; - ServerThroughputControlGroupInternal that = (ServerThroughputControlGroupInternal) o; + ServerThroughputControlGroup that = (ServerThroughputControlGroup) o; return isDefault == that.isDefault && Objects.equals(groupName, that.groupName) - && Objects.equals(targetContainer.getId(), that.targetContainer.getId()) - && Objects.equals(priorityLevel, that.priorityLevel); + && Objects.equals(targetContainer, that.targetContainer) + && Objects.equals(priorityLevel, that.priorityLevel) + && Objects.equals(throughputBucket, that.throughputBucket); } @Override public int hashCode() { - return Objects.hash(groupName, isDefault, targetContainer, priorityLevel); + return Objects.hash(groupName, isDefault, targetContainer, priorityLevel, throughputBucket); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ThroughputBucketControlGroup.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ThroughputBucketControlGroup.java deleted file mode 100644 index ced6a97151e5..000000000000 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ThroughputBucketControlGroup.java +++ /dev/null @@ -1,45 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -package com.azure.cosmos.implementation.throughputControl.server.config; - -import com.azure.cosmos.CosmosAsyncContainer; -import com.azure.cosmos.models.PriorityLevel; - -import java.util.Objects; - -import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; - -public class ThroughputBucketControlGroup extends ServerThroughputControlGroupInternal { - private final Integer throughputBucket; - - public ThroughputBucketControlGroup( - String groupName, - boolean isDefault, - PriorityLevel priorityLevel, - Integer throughputBucket, - CosmosAsyncContainer targetContainer) { - - super(groupName, isDefault, priorityLevel, targetContainer); - - checkArgument(throughputBucket != null && throughputBucket > 0, "Target throughput should be greater than 0"); - this.throughputBucket = throughputBucket; - } - - public Integer getThroughputBucket() { - return this.throughputBucket; - } - - @Override - public boolean equals(Object o) { - if (o == null || getClass() != o.getClass()) return false; - if (!super.equals(o)) return false; - ThroughputBucketControlGroup that = (ThroughputBucketControlGroup) o; - return Objects.equals(throughputBucket, that.throughputBucket); - } - - @Override - public int hashCode() { - return Objects.hash(super.hashCode(), throughputBucket); - } -} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputContainerController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputContainerController.java index 8ce7df52fa28..8cf6d76bed5c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputContainerController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputContainerController.java @@ -8,8 +8,7 @@ import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.implementation.caches.AsyncCache; import com.azure.cosmos.implementation.throughputControl.IThroughputContainerController; -import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroupInternal; -import com.azure.cosmos.implementation.throughputControl.server.config.ThroughputBucketControlGroup; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroup; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; @@ -23,11 +22,11 @@ public class ServerThroughputContainerController implements IThroughputContainerController { private static final Logger logger = LoggerFactory.getLogger(ServerThroughputContainerController.class); - private final AsyncCache groupControllerCache; - private final Map groups; + private final AsyncCache groupControllerCache; + private final Map groups; - private ServerThroughputGroupControllerBase defaultGroupController; - public ServerThroughputContainerController(Map groups) { + private ServerThroughputGroupController defaultGroupController; + public ServerThroughputContainerController(Map groups) { checkArgument(groups != null && !groups.isEmpty(), "Throughput groups can not be null or empty"); @@ -56,14 +55,14 @@ public Mono processRequest(RxDocumentServiceRequest request, Mono orig }); } - private Mono> getOrCreateThroughputGroupController(String groupName) { + private Mono> getOrCreateThroughputGroupController(String groupName) { // If there is no control group defined, using the default group controller if (StringUtils.isEmpty(groupName)) { return Mono.just(new Utils.ValueHolder<>(this.defaultGroupController)); } - ServerThroughputControlGroupInternal group = this.groups.get(groupName); + ServerThroughputControlGroup group = this.groups.get(groupName); if (group == null) { // If the request is associated with a group not enabled, will fall back to the default one. return Mono.just(new Utils.ValueHolder<>(this.defaultGroupController)); @@ -83,18 +82,18 @@ private Mono createAndInitializeGroupContro .then(Mono.just(this)); } - private Mono resolveThroughputGroupController(ServerThroughputControlGroupInternal group) { + private Mono resolveThroughputGroupController(ServerThroughputControlGroup group) { return this.groupControllerCache.getAsync( group.getGroupName(), null, () -> this.createAndInitializeGroupController(group)); } - private Mono createAndInitializeGroupController(ServerThroughputControlGroupInternal group) { - if (group instanceof ThroughputBucketControlGroup) { + private Mono createAndInitializeGroupController(ServerThroughputControlGroup group) { + if (group instanceof ServerThroughputControlGroup) { // create throughput bucket group controller - ThroughputBucketGroupController throughputBucketGroupController = - new ThroughputBucketGroupController((ThroughputBucketControlGroup) group); + ServerThroughputGroupController throughputBucketGroupController = + new ServerThroughputGroupController(group); if (throughputBucketGroupController.isDefault()) { this.defaultGroupController = throughputBucketGroupController; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ThroughputBucketGroupController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupController.java similarity index 51% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ThroughputBucketGroupController.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupController.java index 138a5b1fe949..38c1666408fe 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ThroughputBucketGroupController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupController.java @@ -1,17 +1,18 @@ package com.azure.cosmos.implementation.throughputControl.server.controller; import com.azure.cosmos.implementation.RxDocumentServiceRequest; -import com.azure.cosmos.implementation.throughputControl.server.config.ThroughputBucketControlGroup; +import com.azure.cosmos.implementation.throughputControl.IThroughputController; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroup; import reactor.core.publisher.Mono; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; -public class ThroughputBucketGroupController extends ServerThroughputGroupControllerBase { - private final ThroughputBucketControlGroup throughputBucketControlGroup; +public class ServerThroughputGroupController implements IThroughputController { + private final ServerThroughputControlGroup serverThroughputControlGroup; - public ThroughputBucketGroupController(ThroughputBucketControlGroup throughputBucketControlGroup) { - checkNotNull(throughputBucketControlGroup, "Argument 'throughputBucketControlGroup' cannot be null."); - this.throughputBucketControlGroup = throughputBucketControlGroup; + public ServerThroughputGroupController(ServerThroughputControlGroup serverThroughputControlGroup) { + checkNotNull(serverThroughputControlGroup, "Argument 'serverThroughputControlGroup' cannot be null."); + this.serverThroughputControlGroup = serverThroughputControlGroup; } @Override @@ -27,18 +28,18 @@ public Mono init() { @Override public Mono processRequest(RxDocumentServiceRequest request, Mono originalRequestMono) { - if (this.throughputBucketControlGroup.getPriorityLevel() != null) { - request.setPriorityLevel(this.throughputBucketControlGroup.getPriorityLevel()); + if (this.serverThroughputControlGroup.getPriorityLevel() != null) { + request.setPriorityLevel(this.serverThroughputControlGroup.getPriorityLevel()); } - if (this.throughputBucketControlGroup.getThroughputBucket() != null) { - request.setThroughputBucket(this.throughputBucketControlGroup.getThroughputBucket()); + if (this.serverThroughputControlGroup.getThroughputBucket() != null) { + request.setThroughputBucket(this.serverThroughputControlGroup.getThroughputBucket()); } return originalRequestMono; } public boolean isDefault() { - return this.throughputBucketControlGroup.isDefault(); + return this.serverThroughputControlGroup.isDefault(); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupControllerBase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupControllerBase.java deleted file mode 100644 index 289d33622b9c..000000000000 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupControllerBase.java +++ /dev/null @@ -1,9 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -package com.azure.cosmos.implementation.throughputControl.server.controller; - -import com.azure.cosmos.implementation.throughputControl.IThroughputController; - -public abstract class ServerThroughputGroupControllerBase implements IThroughputController { -} From 763f754e07de584d5453661e6e59d12f6bf58c70 Mon Sep 17 00:00:00 2001 From: annie-mac Date: Fri, 25 Jul 2025 21:11:10 -0700 Subject: [PATCH 07/12] change --- .../main/java/com/azure/cosmos/CosmosAsyncContainer.java | 2 +- .../azure/cosmos/ThroughputControlGroupConfigBuilder.java | 5 +++-- .../azure/cosmos/implementation/RxDocumentClientImpl.java | 6 +++--- 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 3d40ad1b81a7..3b797fffcc7a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -2846,7 +2846,7 @@ void enableGlobalThroughputControlGroup( */ public void enableServerThroughputControlGroup(ThroughputControlGroupConfig groupConfig) { if (groupConfig.getPriorityLevel() == null && groupConfig.getThroughputBucket() == null) { - throw new IllegalArgumentException("Config 'priorityLevel' and 'throughputBucket' can not both are null."); + throw new IllegalArgumentException("Config 'priorityLevel' and 'throughputBucket' can not be null for both."); } ServerThroughputControlGroup serverThroughputControlGroup = diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfigBuilder.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfigBuilder.java index 9df41d5ce8e4..daeef3c725dc 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfigBuilder.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfigBuilder.java @@ -198,8 +198,9 @@ public ThroughputControlGroupConfig build() { if (this.targetThroughput == null && this.targetThroughputThreshold == null && this.priorityLevel == null - && this.throughputBucket == null) { // TODO: any exclusiveness between this two? - throw new IllegalArgumentException("All targetThroughput, targetThroughputThreshold, priorityLevel and throughput bucket cannot be null or empty."); + && this.throughputBucket == null) { + throw new IllegalArgumentException( + "All targetThroughput, targetThroughputThreshold, priorityLevel and throughput bucket cannot be null or empty."); } if (this.targetThroughput == null && this.targetThroughputThreshold == null) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 0aa24f5b1086..902e457cd353 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -6254,7 +6254,7 @@ public void close() { } } @Override - public void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono throughputQueryMono) { + public synchronized void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono throughputQueryMono) { checkNotNull(group, "Throughput control group can not be null"); this.enableThroughputControlStore(); @@ -6262,14 +6262,14 @@ public void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal gr } @Override - public void enableServerThroughputControlGroup(ServerThroughputControlGroup group) { + public synchronized void enableServerThroughputControlGroup(ServerThroughputControlGroup group) { checkNotNull(group, "Argument 'group' can not be null"); this.enableThroughputControlStore(); this.throughputControlStore.enableServerThroughputControlGroup(group); } - private synchronized void enableThroughputControlStore() { + private void enableThroughputControlStore() { if (this.throughputControlEnabled.compareAndSet(false, true)) { this.throughputControlStore = new ThroughputControlStore( From c713a38dc36d6c89d0a0a7485dca1bb29f472118 Mon Sep 17 00:00:00 2001 From: annie-mac Date: Mon, 28 Jul 2025 15:59:22 -0700 Subject: [PATCH 08/12] change --- ...InjectionServerErrorRuleOnDirectTests.java | 2 +- .../ThroughputControlGroupConfigTests.java | 45 +++++++++++++ .../ThroughputControlTests.java | 65 ++++++++++++++++++- .../cosmos/rx/ClientRetryPolicyE2ETests.java | 1 - .../azure/cosmos/CosmosAsyncContainer.java | 9 ++- .../com/azure/cosmos/CosmosContainer.java | 13 ++-- .../cosmos/ThroughputControlGroupConfig.java | 2 + .../ThroughputControlGroupConfigBuilder.java | 3 +- .../implementation/RxDocumentClientImpl.java | 6 +- .../ThroughputControlStore.java | 4 +- .../config/ServerThroughputControlGroup.java | 5 +- .../ServerThroughputContainerController.java | 16 ++--- .../ServerThroughputGroupController.java | 3 + .../main/java/com/azure/cosmos/util/Beta.java | 4 +- .../cosmos/ThroughputControlCodeSnippet.java | 12 ++++ 15 files changed, 159 insertions(+), 31 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupConfigTests.java diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java index 0cdae62ac2ed..079768788bfc 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java @@ -920,7 +920,7 @@ public void faultInjectionServerErrorRuleTests_LeaseNotFound(OperationType opera shouldRetryCrossRegion = true; } - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); String ruleId = "serverErrorRule-" + FaultInjectionServerErrorType.LEASE_NOT_FOUND + "-" + UUID.randomUUID(); FaultInjectionRule serverErrorRule = diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupConfigTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupConfigTests.java new file mode 100644 index 000000000000..0700f857d8dc --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupConfigTests.java @@ -0,0 +1,45 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl; + +import com.azure.cosmos.ThroughputControlGroupConfig; +import com.azure.cosmos.ThroughputControlGroupConfigBuilder; +import org.testng.annotations.Test; + +import java.util.UUID; + +import static org.assertj.core.api.AssertionsForClassTypes.assertThat; +import static org.assertj.core.api.AssertionsForClassTypes.assertThatThrownBy; + +public class ThroughputControlGroupConfigTests { + + @Test(groups = "unit") + public void throughputControlGroup_throughputBucket_priorityLevel() { + // validate throughputBucket >= 0 + assertThatThrownBy( + () -> new ThroughputControlGroupConfigBuilder() + .groupName("throughputBucket-" + UUID.randomUUID()) + .throughputBucket(-1) + .build()) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Throughput bucket should be no smaller than 0"); + + // valid config + ThroughputControlGroupConfig throughputControlGroupConfig = + new ThroughputControlGroupConfigBuilder() + .groupName("throughputBucket-" + UUID.randomUUID()) + .throughputBucket(1) + .build(); + + assertThat(throughputControlGroupConfig.getThroughputBucket()).isEqualTo(1); + + // neither priorityLevel nor throughput bucket configured + assertThatThrownBy( + () -> new ThroughputControlGroupConfigBuilder() + .groupName("throughputBucket-" + UUID.randomUUID()) + .build()) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("All targetThroughput, targetThroughputThreshold, priorityLevel and throughput bucket cannot be null or empty."); + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java index 838210d19656..cbd4d2cbdaf5 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java @@ -672,9 +672,9 @@ public void enableSameGroupMultipleTimes() { } @Test(groups = {"emulator"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT) - public void serverThroughputControl(OperationType operationType) { + public void serverThroughputControl_throughputBucket(OperationType operationType) { // TODO: currently there is no easy way to do e2e testing, so the testing here is to just verify that - // server throughput can be enabled on the container + // server throughput can be enabled on the container with throughput control this.ensureContainer(); ThroughputControlGroupConfig serverThroughputControlGroup = @@ -701,6 +701,67 @@ public void serverThroughputControl(OperationType operationType) { serverThroughputControlGroup.getGroupName()); } + @Test(groups = {"emulator"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT) + public void serverThroughputControl_priorityLevel(OperationType operationType) { + // TODO: currently there is no easy way to do e2e testing, so the testing here is to just verify that + // server throughput can be enabled on the container with priority level + this.ensureContainer(); + + ThroughputControlGroupConfig serverThroughputControlGroup = + new ThroughputControlGroupConfigBuilder() + .groupName("serverThroughputControl") + .priorityLevel(PriorityLevel.LOW) + .build(); + container.enableServerThroughputControlGroup(serverThroughputControlGroup); + + CosmosItemRequestOptions requestOptions = new CosmosItemRequestOptions(); + requestOptions.setContentResponseOnWriteEnabled(true); + requestOptions.setThroughputControlGroupName(serverThroughputControlGroup.getGroupName()); + + CosmosItemResponse createItemResponse = container.createItem(getDocumentDefinition(), requestOptions).block(); + TestItem createdItem = createItemResponse.getItem(); + this.validateRequestNotThrottled( + createItemResponse.getDiagnostics().toString(), + BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + + performDocumentOperation( + this.container, + operationType, + createdItem, + serverThroughputControlGroup.getGroupName()); + } + + @Test(groups = {"emulator"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT) + public void serverThroughputControl_priorityLevel_throughputBucket(OperationType operationType) { + // TODO: currently there is no easy way to do e2e testing, so the testing here is to just verify that + // server throughput can be enabled on the container with priority level and throughput bucket + this.ensureContainer(); + + ThroughputControlGroupConfig serverThroughputControlGroup = + new ThroughputControlGroupConfigBuilder() + .groupName("serverThroughputControl") + .priorityLevel(PriorityLevel.LOW) + .throughputBucket(3) + .build(); + container.enableServerThroughputControlGroup(serverThroughputControlGroup); + + CosmosItemRequestOptions requestOptions = new CosmosItemRequestOptions(); + requestOptions.setContentResponseOnWriteEnabled(true); + requestOptions.setThroughputControlGroupName(serverThroughputControlGroup.getGroupName()); + + CosmosItemResponse createItemResponse = container.createItem(getDocumentDefinition(), requestOptions).block(); + TestItem createdItem = createItemResponse.getItem(); + this.validateRequestNotThrottled( + createItemResponse.getDiagnostics().toString(), + BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + + performDocumentOperation( + this.container, + operationType, + createdItem, + serverThroughputControlGroup.getGroupName()); + } + @Test(groups = {"emulator"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT) public void throughputControl_LocalAndServer_requestOptions(OperationType operationType) { this.ensureContainer(); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ClientRetryPolicyE2ETests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ClientRetryPolicyE2ETests.java index 6b3f8c17ba46..0db09f46aed8 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ClientRetryPolicyE2ETests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ClientRetryPolicyE2ETests.java @@ -687,7 +687,6 @@ private Mono performDocumentOperation( Function extractPartitionKeyFunc, boolean isReadMany) { - try { if (operationType == OperationType.Query && isReadMany) { CosmosQueryRequestOptions queryRequestOptions = new CosmosQueryRequestOptions(); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 682cc5d051ea..b3fb59f595ee 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -71,6 +71,7 @@ import com.azure.cosmos.models.SqlQuerySpec; import com.azure.cosmos.models.ThroughputProperties; import com.azure.cosmos.models.ThroughputResponse; +import com.azure.cosmos.util.Beta; import com.azure.cosmos.util.CosmosPagedFlux; import com.azure.cosmos.util.UtilBridgeInternal; import org.slf4j.Logger; @@ -2825,12 +2826,14 @@ void enableGlobalThroughputControlGroup( /*** * Enable the server throughput bucket control group. - *

- * For more information about throughput bucket please visit - * Throughput buckets in Azure Cosmos DB + *
+ * + * + * * * @param groupConfig the throughput control group config, see {@link ThroughputControlGroupConfig}. */ + @Beta(value = Beta.SinceVersion.V4_74_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) public void enableServerThroughputControlGroup(ThroughputControlGroupConfig groupConfig) { if (groupConfig.getPriorityLevel() == null && groupConfig.getThroughputBucket() == null) { throw new IllegalArgumentException("Config 'priorityLevel' and 'throughputBucket' can not be null for both."); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java index 35aeac425b62..cf22b8455f3d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java @@ -28,6 +28,7 @@ import com.azure.cosmos.models.SqlQuerySpec; import com.azure.cosmos.models.ThroughputProperties; import com.azure.cosmos.models.ThroughputResponse; +import com.azure.cosmos.util.Beta; import com.azure.cosmos.util.CosmosPagedFlux; import com.azure.cosmos.util.CosmosPagedIterable; import org.slf4j.Logger; @@ -948,7 +949,7 @@ public CosmosScripts getScripts() { } // TODO: should make partitionkey public in CosmosAsyncItem and fix the below call - + private CosmosPagedIterable getCosmosPagedIterable(CosmosPagedFlux cosmosPagedFlux) { return new CosmosPagedIterable<>(cosmosPagedFlux); } @@ -1031,13 +1032,15 @@ public void enableGlobalThroughputControlGroup(ThroughputControlGroupConfig grou } /*** - * Enable the server throughput bucket control group. - *

- * For more information about throughput bucket please visit - * Throughput buckets in Azure Cosmos DB + * Enable the server throughput control group. + *
+ * + * + * * * @param groupConfig the throughput control group config, see {@link ThroughputControlGroupConfig}. */ + @Beta(value = Beta.SinceVersion.V4_74_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) public void enableServerThroughputControlGroup(ThroughputControlGroupConfig groupConfig) { this.asyncContainer.enableServerThroughputControlGroup(groupConfig); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfig.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfig.java index de626b24e708..b5f3cedac8fe 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfig.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfig.java @@ -4,6 +4,7 @@ package com.azure.cosmos; import com.azure.cosmos.models.PriorityLevel; +import com.azure.cosmos.util.Beta; /** * Throughput control group configuration. @@ -88,6 +89,7 @@ public Double getTargetThroughputThreshold() { * * @return the throughput bucket of the throughput control group. */ + @Beta(value = Beta.SinceVersion.V4_74_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) public Integer getThroughputBucket() { return this.throughputBucket; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfigBuilder.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfigBuilder.java index daeef3c725dc..49f833626428 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfigBuilder.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfigBuilder.java @@ -167,8 +167,9 @@ public ThroughputControlGroupConfigBuilder defaultControlGroup(boolean aDefault) * @param throughputBucket the throughput bucket id. * @return The {@link ThroughputControlGroupConfigBuilder}. */ + @Beta(value = Beta.SinceVersion.V4_74_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) public ThroughputControlGroupConfigBuilder throughputBucket(int throughputBucket) { - checkArgument(throughputBucket > 0, "Throughput Bucket should be greater than 0"); + checkArgument(throughputBucket >= 0, "Throughput bucket should be no smaller than 0"); this.throughputBucket = throughputBucket; return this; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 6742e8ce7e32..f7727b95f9d0 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -6363,7 +6363,7 @@ public void close() { } } @Override - public synchronized void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono throughputQueryMono) { + public void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono throughputQueryMono) { checkNotNull(group, "Throughput control group can not be null"); this.enableThroughputControlStore(); @@ -6371,14 +6371,14 @@ public synchronized void enableSDKThroughputControlGroup(SDKThroughputControlGro } @Override - public synchronized void enableServerThroughputControlGroup(ServerThroughputControlGroup group) { + public void enableServerThroughputControlGroup(ServerThroughputControlGroup group) { checkNotNull(group, "Argument 'group' can not be null"); this.enableThroughputControlStore(); this.throughputControlStore.enableServerThroughputControlGroup(group); } - private void enableThroughputControlStore() { + private synchronized void enableThroughputControlStore() { if (this.throughputControlEnabled.compareAndSet(false, true)) { this.throughputControlStore = new ThroughputControlStore( diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlStore.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlStore.java index 250dbd466ae7..feb207b9abb5 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlStore.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlStore.java @@ -38,7 +38,7 @@ public ThroughputControlStore( this.serverThroughputControlStore = new ServerThroughputControlStore(); } - public void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono throughputQueryMono) { + public synchronized void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono throughputQueryMono) { checkNotNull(group, "Throughput control group cannot be null"); if (group.isDefault()) { @@ -52,7 +52,7 @@ public void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal gr this.sdkThroughputControlStore.enableThroughputControlGroup(group, throughputQueryMono); } - public void enableServerThroughputControlGroup(ServerThroughputControlGroup group) { + public synchronized void enableServerThroughputControlGroup(ServerThroughputControlGroup group) { checkNotNull(group, "Throughput control group cannot be null"); if (group.isDefault()) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ServerThroughputControlGroup.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ServerThroughputControlGroup.java index f2dac5d227ad..697529a93cfd 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ServerThroughputControlGroup.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ServerThroughputControlGroup.java @@ -28,7 +28,10 @@ public ServerThroughputControlGroup( checkArgument(StringUtils.isNotEmpty(groupName), "Argument 'groupName' cannot be null or empty."); checkNotNull(targetContainer, "Argument 'targetContainer' can not be null"); - checkArgument(throughputBucket != null && throughputBucket > 0, "Target throughput should be greater than 0"); + checkArgument(throughputBucket != null && throughputBucket >= 0, "Target throughput should be no less than 0"); + checkArgument( + priorityLevel != null || throughputBucket != null, + "At least one of 'priorityLevel' or 'throughputBucket' must be provided."); this.groupName = groupName; this.isDefault = isDefault; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputContainerController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputContainerController.java index 8cf6d76bed5c..4aa3e51a97d9 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputContainerController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputContainerController.java @@ -90,17 +90,11 @@ private Mono resolveThroughputGroupController(S } private Mono createAndInitializeGroupController(ServerThroughputControlGroup group) { - if (group instanceof ServerThroughputControlGroup) { - // create throughput bucket group controller - ServerThroughputGroupController throughputBucketGroupController = - new ServerThroughputGroupController(group); - if (throughputBucketGroupController.isDefault()) { - this.defaultGroupController = throughputBucketGroupController; - } - return Mono.just(throughputBucketGroupController); + // create throughput bucket group controller + ServerThroughputGroupController throughputGroupController = new ServerThroughputGroupController(group); + if (throughputGroupController.isDefault()) { + this.defaultGroupController = throughputGroupController; } - - return Mono.error( - new IllegalArgumentException("Server throughput control group type is not supported: " + group.getClass().getSimpleName())); + return Mono.just(throughputGroupController); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupController.java index 38c1666408fe..e2ad4ad34a58 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupController.java @@ -1,3 +1,6 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + package com.azure.cosmos.implementation.throughputControl.server.controller; import com.azure.cosmos.implementation.RxDocumentServiceRequest; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java index 857cb893ed4b..5ccd99db2fd0 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java @@ -105,6 +105,8 @@ public enum SinceVersion { /** v4.69.0 */ V4_69_0, /** v4.71.0 */ - V4_71_0 + V4_71_0, + /** v4.74.0 */ + V4_74_0 } } diff --git a/sdk/cosmos/azure-cosmos/src/samples/java/com/azure/cosmos/ThroughputControlCodeSnippet.java b/sdk/cosmos/azure-cosmos/src/samples/java/com/azure/cosmos/ThroughputControlCodeSnippet.java index 51d482c8efe9..22c173f928dd 100644 --- a/sdk/cosmos/azure-cosmos/src/samples/java/com/azure/cosmos/ThroughputControlCodeSnippet.java +++ b/sdk/cosmos/azure-cosmos/src/samples/java/com/azure/cosmos/ThroughputControlCodeSnippet.java @@ -53,4 +53,16 @@ public void codeSnippetForEnableGlobalThroughputControl() { container.enableGlobalThroughputControlGroup(groupConfig, globalControlConfig); // END: com.azure.cosmos.throughputControl.globalControl } + + public void codeSnippetForEnableServerThroughputControl() { + // BEGIN: com.azure.cosmos.throughputControl.serverControl + ThroughputControlGroupConfig groupConfig = + new ThroughputControlGroupConfigBuilder() + .groupName("localControlGroup") + .throughputBucket(2) + .build(); + + container.enableServerThroughputControlGroup(groupConfig); + // END: com.azure.cosmos.throughputControl.serverControl + } } From c9658159cdbb59afb4e249ff6b9e40c31ee60043 Mon Sep 17 00:00:00 2001 From: annie-mac Date: Mon, 28 Jul 2025 22:40:17 -0700 Subject: [PATCH 09/12] add throughput control group config in diagnostics --- .../ThroughputControlTests.java | 344 +++++++++++++----- .../sdk/ThroughputRequestThrottlerTests.java | 10 +- .../com/azure/cosmos/CosmosContainer.java | 1 - .../ClientSideRequestStatistics.java | 14 + .../DocumentServiceRequestContext.java | 9 +- .../StoreResponseDiagnostics.java | 16 + .../StoreResultDiagnostics.java | 8 + .../IThroughputControlGroup.java | 8 + .../ThroughputControlRequestContext.java | 33 ++ .../sdk/ThroughputRequestThrottler.java | 6 +- .../SDKThroughputControlGroupInternal.java | 23 +- .../SDKThroughputGroupControllerBase.java | 6 + .../config/ServerThroughputControlGroup.java | 28 +- .../ServerThroughputGroupController.java | 7 + 14 files changed, 410 insertions(+), 103 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/IThroughputControlGroup.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlRequestContext.java diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java index cbd4d2cbdaf5..ca97bcab6645 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java @@ -107,15 +107,24 @@ public void throughputLocalControl_requestOptions(OperationType operationType) { CosmosItemResponse createItemResponse = container.createItem(getDocumentDefinition(), requestOptions).block(); TestItem createdItem = createItemResponse.getItem(); - this.validateRequestNotThrottled( - createItemResponse.getDiagnostics().toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, 6) + ); // second request to group-1. which will get throttled - CosmosDiagnostics cosmosDiagnostics = performDocumentOperation(this.container, operationType, createdItem, groupConfig.getGroupName()); - this.validateRequestThrottled( - cosmosDiagnostics.toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + cosmosDiagnosticsString = performDocumentOperation(this.container, operationType, createdItem, groupConfig.getGroupName()).toString(); + this.validateRequestThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, 6) + ); } @Test(groups = {"emulator"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT) @@ -142,15 +151,24 @@ public void throughputLocalControl_default(OperationType operationType) { CosmosItemResponse createItemResponse = testContainer.createItem(getDocumentDefinition()).block(); TestItem createdItem = createItemResponse.getItem(); - this.validateRequestNotThrottled( - createItemResponse.getDiagnostics().toString(), - BridgeInternal.getContextClient(cosmosAsyncClient).getConnectionPolicy().getConnectionMode()); + + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), true, 6) + ); // second request to group-1. which will get throttled - CosmosDiagnostics cosmosDiagnostics = performDocumentOperation(testContainer, operationType, createdItem, groupConfig.getGroupName()); - this.validateRequestThrottled( - cosmosDiagnostics.toString(), - BridgeInternal.getContextClient(cosmosAsyncClient).getConnectionPolicy().getConnectionMode()); + cosmosDiagnosticsString = + performDocumentOperation(testContainer, operationType, createdItem, groupConfig.getGroupName()).toString(); + this.validateRequestThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), true, 6) + ); } finally { safeClose(cosmosAsyncClient); } @@ -190,15 +208,23 @@ public void throughputLocalControlWithThroughputQuery() { CosmosItemResponse createItemResponse = testContainer.createItem(getDocumentDefinition(), requestOptions).block(); TestItem createdItem = createItemResponse.getItem(); - this.validateRequestNotThrottled( - createItemResponse.getDiagnostics().toString(), - BridgeInternal.getContextClient(cosmosAsyncClient).getConnectionPolicy().getConnectionMode()); + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRUThreshold=%s)", groupConfig.getGroupName(), false, 0.9) + ); // second request to group-1. which will get throttled - CosmosDiagnostics cosmosDiagnostics = performDocumentOperation(testContainer, OperationType.Read, createdItem, groupConfig.getGroupName()); - this.validateRequestThrottled( - cosmosDiagnostics.toString(), - BridgeInternal.getContextClient(cosmosAsyncClient).getConnectionPolicy().getConnectionMode()); + cosmosDiagnosticsString = + performDocumentOperation(testContainer, OperationType.Read, createdItem, groupConfig.getGroupName()).toString(); + this.validateRequestThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRUThreshold=%s)", groupConfig.getGroupName(), false, 0.9) + ); assertThat(throughputQueryMonoCalledCount.get()).isGreaterThanOrEqualTo(1); } finally { @@ -222,9 +248,19 @@ public void throughputLocalControlPriorityLevel(OperationType operationType) { CosmosItemResponse createItemResponse = container.createItem(getDocumentDefinition(), requestOptions).block(); TestItem createdItem = createItemResponse.getItem(); + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, priorityLevel=%s, targetRU=%s)", groupConfig.getGroupName(), false, PriorityLevel.LOW, Integer.MAX_VALUE) + ); - performDocumentOperation(container, operationType, createdItem, groupConfig.getGroupName()); - + cosmosDiagnosticsString = performDocumentOperation(container, operationType, createdItem, groupConfig.getGroupName()).toString(); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, priorityLevel=%s, targetRU=%s)", groupConfig.getGroupName(), false, PriorityLevel.LOW, Integer.MAX_VALUE) + ); assertThat(groupConfig.getTargetThroughput()).isEqualTo(Integer.MAX_VALUE); assertThat(createItemResponse.getStatusCode()).isEqualTo(201); } @@ -257,15 +293,24 @@ public void throughputGlobalControl(OperationType operationType) { CosmosItemResponse createItemResponse = container.createItem(getDocumentDefinition(), requestOptions).block(); TestItem createdItem = createItemResponse.getItem(); - this.validateRequestNotThrottled( - createItemResponse.getDiagnostics().toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, 6) + ); // second request to same group. which will get throttled - CosmosDiagnostics cosmosDiagnostics = performDocumentOperation(this.container, operationType, createdItem, groupConfig.getGroupName()); - this.validateRequestThrottled( - cosmosDiagnostics.toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + cosmosDiagnosticsString = + performDocumentOperation(this.container, operationType, createdItem, groupConfig.getGroupName()).toString(); + this.validateRequestThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, 6) + ); } finally { controlContainer .delete() @@ -314,15 +359,23 @@ public void throughputGlobalControlWithThroughputQuery() { CosmosItemResponse createItemResponse = testContainer.createItem(getDocumentDefinition(), requestOptions).block(); TestItem createdItem = createItemResponse.getItem(); - this.validateRequestNotThrottled( - createItemResponse.getDiagnostics().toString(), - BridgeInternal.getContextClient(cosmosAsyncClient).getConnectionPolicy().getConnectionMode()); + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRUThreshold=%s)", groupConfig.getGroupName(), false, 0.9) + ); // second request to same group. which will get throttled - CosmosDiagnostics cosmosDiagnostics = performDocumentOperation(testContainer, OperationType.Create, createdItem, groupConfig.getGroupName()); - this.validateRequestThrottled( - cosmosDiagnostics.toString(), - BridgeInternal.getContextClient(cosmosAsyncClient).getConnectionPolicy().getConnectionMode()); + cosmosDiagnosticsString = + performDocumentOperation(testContainer, OperationType.Create, createdItem, groupConfig.getGroupName()).toString(); + this.validateRequestThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRUThreshold=%s)", groupConfig.getGroupName(), false, 0.9) + ); assertThat(throughputQueryMonoCalledCount.get()).isGreaterThanOrEqualTo(1); } finally { @@ -378,21 +431,29 @@ public void throughputGlobalControlCanUpdateConfig(OperationType operationType) .createItem(getDocumentDefinition(), requestOptions) .block(); TestItem createdItem = createItemResponse.getItem(); - this.validateRequestNotThrottled( - createItemResponse.getDiagnostics().toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); - CosmosDiagnostics cosmosDiagnostics = performDocumentOperation(this.container, operationType, createdItem, groupConfig.getGroupName()); + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, targetThroughput) + ); + + cosmosDiagnosticsString = + performDocumentOperation(this.container, operationType, createdItem, groupConfig.getGroupName()).toString(); if (shouldThrottleSecondRequest) { - this.validateRequestThrottled( - cosmosDiagnostics.toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + this.validateRequestThrottled(cosmosDiagnosticsString); } else { - this.validateRequestNotThrottled( - cosmosDiagnostics.toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + this.validateRequestNotThrottled(cosmosDiagnosticsString); } + + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, targetThroughput) + ); } } finally { safeDeleteCollection(database.getContainer(controlContainerId)); @@ -436,9 +497,14 @@ public void throughputLocalControlForContainerCreateDeleteWithSameName(Operation // Step2: first request to group-1, which will not get throttled, but will consume all the rus of the throughput control group. CosmosItemResponse createItemResponse = createdContainer.createItem(getDocumentDefinition(), requestOptions).block(); TestItem createdItem = createItemResponse.getItem(); - this.validateRequestNotThrottled( - createItemResponse.getDiagnostics().toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, 1) + ); // Step 3: delete the container safeDeleteCollection(createdContainer); @@ -452,16 +518,24 @@ public void throughputLocalControlForContainerCreateDeleteWithSameName(Operation createdItem = createdContainer.createItem(getDocumentDefinition()).block().getItem(); // Step 6: second request to group-1. which will not get throttled because new container controller will be built. - CosmosDiagnostics cosmosDiagnostics = performDocumentOperation(createdContainer, operationType, createdItem, groupConfig.getGroupName()); - this.validateRequestNotThrottled( - cosmosDiagnostics.toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + cosmosDiagnosticsString = + performDocumentOperation(createdContainer, operationType, createdItem, groupConfig.getGroupName()).toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, 1) + ); // Step 7: third request to group-1, which will get throttled - cosmosDiagnostics = performDocumentOperation(createdContainer, operationType, createdItem, groupConfig.getGroupName()); - this.validateRequestThrottled( - cosmosDiagnostics.toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + cosmosDiagnosticsString = + performDocumentOperation(createdContainer, operationType, createdItem, groupConfig.getGroupName()).toString(); + this.validateRequestThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, 1) + ); } finally { createdContainer.delete().block(); } @@ -484,9 +558,14 @@ public void throughputLocalControl_createItem() throws InterruptedException { requestOptions.setThroughputControlGroupName(groupConfig.getGroupName()); CosmosItemResponse createItemResponse = container.createItem(getDocumentDefinition(), requestOptions).block(); - this.validateRequestNotThrottled( - createItemResponse.getDiagnostics().toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, 6) + ); // second request to same group. which will get throttled TestItem itemGetThrottled = getDocumentDefinition(createItemResponse.getItem().getMypk()); @@ -544,7 +623,7 @@ public void throughputControlContinueOnInitError(boolean continueOnInitError) { } @Test(groups = {"emulator"}, timeOut = TIMEOUT * 4) - public void throughputGlobalControlMultipleClients() throws InterruptedException { + public void throughputGlobalControlMultipleClients() { this.ensureContainer(); List cosmosAsyncClients = new ArrayList<>(); // and do not enable ttl on the container so to test how many items are created. @@ -674,7 +753,7 @@ public void enableSameGroupMultipleTimes() { @Test(groups = {"emulator"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT) public void serverThroughputControl_throughputBucket(OperationType operationType) { // TODO: currently there is no easy way to do e2e testing, so the testing here is to just verify that - // server throughput can be enabled on the container with throughput control + // server throughput can be enabled on the container with throughput bucket this.ensureContainer(); ThroughputControlGroupConfig serverThroughputControlGroup = @@ -690,15 +769,25 @@ public void serverThroughputControl_throughputBucket(OperationType operationType CosmosItemResponse createItemResponse = container.createItem(getDocumentDefinition(), requestOptions).block(); TestItem createdItem = createItemResponse.getItem(); - this.validateRequestNotThrottled( - createItemResponse.getDiagnostics().toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); - performDocumentOperation( + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + serverThroughputControlGroup.getGroupName(), + String.format("(name=%s, default=%s, throughputBucket=%s)", serverThroughputControlGroup.getGroupName(), false, 2) + ); + + cosmosDiagnosticsString = performDocumentOperation( this.container, operationType, createdItem, - serverThroughputControlGroup.getGroupName()); + serverThroughputControlGroup.getGroupName()).toString(); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + serverThroughputControlGroup.getGroupName(), + String.format("(name=%s, default=%s, throughputBucket=%s)", serverThroughputControlGroup.getGroupName(), false, 2) + ); } @Test(groups = {"emulator"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT) @@ -720,15 +809,25 @@ public void serverThroughputControl_priorityLevel(OperationType operationType) { CosmosItemResponse createItemResponse = container.createItem(getDocumentDefinition(), requestOptions).block(); TestItem createdItem = createItemResponse.getItem(); - this.validateRequestNotThrottled( - createItemResponse.getDiagnostics().toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); - performDocumentOperation( + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + serverThroughputControlGroup.getGroupName(), + String.format("(name=%s, default=%s, priorityLevel=%s)", serverThroughputControlGroup.getGroupName(), false, PriorityLevel.LOW) + ); + + cosmosDiagnosticsString = performDocumentOperation( this.container, operationType, createdItem, - serverThroughputControlGroup.getGroupName()); + serverThroughputControlGroup.getGroupName()).toString(); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + serverThroughputControlGroup.getGroupName(), + String.format("(name=%s, default=%s, priorityLevel=%s)", serverThroughputControlGroup.getGroupName(), false, PriorityLevel.LOW) + ); } @Test(groups = {"emulator"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT) @@ -751,15 +850,35 @@ public void serverThroughputControl_priorityLevel_throughputBucket(OperationType CosmosItemResponse createItemResponse = container.createItem(getDocumentDefinition(), requestOptions).block(); TestItem createdItem = createItemResponse.getItem(); - this.validateRequestNotThrottled( - createItemResponse.getDiagnostics().toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); - performDocumentOperation( + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + serverThroughputControlGroup.getGroupName(), + String.format( + "(name=%s, default=%s, priorityLevel=%s, throughputBucket=%s)", + serverThroughputControlGroup.getGroupName(), + false, + PriorityLevel.LOW, + 3) + ); + + cosmosDiagnosticsString = performDocumentOperation( this.container, operationType, createdItem, - serverThroughputControlGroup.getGroupName()); + serverThroughputControlGroup.getGroupName()).toString(); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + serverThroughputControlGroup.getGroupName(), + String.format( + "(name=%s, default=%s, priorityLevel=%s, throughputBucket=%s)", + serverThroughputControlGroup.getGroupName(), + false, + PriorityLevel.LOW, + 3) + ); } @Test(groups = {"emulator"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT) @@ -789,19 +908,34 @@ public void throughputControl_LocalAndServer_requestOptions(OperationType operat CosmosItemResponse createItemResponse = container.createItem(getDocumentDefinition(), requestOptions).block(); TestItem createdItem = createItemResponse.getItem(); - this.validateRequestNotThrottled( - createItemResponse.getDiagnostics().toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, 6) + ); // second request to group-1. which will get throttled - CosmosDiagnostics cosmosDiagnostics = performDocumentOperation(this.container, operationType, createdItem, groupConfig.getGroupName()); - this.validateRequestThrottled( - cosmosDiagnostics.toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + cosmosDiagnosticsString = + performDocumentOperation(this.container, operationType, createdItem, groupConfig.getGroupName()).toString(); + this.validateRequestThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, 6) + ); // third request to server group, which will not get throttled - requestOptions.setThroughputControlGroupName(serverGroupConfig.getGroupName()); - performDocumentOperation(this.container, operationType, createdItem, groupConfig.getGroupName()); + cosmosDiagnosticsString = + performDocumentOperation(this.container, operationType, createdItem, serverGroupConfig.getGroupName()).toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + serverGroupConfig.getGroupName(), + String.format("(name=%s, default=%s, throughputBucket=%s)", serverGroupConfig.getGroupName(), false, 3) + ); } @Test(groups = {"emulator"}, timeOut = TIMEOUT) @@ -830,6 +964,26 @@ public void throughputControlDefaultGroup_LocalAndServer_requestOptions() { .hasMessage("A default group already exists"); } + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void throughputControl_noThroughputControlGroupEnabled() { + this.ensureContainer(); + + // This test is verify that request will succeed if no throughput control group defined + String throughputControlGroupName = "throughputControlGroup-" + UUID.randomUUID(); + CosmosItemRequestOptions itemRequestOptions = new CosmosItemRequestOptions(); + itemRequestOptions.setThroughputControlGroupName(throughputControlGroupName); + String cosmosDiagnosticsString = + container + .createItem(getDocumentDefinition(), itemRequestOptions) + .block() + .getDiagnostics() + .toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + assertThat(cosmosDiagnosticsString).contains("requestTCG"); + assertThat(cosmosDiagnosticsString).contains(throughputControlGroupName); + assertThat(cosmosDiagnosticsString).doesNotContain("requestTCGConfig"); + } + @BeforeClass(groups = { "emulator" }, timeOut = 4 * SETUP_TIMEOUT) public void before_ThroughputControllerTest() { this.ensureContainer(); @@ -874,18 +1028,30 @@ private static TestItem getDocumentDefinition(String partitionKey) { ); } - private void validateRequestThrottled(String cosmosDiagnostics, ConnectionMode connectionMode) { + private void validateRequestThrottled(String cosmosDiagnostics) { assertThat(cosmosDiagnostics).isNotEmpty(); assertThat(cosmosDiagnostics).contains("\"statusCode\":429"); assertThat(cosmosDiagnostics).contains("\"subStatusCode\":10003"); } - private void validateRequestNotThrottled(String cosmosDiagnostics, ConnectionMode connectionMode) { + private void validateRequestNotThrottled(String cosmosDiagnostics) { assertThat(cosmosDiagnostics).isNotEmpty(); assertThat(cosmosDiagnostics).doesNotContain("\"statusCode\":429"); assertThat(cosmosDiagnostics).doesNotContain("\"subStatusCode\":10003"); } + private void validateThroughputControlDiagnostics( + String cosmosDiagnostics, + String groupName, + String groupConfig) { + + assertThat(cosmosDiagnostics).isNotEmpty(); + assertThat(cosmosDiagnostics).contains("requestTCG"); + assertThat(cosmosDiagnostics).contains(groupName); + assertThat(cosmosDiagnostics).contains("requestTCGConfig"); + assertThat(cosmosDiagnostics).contains(groupConfig); + } + private CosmosDiagnostics performDocumentOperation( CosmosAsyncContainer cosmosAsyncContainer, OperationType operationType, diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputRequestThrottlerTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputRequestThrottlerTests.java index 21e3a0f06f34..6b93d901d6b8 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputRequestThrottlerTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputRequestThrottlerTests.java @@ -53,7 +53,7 @@ public void processRequest() { this.assertRequestThrottlerState(requestThrottler, availableThroughput, scheduledThroughput); // Request2: will get throttled since there is no available throughput - requestMock.requestContext.throughputControlCycleId = StringUtils.EMPTY; + requestMock.requestContext.throughputControlRequestContext.setThroughputControlCycleId(StringUtils.EMPTY); TestPublisher requestPublisher2 = TestPublisher.create(); StepVerifier.create(requestThrottler.processRequest(requestMock, requestPublisher2.mono())) .verifyError(RequestRateTooLargeException.class); @@ -66,7 +66,7 @@ public void processRequest() { assertThat(requestThrottler.getAvailableThroughput()).isEqualTo(availableThroughput); // Request 3: will get throttled since there is no available throughput - requestMock.requestContext.throughputControlCycleId = StringUtils.EMPTY; + requestMock.requestContext.throughputControlRequestContext.setThroughputControlCycleId(StringUtils.EMPTY); TestPublisher requestPublisher3 = TestPublisher.create(); StepVerifier.create(requestThrottler.processRequest(requestMock, requestPublisher3.mono())) .verifyErrorSatisfies((t) -> { @@ -87,7 +87,7 @@ public void processRequest() { Mockito.doReturn(mockHeaders).when(bulkRequestMock).getHeaders(); bulkRequestMock.requestContext = new DocumentServiceRequestContext(); - bulkRequestMock.requestContext.throughputControlCycleId = StringUtils.EMPTY; + bulkRequestMock.requestContext.throughputControlRequestContext.setThroughputControlCycleId(StringUtils.EMPTY); TestPublisher requestPublisher4 = TestPublisher.create(); StepVerifier.create(requestThrottler.processRequest(bulkRequestMock, requestPublisher4.mono())) .verifyErrorSatisfies((t) -> { @@ -105,7 +105,7 @@ public void processRequest() { assertThat(requestThrottler.getAvailableThroughput()).isEqualTo(availableThroughput); // Request 5: will pass the request, and record the charge from exception - requestMock.requestContext.throughputControlCycleId = StringUtils.EMPTY; + requestMock.requestContext.throughputControlRequestContext.setThroughputControlCycleId(StringUtils.EMPTY); NotFoundException notFoundException = Mockito.mock(NotFoundException.class); Mockito.doReturn(requestChargePerRequest).when(notFoundException).getRequestCharge(); TestPublisher requestPublisher5 = TestPublisher.create(); @@ -137,7 +137,7 @@ public void responseOutOfCycle() { TestPublisher requestPublisher1 = TestPublisher.create(); StepVerifier.create(requestThrottler.processRequest(requestMock, requestPublisher1.mono())) .then(() -> { - requestMock.requestContext.throughputControlCycleId = UUID.randomUUID().toString(); + requestMock.requestContext.throughputControlRequestContext.setThroughputControlCycleId(UUID.randomUUID().toString()); requestPublisher1.emit(responseMock); }) .expectNext(responseMock) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java index cf22b8455f3d..a03f8259b734 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java @@ -949,7 +949,6 @@ public CosmosScripts getScripts() { } // TODO: should make partitionkey public in CosmosAsyncItem and fix the below call - private CosmosPagedIterable getCosmosPagedIterable(CosmosPagedFlux cosmosPagedFlux) { return new CosmosPagedIterable<>(cosmosPagedFlux); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientSideRequestStatistics.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientSideRequestStatistics.java index 3503a46e11c1..fbfaf776edc8 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientSideRequestStatistics.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientSideRequestStatistics.java @@ -270,6 +270,8 @@ public void recordGatewayResponse( gatewayStatistics.faultInjectionRuleId = storeResponseDiagnostics.getFaultInjectionRuleId(); gatewayStatistics.faultInjectionEvaluationResults = storeResponseDiagnostics.getFaultInjectionEvaluationResults(); gatewayStatistics.endpoint = storeResponseDiagnostics.getEndpoint(); + gatewayStatistics.requestThroughputControlGroupName = storeResponseDiagnostics.getRequestThroughputControlGroupName(); + gatewayStatistics.requestThroughputControlGroupConfig = storeResponseDiagnostics.getRequestThroughputControlGroupConfig(); this.activityId = storeResponseDiagnostics.getActivityId() != null ? storeResponseDiagnostics.getActivityId() : rxDocumentServiceRequest.getActivityId().toString(); @@ -910,6 +912,8 @@ public static class GatewayStatistics { private PerPartitionCircuitBreakerInfoHolder perPartitionCircuitBreakerInfoHolder; private PerPartitionFailoverInfoHolder perPartitionFailoverInfoHolder; private String endpoint; + private String requestThroughputControlGroupName; + private String requestThroughputControlGroupConfig; public String getSessionToken() { return sessionToken; @@ -979,6 +983,14 @@ public String getEndpoint() { return this.endpoint; } + public String getRequestThroughputControlGroupName() { + return this.requestThroughputControlGroupName; + } + + public String getRequestThroughputControlGroupConfig() { + return this.requestThroughputControlGroupConfig; + } + public static class GatewayStatisticsSerializer extends StdSerializer { private static final long serialVersionUID = 1L; @@ -1016,6 +1028,8 @@ public void serialize(GatewayStatistics gatewayStatistics, this.writeNonNullObjectField(jsonGenerator, "perPartitionCircuitBreakerInfoHolder", gatewayStatistics.getPerPartitionCircuitBreakerInfoHolder()); this.writeNonNullObjectField(jsonGenerator, "perPartitionFailoverInfoHolder", gatewayStatistics.getPerPartitionFailoverInfoHolder()); + this.writeNonNullStringField(jsonGenerator, "requestTCG", gatewayStatistics.getRequestThroughputControlGroupName()); + this.writeNonNullStringField(jsonGenerator, "requestTCGConfig", gatewayStatistics.getRequestThroughputControlGroupConfig()); jsonGenerator.writeEndObject(); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DocumentServiceRequestContext.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DocumentServiceRequestContext.java index 426c7493091a..18da5250c458 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DocumentServiceRequestContext.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DocumentServiceRequestContext.java @@ -18,6 +18,7 @@ import com.azure.cosmos.implementation.directconnectivity.Uri; import com.azure.cosmos.implementation.routing.PartitionKeyInternal; import com.azure.cosmos.implementation.routing.RegionalRoutingContext; +import com.azure.cosmos.implementation.throughputControl.ThroughputControlRequestContext; import java.util.ArrayList; import java.util.List; @@ -54,7 +55,7 @@ public class DocumentServiceRequestContext implements Cloneable { public volatile PartitionKeyInternal effectivePartitionKey; public volatile CosmosDiagnostics cosmosDiagnostics; public volatile String resourcePhysicalAddress; - public volatile String throughputControlCycleId; + public ThroughputControlRequestContext throughputControlRequestContext; public volatile boolean replicaAddressValidationEnabled = Configs.isReplicaAddressValidationEnabled(); private final Set failedEndpoints = ConcurrentHashMap.newKeySet(); private CosmosEndToEndOperationLatencyPolicyConfig endToEndOperationLatencyPolicyConfig; @@ -132,6 +133,10 @@ public void addToFailedEndpoints(Exception exception, Uri address) { } } + public void setThroughputControlRequestContext(ThroughputControlRequestContext throughputControlRequestContext) { + this.throughputControlRequestContext = throughputControlRequestContext; + } + @Override public DocumentServiceRequestContext clone() { DocumentServiceRequestContext context = new DocumentServiceRequestContext(); @@ -158,7 +163,7 @@ public DocumentServiceRequestContext clone() { context.performedBackgroundAddressRefresh = this.performedBackgroundAddressRefresh; context.cosmosDiagnostics = this.cosmosDiagnostics; context.resourcePhysicalAddress = this.resourcePhysicalAddress; - context.throughputControlCycleId = this.throughputControlCycleId; + context.throughputControlRequestContext = this.throughputControlRequestContext; context.replicaAddressValidationEnabled = this.replicaAddressValidationEnabled; context.endToEndOperationLatencyPolicyConfig = this.endToEndOperationLatencyPolicyConfig; context.unavailableRegionsForPartition = this.unavailableRegionsForPartition; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResponseDiagnostics.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResponseDiagnostics.java index eebcda67db69..936913f98cb9 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResponseDiagnostics.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResponseDiagnostics.java @@ -52,6 +52,8 @@ public class StoreResponseDiagnostics { private final String faultInjectionRuleId; private final List faultInjectionEvaluationResults; private final String endpoint; + private final String requestThroughputControlGroupName; + private final String requestThroughputControlGroupConfig; public static StoreResponseDiagnostics createStoreResponseDiagnostics( StoreResponse storeResponse, @@ -91,6 +93,9 @@ private StoreResponseDiagnostics(StoreResponse storeResponse, RxDocumentServiceR this.faultInjectionRuleId = storeResponse.getFaultInjectionRuleId(); this.faultInjectionEvaluationResults = storeResponse.getFaultInjectionRuleEvaluationResults(); this.endpoint = storeResponse.getEndpoint(); + this.requestThroughputControlGroupName = rxDocumentServiceRequest.throughputControlGroupName; + this.requestThroughputControlGroupConfig = + rxDocumentServiceRequest.requestContext.throughputControlRequestContext != null ? rxDocumentServiceRequest.requestContext.throughputControlRequestContext.getConfigString() : null; } private StoreResponseDiagnostics(CosmosException e, RxDocumentServiceRequest rxDocumentServiceRequest) { @@ -122,6 +127,9 @@ private StoreResponseDiagnostics(CosmosException e, RxDocumentServiceRequest rxD } else { this.endpoint = ""; } + this.requestThroughputControlGroupName = rxDocumentServiceRequest.throughputControlGroupName; + this.requestThroughputControlGroupConfig = + rxDocumentServiceRequest.requestContext.throughputControlRequestContext != null ? rxDocumentServiceRequest.requestContext.throughputControlRequestContext.getConfigString() : null; } public int getStatusCode() { @@ -205,4 +213,12 @@ public List getFaultInjectionEvaluationResults() { public String getEndpoint() { return this.endpoint; } + + public String getRequestThroughputControlGroupName() { + return this.requestThroughputControlGroupName; + } + + public String getRequestThroughputControlGroupConfig() { + return this.requestThroughputControlGroupConfig; + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResultDiagnostics.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResultDiagnostics.java index 6770ea63d09a..c6c7bd76ef81 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResultDiagnostics.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResultDiagnostics.java @@ -240,6 +240,14 @@ public void serialize(StoreResultDiagnostics storeResultDiagnostics, jsonGenerator.writeObjectField("channelStatistics", storeResponseDiagnostics.getRntbdChannelStatistics()); jsonGenerator.writeObjectField("serviceEndpointStatistics", storeResponseDiagnostics.getRntbdEndpointStatistics()); + this.writeNonNullStringField( + jsonGenerator, + "requestTCG", + storeResponseDiagnostics.getRequestThroughputControlGroupName()); + this.writeNonNullStringField( + jsonGenerator, + "requestTCGConfig", + storeResponseDiagnostics.getRequestThroughputControlGroupConfig()); jsonGenerator.writeEndObject(); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/IThroughputControlGroup.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/IThroughputControlGroup.java new file mode 100644 index 000000000000..46fbef4cb58b --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/IThroughputControlGroup.java @@ -0,0 +1,8 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl; + +public interface IThroughputControlGroup { + String getDiagnosticsString(); +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlRequestContext.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlRequestContext.java new file mode 100644 index 000000000000..0f025e280857 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlRequestContext.java @@ -0,0 +1,33 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl; + +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; + +import java.util.concurrent.atomic.AtomicReference; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; + +public class ThroughputControlRequestContext { + private final String configString; + private final AtomicReference throughputControlCycleId; + + public ThroughputControlRequestContext(String configString) { + checkArgument(StringUtils.isNotEmpty(configString), "Argument 'configString' cannot be null or empty."); + this.configString = configString; + this.throughputControlCycleId = new AtomicReference<>(); + } + + public String getConfigString() { + return this.configString; + } + + public String getThroughputControlCycleId() { + return this.throughputControlCycleId.get(); + } + + public void setThroughputControlCycleId(String throughputControlCycleId) { + this.throughputControlCycleId.set(throughputControlCycleId); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputRequestThrottler.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputRequestThrottler.java index 82411afdd66d..1b75c2ffd138 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputRequestThrottler.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputRequestThrottler.java @@ -92,8 +92,8 @@ public Mono processRequest(RxDocumentServiceRequest request, Mono orig })); if (this.availableThroughput.get() > 0) { - if (StringUtils.isEmpty(request.requestContext.throughputControlCycleId)) { - request.requestContext.throughputControlCycleId = this.cycleId; + if (StringUtils.isEmpty(request.requestContext.throughputControlRequestContext.getThroughputControlCycleId())) { + request.requestContext.throughputControlRequestContext.setThroughputControlCycleId(this.cycleId); } trackingUnit.increasePassedRequest(); @@ -178,7 +178,7 @@ private void trackRequestCharge (RxDocumentServiceRequest request, T respons } // If the response comes back in a different cycle, discard it. - if (StringUtils.equals(this.cycleId, request.requestContext.throughputControlCycleId)) { + if (StringUtils.equals(this.cycleId, request.requestContext.throughputControlRequestContext.getThroughputControlCycleId())) { this.availableThroughput.getAndAccumulate(requestCharge, (available, consumed) -> available - consumed); } else { if (trackingUnit != null) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/SDKThroughputControlGroupInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/SDKThroughputControlGroupInternal.java index 6d393866eab2..981fd440dd34 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/SDKThroughputControlGroupInternal.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/SDKThroughputControlGroupInternal.java @@ -5,6 +5,7 @@ import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.implementation.throughputControl.IThroughputControlGroup; import com.azure.cosmos.models.PriorityLevel; import java.util.Objects; @@ -12,7 +13,7 @@ import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; -public abstract class SDKThroughputControlGroupInternal { +public abstract class SDKThroughputControlGroupInternal implements IThroughputControlGroup { private final String groupName; private final String idPrefix; private final String id; @@ -216,6 +217,26 @@ public boolean hasSameIdentity(Object other) { && this.continueOnInitError == that.continueOnInitError; } + @Override + public String getDiagnosticsString() { + StringBuilder sb = new StringBuilder(); + sb.append("("); + sb.append("name=" + this.groupName); + sb.append(", default=" + this.isDefault); + if (this.priorityLevel != null) { + sb.append(", priorityLevel=" + this.priorityLevel); + } + if (this.targetThroughputThreshold != null) { + sb.append(", targetRUThreshold=" + this.targetThroughputThreshold); + } + if (this.targetThroughput != null) { + sb.append(", targetRU=" + this.targetThroughput); + } + + sb.append(")"); + return sb.toString(); + } + @Override public int hashCode() { int hash = 0; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/SDKThroughputGroupControllerBase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/SDKThroughputGroupControllerBase.java index 5bdc9aaed368..c9cc558050e1 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/SDKThroughputGroupControllerBase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/SDKThroughputGroupControllerBase.java @@ -11,6 +11,7 @@ import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.implementation.caches.AsyncCache; import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; +import com.azure.cosmos.implementation.throughputControl.ThroughputControlRequestContext; import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationToken; import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationTokenSource; import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; @@ -155,6 +156,11 @@ public void onContainerMaxThroughputRefresh(int maxContainerThroughput) { @Override public Mono processRequest(RxDocumentServiceRequest request, Mono originalRequestMono) { + if (request.requestContext != null) { + request.requestContext.throughputControlRequestContext = + new ThroughputControlRequestContext(this.group.getDiagnosticsString()); + } + return this.resolveRequestController() .flatMap(requestController -> { if (requestController.canHandleRequest(request)) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ServerThroughputControlGroup.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ServerThroughputControlGroup.java index 697529a93cfd..ee317cb9a096 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ServerThroughputControlGroup.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ServerThroughputControlGroup.java @@ -5,6 +5,7 @@ import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.implementation.throughputControl.IThroughputControlGroup; import com.azure.cosmos.models.PriorityLevel; import java.util.Objects; @@ -12,7 +13,7 @@ import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; -public class ServerThroughputControlGroup { +public class ServerThroughputControlGroup implements IThroughputControlGroup { private final String groupName; private final boolean isDefault; private final CosmosAsyncContainer targetContainer; @@ -28,10 +29,10 @@ public ServerThroughputControlGroup( checkArgument(StringUtils.isNotEmpty(groupName), "Argument 'groupName' cannot be null or empty."); checkNotNull(targetContainer, "Argument 'targetContainer' can not be null"); - checkArgument(throughputBucket != null && throughputBucket >= 0, "Target throughput should be no less than 0"); checkArgument( priorityLevel != null || throughputBucket != null, "At least one of 'priorityLevel' or 'throughputBucket' must be provided."); + checkArgument(throughputBucket == null || throughputBucket >= 0, "Target throughput bucket should be no less than 0"); this.groupName = groupName; this.isDefault = isDefault; @@ -75,4 +76,27 @@ public boolean equals(Object o) { public int hashCode() { return Objects.hash(groupName, isDefault, targetContainer, priorityLevel, throughputBucket); } + + /*** + * Config show up in the diagnostics. + * @return the config string. + */ + @Override + public String getDiagnosticsString() { + StringBuilder sb = new StringBuilder(); + sb.append("("); + sb.append("name=" + this.groupName); + sb.append(", default=" + this.isDefault); + + if (this.priorityLevel != null) { + sb.append(", priorityLevel=" + this.priorityLevel); + } + + if (this.throughputBucket != null) { + sb.append(", throughputBucket=" + this.throughputBucket); + } + + sb.append(")"); + return sb.toString(); + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupController.java index e2ad4ad34a58..b419b5778939 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupController.java @@ -5,6 +5,7 @@ import com.azure.cosmos.implementation.RxDocumentServiceRequest; import com.azure.cosmos.implementation.throughputControl.IThroughputController; +import com.azure.cosmos.implementation.throughputControl.ThroughputControlRequestContext; import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroup; import reactor.core.publisher.Mono; @@ -39,6 +40,12 @@ public Mono processRequest(RxDocumentServiceRequest request, Mono orig request.setThroughputBucket(this.serverThroughputControlGroup.getThroughputBucket()); } + if (request.requestContext != null) { + request.requestContext.setThroughputControlRequestContext( + new ThroughputControlRequestContext(this.serverThroughputControlGroup.getDiagnosticsString()) + ); + } + return originalRequestMono; } From 4d895ee035e859aa57ca74d0c2cb8c9fa8713a37 Mon Sep 17 00:00:00 2001 From: annie-mac Date: Mon, 28 Jul 2025 23:04:28 -0700 Subject: [PATCH 10/12] fix compile --- .../src/main/java/com/azure/cosmos/CosmosAsyncContainer.java | 1 - .../src/main/java/com/azure/cosmos/CosmosContainer.java | 1 - 2 files changed, 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index b3fb59f595ee..bcde51f15aaf 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -2826,7 +2826,6 @@ void enableGlobalThroughputControlGroup( /*** * Enable the server throughput bucket control group. - *
* * * diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java index a03f8259b734..729d407e12ea 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java @@ -1032,7 +1032,6 @@ public void enableGlobalThroughputControlGroup(ThroughputControlGroupConfig grou /*** * Enable the server throughput control group. - *
* * * From b84e9926aecb42dd04fb64d8251f51eb5fef66b7 Mon Sep 17 00:00:00 2001 From: annie-mac Date: Tue, 29 Jul 2025 09:13:24 -0700 Subject: [PATCH 11/12] resolve comments --- .../test/java/com/azure/cosmos/rx/ClientRetryPolicyE2ETests.java | 1 + .../java/com/azure/cosmos/rx/FullFidelityChangeFeedTest.java | 1 - 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ClientRetryPolicyE2ETests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ClientRetryPolicyE2ETests.java index 0db09f46aed8..5144d7bbc47d 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ClientRetryPolicyE2ETests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ClientRetryPolicyE2ETests.java @@ -697,6 +697,7 @@ private Mono performDocumentOperation( .onErrorResume(throwable -> { if (throwable instanceof CosmosException) { CosmosException cosmosException = (CosmosException) throwable; + return Mono.just(cosmosException.getDiagnostics()); } return Mono.error(throwable); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/FullFidelityChangeFeedTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/FullFidelityChangeFeedTest.java index 005e09bb0968..e26765752927 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/FullFidelityChangeFeedTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/FullFidelityChangeFeedTest.java @@ -28,7 +28,6 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.fail; -import static org.assertj.core.api.InstanceOfAssertFactories.ARRAY; public class FullFidelityChangeFeedTest extends TestSuiteBase { From e934d03b305dcabd325b32fdc8e790d5743ce93c Mon Sep 17 00:00:00 2001 From: annie-mac Date: Tue, 29 Jul 2025 12:57:51 -0700 Subject: [PATCH 12/12] update javadoc --- .../main/java/com/azure/cosmos/CosmosAsyncContainer.java | 9 +++++++++ .../src/main/java/com/azure/cosmos/CosmosContainer.java | 9 +++++++++ 2 files changed, 18 insertions(+) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index bcde51f15aaf..2fcf9bfe50cf 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -2828,6 +2828,15 @@ void enableGlobalThroughputControlGroup( * Enable the server throughput bucket control group. * * + * 

+     * ThroughputControlGroupConfig groupConfig =
+     *     new ThroughputControlGroupConfigBuilder()
+     *         .groupName("localControlGroup")
+     *         .throughputBucket(2)
+     *         .build();
+     *
+     * container.enableServerThroughputControlGroup(groupConfig);
+     *
* * * @param groupConfig the throughput control group config, see {@link ThroughputControlGroupConfig}. diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java index 729d407e12ea..04a6060c1927 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java @@ -1034,6 +1034,15 @@ public void enableGlobalThroughputControlGroup(ThroughputControlGroupConfig grou * Enable the server throughput control group. * * + * 
+     * ThroughputControlGroupConfig groupConfig =
+     *     new ThroughputControlGroupConfigBuilder()
+     *         .groupName("localControlGroup")
+     *         .throughputBucket(2)
+     *         .build();
+     *
+     * container.enableServerThroughputControlGroup(groupConfig);
+     *
* * * @param groupConfig the throughput control group config, see {@link ThroughputControlGroupConfig}.