From 39a7a6606b086ca05c1d17de15276937db207dc6 Mon Sep 17 00:00:00 2001 From: Annie Liang <64233642+xinlian12@users.noreply.github.com> Date: Tue, 2 Mar 2021 15:42:33 -0800 Subject: [PATCH] ThroughputControl-Global (#19183) * ThroughputControl-Global Co-authored-by: Annie Liang Co-authored-by: annie-mac Co-authored-by: annie-mac Co-authored-by: annie-mac --- .../java/com/azure/cosmos/BridgeInternal.java | 10 +- .../com/azure/cosmos/CosmosAsyncClient.java | 16 +- .../azure/cosmos/CosmosAsyncContainer.java | 69 +--- .../java/com/azure/cosmos/CosmosClient.java | 13 + .../com/azure/cosmos/CosmosContainer.java | 49 +-- .../cosmos/GlobalThroughputControlConfig.java | 71 ++++ .../GlobalThroughputControlConfigBuilder.java | 81 ++++ .../cosmos/ThroughputControlGroupConfig.java | 73 ++++ .../ThroughputControlGroupConfigBuilder.java | 97 +++++ .../implementation/AsyncDocumentClient.java | 4 +- .../cosmos/implementation/HttpConstants.java | 4 + .../implementation/RxDocumentClientImpl.java | 16 +- .../RxDocumentServiceRequest.java | 6 +- .../implementation/RxGatewayStoreModel.java | 28 +- .../directconnectivity/TransportClient.java | 3 + .../DocumentQueryExecutionContextBase.java | 10 +- ...llelDocumentQueryExecutionContextBase.java | 11 +- .../LinkedCancellationToken.java | 47 +++ .../LinkedCancellationTokenSource.java | 51 +++ .../ThroughputControlMode.java | 10 - .../ThroughputControlStore.java | 140 ++++--- .../ThroughputRequestThrottler.java | 96 +++-- .../config/GlobalThroughputControlGroup.java | 87 +++++ .../config/LocalThroughputControlGroup.java | 18 + .../config/ThroughputControlGroupFactory.java | 47 +++ .../ThroughputControlGroupInternal.java} | 39 +- .../controller/IThroughputController.java | 6 - .../EmptyThroughputContainerController.java | 5 - .../ThroughputContainerController.java | 202 ++++++---- .../ThroughputProvisioningScope.java} | 4 +- .../group/ThroughputGroupControllerBase.java | 145 ++++---- .../ThroughputGroupControllerFactory.java | 38 +- .../group/ThroughputGroupLocalController.java | 23 -- .../GlobalThroughputControlClientItem.java | 62 ++++ .../GlobalThroughputControlConfigItem.java | 107 ++++++ ...lobalThroughputControlGroupController.java | 126 +++++++ .../global/GlobalThroughputControlItem.java | 64 ++++ .../ThroughputControlContainerManager.java | 213 +++++++++++ .../group/global/ThroughputUsageSnapshot.java | 42 +++ ...LocalThroughputControlGroupController.java | 47 +++ .../GlobalThroughputRequestController.java | 44 +-- .../request/IThroughputRequestController.java | 8 +- .../PkRangesThroughputRequestController.java | 84 ++--- .../CosmosChangeFeedRequestOptions.java | 24 ++ .../models/CosmosItemRequestOptions.java | 4 +- .../models/CosmosQueryRequestOptions.java | 24 ++ .../src/main/java/module-info.java | 1 + .../cosmos/ThroughputControlCodeSnippet.java | 56 +++ .../directconnectivity/ReflectionUtils.java | 16 +- .../LinkedCancellationTokenSourceTests.java | 28 ++ ...tControlGroupConfigConfigurationTests.java | 54 +++ ...oughputControlGroupConfigurationTests.java | 41 -- .../ThroughputControlTests.java | 351 +++++++++++++----- ...lobalThroughputRequestControllerTests.java | 77 +--- ...angesThroughputRequestControllerTests.java | 74 +--- .../com/azure/cosmos/rx/TestSuiteBase.java | 111 ++++-- 56 files changed, 2339 insertions(+), 838 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/GlobalThroughputControlConfig.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/GlobalThroughputControlConfigBuilder.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfig.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfigBuilder.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationToken.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationTokenSource.java delete mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlMode.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/GlobalThroughputControlGroup.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/LocalThroughputControlGroup.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/ThroughputControlGroupFactory.java rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/{ThroughputControlGroup.java => implementation/throughputControl/config/ThroughputControlGroupInternal.java} (71%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/{ThroughputResolveLevel.java => controller/container/ThroughputProvisioningScope.java} (53%) delete mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/ThroughputGroupLocalController.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlClientItem.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlConfigItem.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlGroupController.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlItem.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/ThroughputControlContainerManager.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/ThroughputUsageSnapshot.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/local/LocalThroughputControlGroupController.java create mode 100644 sdk/cosmos/azure-cosmos/src/samples/java/com/azure/cosmos/ThroughputControlCodeSnippet.java create mode 100644 sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationTokenSourceTests.java create mode 100644 sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupConfigConfigurationTests.java delete mode 100644 sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupConfigurationTests.java diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java index ef2cfec1ed47a..f0e2d10dbc572 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java @@ -36,7 +36,6 @@ import com.azure.cosmos.implementation.query.QueryInfo; import com.azure.cosmos.implementation.query.metrics.ClientSideMetrics; import com.azure.cosmos.implementation.routing.PartitionKeyInternal; -import com.azure.cosmos.implementation.throughputControl.ThroughputControlMode; import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.CosmosStoredProcedureProperties; import com.azure.cosmos.models.FeedResponse; @@ -796,18 +795,13 @@ public static List getPatchOperationsFromCosmosPatch(CosmosPatch return cosmosPatchOperations.getPatchOperations(); } - @Warning(value = INTERNAL_USE_ONLY_WARNING) - public static ThroughputControlMode getThroughputControlMode(ThroughputControlGroup throughputControlGroup) { - return throughputControlGroup.getControlMode(); - } - @Warning(value = INTERNAL_USE_ONLY_WARNING) public static SqlQuerySpec getOfferQuerySpecFromResourceId(CosmosAsyncContainer container, String resourceId) { return container.getDatabase().getOfferQuerySpecFromResourceId(resourceId); } @Warning(value = INTERNAL_USE_ONLY_WARNING) - public static CosmosAsyncContainer getTargetContainerFromThroughputControlGroup(ThroughputControlGroup controlGroup) { - return controlGroup.getTargetContainer(); + public static CosmosAsyncContainer getControlContainerFromThroughputGlobalControlConfig(GlobalThroughputControlConfig globalControlConfig) { + return globalControlConfig.getControlContainer(); } } 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 3dcb0f086e4bd..affc31be4eb38 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 @@ -15,6 +15,7 @@ import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.TracerProvider; import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdMetrics; +import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal; import com.azure.cosmos.models.CosmosDatabaseProperties; import com.azure.cosmos.models.CosmosDatabaseRequestOptions; import com.azure.cosmos.models.CosmosDatabaseResponse; @@ -23,6 +24,7 @@ import com.azure.cosmos.models.ModelBridgeInternal; import com.azure.cosmos.models.SqlQuerySpec; import com.azure.cosmos.models.ThroughputProperties; +import com.azure.cosmos.util.Beta; import com.azure.cosmos.util.CosmosPagedFlux; import com.azure.cosmos.util.UtilBridgeInternal; import io.micrometer.core.instrument.MeterRegistry; @@ -471,11 +473,23 @@ TracerProvider getTracerProvider(){ * * @param group Throughput control group going to be enabled. */ - void enableThroughputControlGroup(ThroughputControlGroup group) { + void enableThroughputControlGroup(ThroughputControlGroupInternal group) { checkNotNull(group, "Throughput control group cannot be null"); this.asyncDocumentClient.enableThroughputControlGroup(group); } + /** + * Create global throughput control config builder which will be used to build {@link GlobalThroughputControlConfig}. + * + * @param databaseId The database id of the control container. + * @param containerId The container id of the control container. + * @return A {@link GlobalThroughputControlConfigBuilder}. + */ + @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public GlobalThroughputControlConfigBuilder createGlobalThroughputControlConfigBuilder(String databaseId, String containerId) { + return new GlobalThroughputControlConfigBuilder(this, databaseId, containerId); + } + private CosmosPagedFlux queryDatabasesInternal(SqlQuerySpec querySpec, CosmosQueryRequestOptions options){ return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { pagedFluxOptions.setTracerInformation(this.tracerProvider, "queryDatabases", this.serviceEndpoint, null); 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 1f65517b9cfe1..2a2717895e4e9 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 @@ -21,7 +21,9 @@ import com.azure.cosmos.implementation.batch.BatchExecutor; import com.azure.cosmos.implementation.batch.BulkExecutor; import com.azure.cosmos.implementation.query.QueryInfo; -import com.azure.cosmos.implementation.throughputControl.ThroughputControlMode; +import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupFactory; +import com.azure.cosmos.implementation.throughputControl.config.GlobalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.config.LocalThroughputControlGroup; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; import com.azure.cosmos.models.CosmosConflictProperties; import com.azure.cosmos.models.CosmosContainerProperties; @@ -51,7 +53,6 @@ import static com.azure.core.util.FluxUtil.withContext; import static com.azure.cosmos.implementation.Utils.getEffectiveCosmosChangeFeedRequestOptions; import static com.azure.cosmos.implementation.Utils.setContinuationTokenAndMaxItemCount; -import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; /** @@ -1436,65 +1437,35 @@ public Mono> getFeedRanges() { } /** + * Enable the throughput control group with local control mode. * - * @param groupName The throughput control group name. - * @param targetThroughput The target throughput for the control group. + * {@codesnippet com.azure.cosmos.throughputControl.localControl} * - * @return A {@link ThroughputControlGroup}. + * @param groupConfig A {@link ThroughputControlGroupConfig}. */ @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) - ThroughputControlGroup enableThroughputLocalControlGroup(String groupName, int targetThroughput) { - return this.enableThroughputLocalControlGroup(groupName, targetThroughput, false); + public void enableLocalThroughputControlGroup(ThroughputControlGroupConfig groupConfig) { + LocalThroughputControlGroup localControlGroup = ThroughputControlGroupFactory.createThroughputLocalControlGroup(groupConfig, this); + this.database.getClient().enableThroughputControlGroup(localControlGroup); } /** + * Enable the throughput control group with global control mode. + * The defined throughput limit will be shared across different clients. * - * @param groupName The throughput control group name. - * @param targetThroughput The target throughput for the control group. - * @param isDefault Flag to indicate whether this group will be used as default. + * {@codesnippet com.azure.cosmos.throughputControl.globalControl} * - * @return A {@link ThroughputControlGroup}. + * @param groupConfig The throughput control group configuration, see {@link GlobalThroughputControlGroup}. + * @param globalControlConfig The global throughput control configuration, see {@link GlobalThroughputControlConfig}. */ @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) - ThroughputControlGroup enableThroughputLocalControlGroup(String groupName, int targetThroughput, boolean isDefault) { - return this.enableThroughputControlGroup(groupName, targetThroughput, null, ThroughputControlMode.LOCAL, isDefault); - } - - /** - * - * @param groupName The throughput control group name. - * @param targetThroughputThreshold The target throughput threshold for the control group. - * - * @return A {@link ThroughputControlGroup}. - */ - @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) - ThroughputControlGroup enableThroughputLocalControlGroup(String groupName, double targetThroughputThreshold) { - return this.enableThroughputLocalControlGroup(groupName, targetThroughputThreshold, false); - } - - /** - * - * @param groupName The throughput control group name. - * @param targetThroughputThreshold The target throughput threshold for the control group. - * @param isDefault Flag to indicate whether this group will be used as default. - * @return A {@link ThroughputControlGroup}. - */ - @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) - ThroughputControlGroup enableThroughputLocalControlGroup(String groupName, double targetThroughputThreshold, boolean isDefault) { - return this.enableThroughputControlGroup(groupName, null, targetThroughputThreshold, ThroughputControlMode.LOCAL, isDefault); - } - - private ThroughputControlGroup enableThroughputControlGroup( - String groupName, - Integer targetThroughput, - Double targetThroughputThreshold, - ThroughputControlMode controlMode, - boolean isDefault) { + public void enableGlobalThroughputControlGroup( + ThroughputControlGroupConfig groupConfig, + GlobalThroughputControlConfig globalControlConfig) { - ThroughputControlGroup throughputControlGroup = new ThroughputControlGroup( - groupName, this, targetThroughput, targetThroughputThreshold, controlMode, isDefault); - this.database.getClient().enableThroughputControlGroup(throughputControlGroup); + GlobalThroughputControlGroup globalControlGroup = + ThroughputControlGroupFactory.createThroughputGlobalControlGroup(groupConfig, globalControlConfig, this); - return throughputControlGroup; + this.database.getClient().enableThroughputControlGroup(globalControlGroup); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClient.java index 5098a5966e5e1..a10946877a564 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClient.java @@ -10,6 +10,7 @@ import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.SqlQuerySpec; import com.azure.cosmos.models.ThroughputProperties; +import com.azure.cosmos.util.Beta; import com.azure.cosmos.util.CosmosPagedFlux; import com.azure.cosmos.util.CosmosPagedIterable; import com.azure.cosmos.util.UtilBridgeInternal; @@ -213,4 +214,16 @@ public void close() { private CosmosPagedIterable getCosmosPagedIterable(CosmosPagedFlux cosmosPagedFlux) { return UtilBridgeInternal.createCosmosPagedIterable(cosmosPagedFlux); } + + /** + * Create global throughput control config builder which will be used to build {@link GlobalThroughputControlConfig}. + * + * @param databaseId The database id of the control container. + * @param containerId The container id of the control container. + * @return A {@link GlobalThroughputControlConfigBuilder}. + */ + @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public GlobalThroughputControlConfigBuilder createGlobalThroughputControlConfigBuilder(String databaseId, String containerId) { + return new GlobalThroughputControlConfigBuilder(this.asyncClientWrapper, databaseId, containerId); + } } 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 2788b415df595..5d85842394633 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,6 +3,7 @@ package com.azure.cosmos; +import com.azure.cosmos.implementation.throughputControl.config.GlobalThroughputControlGroup; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; import com.azure.cosmos.models.CosmosItemIdentity; import com.azure.cosmos.models.CosmosItemResponse; @@ -744,52 +745,28 @@ public List getFeedRanges() { } /** + * Enable the throughput control group with local control mode. * - * @param groupName The throughput control group name. - * @param targetThroughput The target throughput for the control group. + * {@codesnippet com.azure.cosmos.throughputControl.localControl} * - * @return A {@link ThroughputControlGroup}. + * @param groupConfig A {@link GlobalThroughputControlConfig}. */ @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) - ThroughputControlGroup enableThroughputLocalControlGroup(String groupName, int targetThroughput) { - return this.asyncContainer.enableThroughputLocalControlGroup(groupName, targetThroughput); + public void enableLocalThroughputControlGroup(ThroughputControlGroupConfig groupConfig) { + this.asyncContainer.enableLocalThroughputControlGroup(groupConfig); } /** + * Enable the throughput control group with global control mode. + * The defined throughput limit will be shared across different clients. * - * @param groupName The throughput control group name. - * @param targetThroughput The target throughput for the control group. - * @param isDefault Flag to indicate whether this group will be used as default. + * {@codesnippet com.azure.cosmos.throughputControl.globalControl} * - * @return A {@link ThroughputControlGroup}. + * @param groupConfig The throughput control group configuration, see {@link GlobalThroughputControlGroup}. + * @param globalControlConfig The global throughput control configuration, see {@link GlobalThroughputControlConfig}. */ @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) - ThroughputControlGroup enableThroughputLocalControlGroup(String groupName, int targetThroughput, boolean isDefault) { - return this.asyncContainer.enableThroughputLocalControlGroup(groupName, targetThroughput, isDefault); - } - - /** - * - * @param groupName The throughput control group name. - * @param targetThroughputThreshold The target throughput threshold for the control group. - * - * @return A {@link ThroughputControlGroup}. - */ - @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) - ThroughputControlGroup enableThroughputLocalControlGroup(String groupName, double targetThroughputThreshold) { - return this.asyncContainer.enableThroughputLocalControlGroup(groupName, targetThroughputThreshold); - } - - /** - * - * @param groupName The throughput control group name. - * @param targetThroughputThreshold The target throughput threshold for the control group. - * @param isDefault Flag to indicate whether this group will be used as default. - * - * @return A {@link ThroughputControlGroup}. - */ - @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) - ThroughputControlGroup enableThroughputLocalControlGroup(String groupName, double targetThroughputThreshold, boolean isDefault) { - return this.asyncContainer.enableThroughputLocalControlGroup(groupName, targetThroughputThreshold, isDefault); + public void enableGlobalThroughputControlGroup(ThroughputControlGroupConfig groupConfig, GlobalThroughputControlConfig globalControlConfig) { + this.asyncContainer.enableGlobalThroughputControlGroup(groupConfig, globalControlConfig); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/GlobalThroughputControlConfig.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/GlobalThroughputControlConfig.java new file mode 100644 index 0000000000000..28acc8fca38e0 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/GlobalThroughputControlConfig.java @@ -0,0 +1,71 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos; + +import com.azure.cosmos.util.Beta; + +import java.time.Duration; + +/** + * This configuration is used for throughput global control mode. + * It contains configuration about the extra container which will track all the clients throughput usage for a certain control group. + */ +@Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) +public class GlobalThroughputControlConfig { + private final CosmosAsyncContainer controlContainer; + private final Duration controlItemRenewInterval; + private final Duration controlItemExpireInterval; + + GlobalThroughputControlConfig( + CosmosAsyncContainer controlContainer, + Duration controlItemRenewInterval, + Duration controlItemExpireInterval) { + + this.controlContainer = controlContainer; + this.controlItemRenewInterval = controlItemRenewInterval; + this.controlItemExpireInterval = controlItemExpireInterval; + } + + /** + * Get the control container. + * This is the container to track all other clients throughput usage. + * + * @return The {@link CosmosAsyncContainer}. + */ + CosmosAsyncContainer getControlContainer() { + return controlContainer; + } + + /** + * Get the control item renew interval. + * + * This controls how often the client is going to update the throughput usage of itself + * and adjust its own throughput share based on the throughput usage of other clients. + * + * In short words, it controls how quickly the shared throughput will reload balanced across different clients. + * + * By default, it is 5s. + * + * @return The control item renew interval. + */ + @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public Duration getControlItemRenewInterval() { + return this.controlItemRenewInterval; + } + + /** + * Get the control item expire interval. + * + * A client may be offline due to various reasons (being shutdown, network issue... ). + * This controls how quickly we will detect the client has been offline and hence allow its throughput share to be taken by other clients. + * + * By default, it is 11s. + * + * @return The control item renew interval. + */ + @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public Duration getControlItemExpireInterval() { + return this.controlItemExpireInterval; + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/GlobalThroughputControlConfigBuilder.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/GlobalThroughputControlConfigBuilder.java new file mode 100644 index 0000000000000..e7cca8226c8d5 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/GlobalThroughputControlConfigBuilder.java @@ -0,0 +1,81 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos; + +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.util.Beta; + +import java.time.Duration; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +/** + * Throughput global control config builder. + */ +@Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) +public class GlobalThroughputControlConfigBuilder { + private final CosmosAsyncContainer controlContainer; + + private Duration controlItemRenewInterval; + private Duration controlItemExpireInterval; + + GlobalThroughputControlConfigBuilder(CosmosAsyncClient client, String databaseId, String containerId) { + checkNotNull(client, "Client can not be null"); + checkArgument(StringUtils.isNotEmpty(databaseId), "DatabaseId cannot be null nor empty"); + checkArgument(StringUtils.isNotEmpty(containerId), "ContainerId cannot be null nor empty"); + + controlContainer = client.getDatabase(databaseId).getContainer(containerId); + } + + /** + * Set the control item renew interval. + * + * This controls how often the client is going to update the throughput usage of itself + * and adjust its own throughput share based on the throughput usage of other clients. + * + * In short words, it controls how quickly the shared throughput will reload balanced across different clients. + * + * @param controlItemRenewInterval The control item renewal interval. + * @return The {@link GlobalThroughputControlConfigBuilder} + */ + @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public GlobalThroughputControlConfigBuilder setControlItemRenewInterval(Duration controlItemRenewInterval) { + this.controlItemRenewInterval = controlItemRenewInterval; + return this; + } + + /** + * Set the control item expire interval. + * + * A client may be offline due to various reasons (being shutdown, network issue... ). + * This controls how quickly we will detect the client has been offline and hence allow its throughput share to be taken by other clients. + * + * @param controlItemExpireInterval The control item expire interval. + * @return The {@link GlobalThroughputControlConfigBuilder} + */ + @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public GlobalThroughputControlConfigBuilder setControlItemExpireInterval(Duration controlItemExpireInterval) { + this.controlItemExpireInterval = controlItemExpireInterval; + return this; + } + + /** + * Validate the throughput global control configuration and create a new throughput global control config item. + * + * @return A new {@link GlobalThroughputControlConfig}. + */ + @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public GlobalThroughputControlConfig build() { + if (this.controlItemExpireInterval != null && this.controlItemRenewInterval != null) { + // expireInterval will be used as the ttl of ThroughputGlobalControlClientItem, it should be a reasonable value + // to make sure the item is not get deleted too fast + if (this.controlItemExpireInterval.getSeconds() < 2 * this.controlItemRenewInterval.getSeconds() + 1) { + throw new IllegalArgumentException("ControlItemExpireInterval is too small compared to ControlItemExpireInterval"); + } + } + + return new GlobalThroughputControlConfig(this.controlContainer, this.controlItemRenewInterval, this.controlItemExpireInterval); + } +} 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 new file mode 100644 index 0000000000000..bef6e91b55542 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfig.java @@ -0,0 +1,73 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos; + +import com.azure.cosmos.util.Beta; + +/** + * Throughput control group configuration. + */ +@Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) +public final class ThroughputControlGroupConfig { + private final String groupName; + private final Integer targetThroughput; + private final Double targetThroughputThreshold; + private final boolean isDefault; + + ThroughputControlGroupConfig(String groupName, Integer targetThroughput, Double targetThroughputThreshold, boolean isDefault) { + this.groupName= groupName; + this.targetThroughput = targetThroughput; + this.targetThroughputThreshold = targetThroughputThreshold; + this.isDefault = isDefault; + } + + /** + * Get the throughput control group name. + * + * @return the group name. + */ + @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public String getGroupName() { + return this.groupName; + } + + /** + * Get throughput control group target throughput. + * Since we allow either TargetThroughput or TargetThroughputThreshold, this value can be null. + * + * By default, it is null. + * + * @return the target throughput. + */ + @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public Integer getTargetThroughput() { + return this.targetThroughput; + } + + /** + * Get the throughput control group target throughput threshold. + * Since we allow either TargetThroughput or TargetThroughputThreshold, this value can be null. + * + * By default, this value is null. + * + * @return the target throughput threshold. + */ + @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public Double getTargetThroughputThreshold() { + return this.targetThroughputThreshold; + } + + /** + * Get whether this throughput control group will be used by default. + * + * By default, it is false. + * If it is true, requests without explicit override of the throughput control group will be routed to this group. + * + * @return {@code true} this throughput control group will be used by default unless being override. {@code false} otherwise. + */ + @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public boolean isDefault() { + return this.isDefault; + } +} 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 new file mode 100644 index 0000000000000..51870e5be734d --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfigBuilder.java @@ -0,0 +1,97 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos; + +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.util.Beta; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; + +/** + * The throughput control group config builder. + */ +@Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) +public class ThroughputControlGroupConfigBuilder { + private String groupName; + private Integer targetThroughput; + private Double targetThroughputThreshold; + private boolean isDefault; + + /** + * Set the throughput control group name. + * + * @param groupName The throughput control group name. + * @return The {@link ThroughputControlGroupConfigBuilder}. + */ + @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public ThroughputControlGroupConfigBuilder setGroupName(String groupName) { + checkArgument(StringUtils.isNotEmpty(groupName), "Group name cannot be null nor empty"); + + this.groupName = groupName; + return this; + } + + /** + * Set the throughput control group target throughput. + * + * The target throughput value should be greater than 0. + * + * @param targetThroughput The target throughput for the control group. + * @return The {@link ThroughputControlGroupConfigBuilder}. + */ + @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public ThroughputControlGroupConfigBuilder setTargetThroughput(int targetThroughput) { + checkArgument(targetThroughput > 0, "Target throughput should be greater than 0"); + + this.targetThroughput = targetThroughput; + return this; + } + + /** + * Set the throughput control group target throughput threshold. + * + * The target throughput threshold value should be between (0, 1]. + * + * @param targetThroughputThreshold The target throughput threshold for the control group. + * @return The {@link ThroughputControlGroupConfigBuilder}. + */ + @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public ThroughputControlGroupConfigBuilder setTargetThroughputThreshold(double targetThroughputThreshold) { + checkArgument(targetThroughputThreshold > 0 && targetThroughputThreshold <= 1, "Target throughput threshold should between (0, 1]"); + + this.targetThroughputThreshold = targetThroughputThreshold; + return this; + } + + + /** + * Set whether this throughput control group will be used by default. + * If set to true, requests without explicit override of the throughput control group will be routed to this group. + * + * @param aDefault The flag to indicate whether the throughput control group will be used by default. + * @return The {@link ThroughputControlGroupConfigBuilder}. + */ + @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public ThroughputControlGroupConfigBuilder setDefault(boolean aDefault) { + isDefault = aDefault; + return this; + } + + /** + * Validate the throughput configuration and create a new throughput control group config item. + * + * @return A new {@link ThroughputControlGroupConfig}. + */ + @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + 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) { + throw new IllegalArgumentException("Neither targetThroughput nor targetThroughputThreshold is defined."); + } + + return new ThroughputControlGroupConfig(groupName, this.targetThroughput, this.targetThroughputThreshold, isDefault); + } +} 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 39e6193e6840b..b6a614ca97cc6 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 @@ -6,7 +6,6 @@ import com.azure.core.credential.TokenCredential; import com.azure.cosmos.ConsistencyLevel; import com.azure.cosmos.CosmosPatchOperations; -import com.azure.cosmos.ThroughputControlGroup; import com.azure.cosmos.TransactionalBatchResponse; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.implementation.batch.ServerBatchRequest; @@ -14,6 +13,7 @@ import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; import com.azure.cosmos.implementation.clientTelemetry.ClientTelemetry; import com.azure.cosmos.implementation.query.PartitionedQueryExecutionInfo; +import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; import com.azure.cosmos.models.CosmosItemIdentity; import com.azure.cosmos.models.CosmosQueryRequestOptions; @@ -1510,5 +1510,5 @@ Flux> readAllDocuments( * * @param group the throughput control group. */ - void enableThroughputControlGroup(ThroughputControlGroup group); + void enableThroughputControlGroup(ThroughputControlGroupInternal group); } 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 6c9648a59f113..1007d36f19aaf 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 @@ -316,6 +316,7 @@ public static class SubStatusCodes { // 404: LSN in session token is higher public static final int READ_SESSION_NOT_AVAILABLE = 1002; + public static final int OWNER_RESOURCE_NOT_EXISTS = 1003; // Client generated gateway network error substatus public static final int GATEWAY_ENDPOINT_UNAVAILABLE = 10001; @@ -325,6 +326,9 @@ public static class SubStatusCodes { // Client generated request rate too large exception public static final int THROUGHPUT_CONTROL_REQUEST_RATE_TOO_LARGE = 10003; + + // Client generated offer not configured exception + public static final int OFFER_NOT_CONFIGURED = 10004; } public static class HeaderValues { 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 878ffb10886b7..78216e244c336 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 @@ -12,7 +12,6 @@ import com.azure.cosmos.CosmosDiagnostics; import com.azure.cosmos.CosmosPatchOperations; import com.azure.cosmos.DirectConnectionConfig; -import com.azure.cosmos.ThroughputControlGroup; import com.azure.cosmos.TransactionalBatchResponse; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.implementation.batch.BatchResponseParser; @@ -50,6 +49,7 @@ import com.azure.cosmos.implementation.routing.PartitionKeyRangeIdentity; import com.azure.cosmos.implementation.routing.Range; import com.azure.cosmos.implementation.throughputControl.ThroughputControlStore; +import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; import com.azure.cosmos.models.CosmosItemIdentity; import com.azure.cosmos.models.CosmosQueryRequestOptions; @@ -66,7 +66,6 @@ import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; -import reactor.core.scheduler.Schedulers; import java.io.IOException; import java.io.UnsupportedEncodingException; @@ -280,6 +279,7 @@ private RxDocumentClientImpl(URI serviceEndpoint, this.diagnosticsClientConfig.withConnectionSharingAcrossClientsEnabled(connectionSharingAcrossClientsEnabled); this.diagnosticsClientConfig.withConsistency(consistencyLevel); + this.throughputControlEnabled = new AtomicBoolean(false); logger.info( "Initializing DocumentClient [{}] with" @@ -350,7 +350,6 @@ private RxDocumentClientImpl(URI serviceEndpoint, this.resetSessionTokenRetryPolicy = retryPolicy; CpuMemoryMonitor.register(this); this.queryPlanCache = new ConcurrentHashMap<>(); - this.throughputControlEnabled = new AtomicBoolean(false); } catch (RuntimeException e) { logger.error("unexpected failure in initializing client.", e); close(); @@ -2218,7 +2217,7 @@ public Flux> queryDocumentChangeFeed( this, ResourceType.Document, Document.class, - collection.getSelfLink(), + collection.getAltLink(), collection.getResourceId(), changeFeedOptions); @@ -3759,6 +3758,12 @@ public void close() { LifeCycleUtils.closeQuietly(this.reactorHttpClient); logger.info("Shutting down CpuMonitor ..."); CpuMemoryMonitor.unregister(this); + + if (this.throughputControlEnabled.get()) { + logger.info("Closing ThroughputControlStore ..."); + this.throughputControlStore.close(); + } + logger.info("Shutting down completed."); } else { logger.warn("Already shutdown!"); @@ -3771,7 +3776,7 @@ public ItemDeserializer getItemDeserializer() { } @Override - public void enableThroughputControlGroup(ThroughputControlGroup group) { + public void enableThroughputControlGroup(ThroughputControlGroupInternal group) { checkNotNull(group, "Throughput control group can not be null"); if (this.throughputControlEnabled.compareAndSet(false, true)) { @@ -3779,7 +3784,6 @@ public void enableThroughputControlGroup(ThroughputControlGroup group) { new ThroughputControlStore( this.collectionCache, this.connectionPolicy.getConnectionMode(), - this.globalEndpointManager, this.partitionKeyRangeCache); this.storeModel.enableThroughputControl(throughputControlStore); 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 25d48a219ad74..ce66909dde7f7 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 @@ -1058,9 +1058,11 @@ private static String getThroughputControlGroupName(Object options) { return null; } else if (options instanceof RequestOptions) { return ((RequestOptions) options).getThroughputControlGroupName(); + } else if (options instanceof CosmosQueryRequestOptions) { + return ((CosmosQueryRequestOptions) options).getThroughputControlGroupName(); + } else if (options instanceof CosmosChangeFeedRequestOptions) { + return ((CosmosChangeFeedRequestOptions) options).getThroughputControlGroupName(); } else { - // TODO: add for query and changeFeed - // TODO: tracked by item https://github.com/Azure/azure-sdk-for-java/issues/18775 return null; } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java index 68ca3e6f27e13..321b00d90b5e8 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java @@ -139,14 +139,18 @@ private Mono query(RxDocumentServiceRequest request) public Mono performRequest(RxDocumentServiceRequest request, HttpMethod method) { try { + if (request.requestContext.cosmosDiagnostics == null) { + request.requestContext.cosmosDiagnostics = clientContext.createDiagnostics(); + } + URI uri = getUri(request); request.requestContext.resourcePhysicalAddress = uri.toString(); if (this.throughputControlStore != null) { - return this.throughputControlStore.processRequest(request, performRequestInternal(request, method)); + return this.throughputControlStore.processRequest(request, performRequestInternal(request, method, uri)); } - return this.performRequestInternal(request, method); + return this.performRequestInternal(request, method, uri); } catch (Exception e) { return Mono.error(e); } @@ -157,24 +161,20 @@ public Mono performRequest(RxDocumentServiceRequest r * * @param request * @param method + * @param requestUri * @return Flux */ - public Mono performRequestInternal(RxDocumentServiceRequest request, HttpMethod method) { + public Mono performRequestInternal(RxDocumentServiceRequest request, HttpMethod method, URI requestUri) { try { - if (request.requestContext.cosmosDiagnostics == null) { - request.requestContext.cosmosDiagnostics = clientContext.createDiagnostics(); - } - - URI uri = getUri(request); HttpHeaders httpHeaders = this.getHttpRequestHeaders(request.getHeaders()); Flux contentAsByteArray = request.getContentAsByteArrayFlux(); HttpRequest httpRequest = new HttpRequest(method, - uri, - uri.getPort(), + requestUri, + requestUri.getPort(), httpHeaders, contentAsByteArray); @@ -426,6 +426,11 @@ public Mono processMessage(RxDocumentServiceRequest r this.captureSessionToken(request, dce.getResponseHeaders()); } + if (Exceptions.isThroughputControlRequestRateTooLargeException(dce)) { + BridgeInternal.recordGatewayResponse(request.requestContext.cosmosDiagnostics, request, null, dce); + BridgeInternal.setCosmosDiagnostics(dce, request.requestContext.cosmosDiagnostics); + } + return Mono.error(dce); } ).map(response -> @@ -438,7 +443,8 @@ public Mono processMessage(RxDocumentServiceRequest r @Override public void enableThroughputControl(ThroughputControlStore throughputControlStore) { - this.throughputControlStore = throughputControlStore; + // no-op + // Disable throughput control for gateway mode } private void captureSessionToken(RxDocumentServiceRequest request, Map responseHeaders) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/TransportClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/TransportClient.java index 819d7e4f946e8..c077b2adc29fe 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/TransportClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/TransportClient.java @@ -23,6 +23,9 @@ public void enableThroughputControl(ThroughputControlStore throughputControlStor // Uses requests's ResourceOperation to determine the operation public Mono invokeResourceOperationAsync(Uri physicalAddress, RxDocumentServiceRequest request) { + if (StringUtils.isEmpty(request.requestContext.resourcePhysicalAddress)) { + request.requestContext.resourcePhysicalAddress = physicalAddress.toString(); + } if (this.throughputControlStore != null) { return this.throughputControlStore.processRequest(request, this.invokeStoreAsync(physicalAddress, request)); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextBase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextBase.java index b08899f6890e2..75b4177beee95 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextBase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextBase.java @@ -89,10 +89,13 @@ protected RxDocumentServiceRequest createDocumentServiceRequest(Map>> executeFunc = (request) -> { @@ -147,7 +143,8 @@ protected void initializeReadMany( querySpec, null, feedRangeEpk, - collectionRid); + collectionRid, + cosmosQueryRequestOptions.getThroughputControlGroupName()); }; Function>> executeFunc = (request) -> { 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/LinkedCancellationToken.java new file mode 100644 index 0000000000000..c53ff1316fd06 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationToken.java @@ -0,0 +1,47 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; + +public class LinkedCancellationToken { + private final List childTokenSourceList; + private final LinkedCancellationTokenSource tokenSource; + private final AtomicBoolean cancellationRequested; + + public LinkedCancellationToken(LinkedCancellationTokenSource tokenSource) { + this.childTokenSourceList = new ArrayList<>(); + this.tokenSource = tokenSource; + this.cancellationRequested = new AtomicBoolean(); + } + + public void register(LinkedCancellationTokenSource childTokenSource) { + synchronized (this) { + if (this.cancellationRequested.get()) { + throw new IllegalStateException("The cancellation token has been cancelled"); + } + + this.childTokenSourceList.add(childTokenSource); + } + } + + public void cancel() { + synchronized (this) { + if (this.cancellationRequested.compareAndSet(false, true)) { + for (LinkedCancellationTokenSource childTokenSource : this.childTokenSourceList) { + childTokenSource.close(); + } + + childTokenSourceList.clear(); + } + } + } + + public boolean isCancellationRequested() { + return this.cancellationRequested.get() + || this.tokenSource.isClosed(); + } +} 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/LinkedCancellationTokenSource.java new file mode 100644 index 0000000000000..0bbfd014fb7d2 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationTokenSource.java @@ -0,0 +1,51 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl; + +import java.io.Closeable; + +/** + * Used in the throughput controller hierarchy. + * + * There are many cases a higher level controller need to close the lower level controller and create a new one. + * A {@link LinkedCancellationToken} will be passed down to lower level controller so to cancel all the underlying tasks. + */ +public class LinkedCancellationTokenSource implements Closeable { + private boolean tokenSourceClosed; + private LinkedCancellationToken parentToken; + + public LinkedCancellationTokenSource() { + this(null); + } + + public LinkedCancellationTokenSource(LinkedCancellationToken parent) { + this.tokenSourceClosed = false; + if (parent != null) { + parent.register(this); + this.parentToken = parent; + } + } + + public LinkedCancellationToken getToken() { + synchronized (this) { + if (this.tokenSourceClosed) { + throw new IllegalStateException("The cancellation token resource has been closed"); + } + + return new LinkedCancellationToken(this); + } + } + + public boolean isClosed() { + synchronized (this) { + return this.tokenSourceClosed || (this.parentToken != null && this.parentToken.isCancellationRequested()); + } + } + @Override + public void close() { + synchronized (this) { + this.tokenSourceClosed = true; + } + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlMode.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlMode.java deleted file mode 100644 index 3e3ba764eaaab..0000000000000 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlMode.java +++ /dev/null @@ -1,10 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -package com.azure.cosmos.implementation.throughputControl; - -public enum ThroughputControlMode { - LOCAL - - // TODO: add distributed -} 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 64448255ac311..31b9ca2f1f7fc 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 @@ -6,7 +6,6 @@ import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.ConnectionMode; import com.azure.cosmos.CosmosException; -import com.azure.cosmos.ThroughputControlGroup; import com.azure.cosmos.implementation.GlobalEndpointManager; import com.azure.cosmos.implementation.ResourceType; import com.azure.cosmos.implementation.RxDocumentServiceRequest; @@ -15,6 +14,7 @@ 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; @@ -22,11 +22,8 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.Exceptions; -import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; -import reactor.core.scheduler.Schedulers; -import java.util.Collections; import java.util.HashSet; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; @@ -84,53 +81,55 @@ public class ThroughputControlStore { private final RxClientCollectionCache collectionCache; private final ConnectionMode connectionMode; private final AsyncCache containerControllerCache; - private final GlobalEndpointManager globalEndpointManager; - private final ConcurrentHashMap> groupMapByContainer; + private final ConcurrentHashMap> groupMapByContainer; private final RxPartitionKeyRangeCache partitionKeyRangeCache; + private final LinkedCancellationTokenSource cancellationTokenSource; + private final ConcurrentHashMap cancellationTokenMap; + public ThroughputControlStore( RxClientCollectionCache collectionCache, ConnectionMode connectionMode, - GlobalEndpointManager globalEndpointManager, RxPartitionKeyRangeCache partitionKeyRangeCache) { checkNotNull(collectionCache,"RxClientCollectionCache can not be null"); - checkNotNull(globalEndpointManager, "Global endpoint manager can not be null"); checkNotNull(partitionKeyRangeCache, "PartitionKeyRangeCache can not be null"); this.collectionCache = collectionCache; this.connectionMode = connectionMode; this.containerControllerCache = new AsyncCache<>(); - this.globalEndpointManager = globalEndpointManager; this.groupMapByContainer = new ConcurrentHashMap<>(); this.partitionKeyRangeCache = partitionKeyRangeCache; + + this.cancellationTokenSource = new LinkedCancellationTokenSource(); + this.cancellationTokenMap = new ConcurrentHashMap<>(); } - public void enableThroughputControlGroup(ThroughputControlGroup group) { + public void enableThroughputControlGroup(ThroughputControlGroupInternal group) { checkNotNull(group, "Throughput control group cannot be null"); - String collectionLink = Utils.trimBeginningAndEndingSlashes( - BridgeInternal.extractContainerSelfLink( - BridgeInternal.getTargetContainerFromThroughputControlGroup(group))); - this.groupMapByContainer.compute(collectionLink, (key, groupSet) -> { + String containerNameLink = Utils.trimBeginningAndEndingSlashes(BridgeInternal.extractContainerSelfLink(group.getTargetContainer())); + this.groupMapByContainer.compute(containerNameLink, (key, groupSet) -> { if (groupSet == null) { - groupSet = new HashSet<>(); + groupSet = ConcurrentHashMap.newKeySet(); } if (group.isDefault()) { - if (groupSet.stream().anyMatch(ThroughputControlGroup::isDefault)) { + if (groupSet.stream().anyMatch( + controlGroup -> group.isDefault() && !StringUtils.equals(group.getId(), controlGroup.getId()))) { throw new IllegalArgumentException("A default group already exists"); } } if (!groupSet.add(group)) { logger.warn("Can not add duplicate group"); + return groupSet; } if (groupSet.size() == 1) { // This is the first enabled group for the target container // Clean the current cache in case we have built EmptyThroughputContainerController. - this.containerControllerCache.remove(collectionLink); + this.containerControllerCache.remove(containerNameLink); } return groupSet; @@ -142,72 +141,93 @@ public Mono processRequest(RxDocumentServiceRequest request, Mono orig 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 more resource type here. + // 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 collectionLink = Utils.getCollectionName(request.getResourceAddress()); - return this.resolveContainerController(collectionLink) + String collectionNameLink = Utils.getCollectionName(request.getResourceAddress()); + return this.resolveContainerController(collectionNameLink) .flatMap(containerController -> { if (containerController.canHandleRequest(request)) { - return Mono.just(containerController); + 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.shouldRefreshContainerController(collectionLink, request) - .flatMap(shouldRefresh -> { - if (shouldRefresh) { - containerController.close().subscribeOn(Schedulers.parallel()).subscribe(); - this.containerControllerCache.refresh(collectionLink, () -> this.createAndInitContainerController(collectionLink)); - return this.resolveContainerController(collectionLink); + return this.updateControllerAndRetry(collectionNameLink, request, originalRequestMono); + }); + } + + 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(); } - // The container container controller is up to date, the request has staled info, will let the request pass - return Mono.just(containerController); + return null; }); - }) - .flatMap(containerController -> { - if (containerController.canHandleRequest(request)) { - return containerController.processRequest(request, originalRequestMono) - .doOnError(throwable -> this.handleException(request, containerController, 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; + + 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 collectionLink) { - checkArgument(StringUtils.isNotEmpty(collectionLink), "Collection link can not be null or empty"); + private Mono resolveContainerController(String containerNameLink) { + checkArgument(StringUtils.isNotEmpty(containerNameLink), "Container name link can not be null or empty"); return this.containerControllerCache.getAsync( - collectionLink, + containerNameLink, null, - () -> this.createAndInitContainerController(collectionLink) + () -> this.createAndInitContainerController(containerNameLink) ); } - private Mono createAndInitContainerController(String containerLink) { - checkArgument(StringUtils.isNotEmpty(containerLink), "Container link should not be null or empty"); + private Mono createAndInitContainerController(String containerNameLink) { + checkArgument(StringUtils.isNotEmpty(containerNameLink), "Container link should not be null or empty"); - if (this.groupMapByContainer.containsKey(containerLink)) { - return Mono.just(this.groupMapByContainer.get(containerLink)) + if (this.groupMapByContainer.containsKey(containerNameLink)) { + return Mono.just(this.groupMapByContainer.get(containerNameLink)) .flatMap(groups -> { + LinkedCancellationToken parentToken = + this.cancellationTokenMap.compute( + containerNameLink, + (key, cancellationToken) -> this.cancellationTokenSource.getToken()); + ThroughputContainerController containerController = new ThroughputContainerController( + this.collectionCache, this.connectionMode, - this.globalEndpointManager, groups, - this.partitionKeyRangeCache); + this.partitionKeyRangeCache, + parentToken); return containerController.init(); }); @@ -224,9 +244,9 @@ private Mono shouldRefreshContainerController(String containerLink, RxD Mono.just(StringUtils.equals(documentCollection.getResourceId(), request.requestContext.resolvedCollectionRid))); } - private void handleException(RxDocumentServiceRequest request, IThroughputController controller, Throwable throwable) { + 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(controller, "Container controller can not be null"); checkNotNull(throwable, "Exception can not be null"); CosmosException cosmosException = Utils.as(Exceptions.unwrap(throwable), CosmosException.class); @@ -234,7 +254,13 @@ private void handleException(RxDocumentServiceRequest request, IThroughputContro if (cosmosException != null && (isNameCacheStale(cosmosException) || isPartitionKeyMismatchException(cosmosException))) { - controller.close().subscribeOn(Schedulers.parallel()).subscribe(); + 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); @@ -244,4 +270,8 @@ private void handleException(RxDocumentServiceRequest request, IThroughputContro ); } } + + public void close() { + this.cancellationTokenSource.close(); + } } 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/ThroughputRequestThrottler.java index cb7504b2ef916..98bd6d575e12b 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/ThroughputRequestThrottler.java @@ -17,6 +17,7 @@ import reactor.core.publisher.Mono; import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.locks.ReentrantReadWriteLock; /** * This is the place where we tracking the RU usage, and make decision whether we should block the request. @@ -26,15 +27,28 @@ public class ThroughputRequestThrottler { private final AtomicReference availableThroughput; private final AtomicReference scheduledThroughput; + private final ReentrantReadWriteLock.WriteLock throughputWriteLock; + private final ReentrantReadWriteLock.ReadLock throughputReadLock; public ThroughputRequestThrottler(double scheduledThroughput) { this.availableThroughput = new AtomicReference<>(scheduledThroughput); this.scheduledThroughput = new AtomicReference<>(scheduledThroughput); + ReentrantReadWriteLock throughputReadWriteLock = new ReentrantReadWriteLock(); + this.throughputWriteLock = throughputReadWriteLock.writeLock(); + this.throughputReadLock = throughputReadWriteLock.readLock(); } - public void renewThroughputUsageCycle(double scheduledThroughput) { - this.scheduledThroughput.set(scheduledThroughput); - this.updateAvailableThroughput(); + public double renewThroughputUsageCycle(double scheduledThroughput) { + try { + this.throughputWriteLock.lock(); + double throughputUsagePercentage = (this.scheduledThroughput.get() - this.availableThroughput.get()) / this.scheduledThroughput.get(); + this.scheduledThroughput.set(scheduledThroughput); + this.updateAvailableThroughput(); + + return throughputUsagePercentage; + } finally { + this.throughputWriteLock.unlock(); + } } private void updateAvailableThroughput() { @@ -44,44 +58,58 @@ private void updateAvailableThroughput() { } public Mono processRequest(RxDocumentServiceRequest request, Mono originalRequestMono) { - if (this.availableThroughput.get() > 0) { - return originalRequestMono - .doOnSuccess(response -> this.trackRequestCharge(response)) - .doOnError(throwable -> this.trackRequestCharge(throwable)); - } else { - // there is no enough throughput left, block request - RequestRateTooLargeException requestRateTooLargeException = new RequestRateTooLargeException(); - - int backoffTimeInMilliSeconds = (int)Math.floor(Math.abs(this.availableThroughput.get() * 1000 / this.scheduledThroughput.get())); - requestRateTooLargeException.getResponseHeaders().put( - HttpConstants.HttpHeaders.RETRY_AFTER_IN_MILLISECONDS, - String.valueOf(backoffTimeInMilliSeconds)); - - requestRateTooLargeException.getResponseHeaders().put( - HttpConstants.HttpHeaders.SUB_STATUS, - String.valueOf(HttpConstants.SubStatusCodes.THROUGHPUT_CONTROL_REQUEST_RATE_TOO_LARGE)); - - if (request.requestContext != null) { - BridgeInternal.setResourceAddress(requestRateTooLargeException, request.requestContext.resourcePhysicalAddress); + try { + this.throughputReadLock.lock(); + if (this.availableThroughput.get() > 0) { + return originalRequestMono + .doOnSuccess(response -> this.trackRequestCharge(response)) + .doOnError(throwable -> this.trackRequestCharge(throwable)); + } else { + // there is no enough throughput left, block request + RequestRateTooLargeException requestRateTooLargeException = new RequestRateTooLargeException(); + + int backoffTimeInMilliSeconds = (int)Math.floor(Math.abs(this.availableThroughput.get() * 1000 / this.scheduledThroughput.get())); + + requestRateTooLargeException.getResponseHeaders().put( + HttpConstants.HttpHeaders.RETRY_AFTER_IN_MILLISECONDS, + String.valueOf(backoffTimeInMilliSeconds)); + + requestRateTooLargeException.getResponseHeaders().put( + HttpConstants.HttpHeaders.SUB_STATUS, + String.valueOf(HttpConstants.SubStatusCodes.THROUGHPUT_CONTROL_REQUEST_RATE_TOO_LARGE)); + + if (request.requestContext != null) { + BridgeInternal.setResourceAddress(requestRateTooLargeException, request.requestContext.resourcePhysicalAddress); + } + + return Mono.error(requestRateTooLargeException); } - - return Mono.error(requestRateTooLargeException); + } finally { + this.throughputReadLock.unlock(); } + } private void trackRequestCharge (T response) { - double requestCharge = 0; - if (response instanceof StoreResponse) { - requestCharge = ((StoreResponse)response).getRequestCharge(); - } else if (response instanceof RxDocumentServiceResponse) { - requestCharge = ((RxDocumentServiceResponse)response).getRequestCharge(); - } else if (response instanceof Throwable) { - CosmosException cosmosException = Utils.as(Exceptions.unwrap((Throwable) response), CosmosException.class); - if (cosmosException != null) { - requestCharge = cosmosException.getRequestCharge(); + try { + // Read lock is enough here. + this.throughputReadLock.lock(); + double requestCharge = 0; + if (response instanceof StoreResponse) { + requestCharge = ((StoreResponse)response).getRequestCharge(); + } else if (response instanceof RxDocumentServiceResponse) { + requestCharge = ((RxDocumentServiceResponse)response).getRequestCharge(); + } else if (response instanceof Throwable) { + CosmosException cosmosException = Utils.as(Exceptions.unwrap((Throwable) response), CosmosException.class); + if (cosmosException != null) { + requestCharge = cosmosException.getRequestCharge(); + } } + this.availableThroughput.getAndAccumulate(requestCharge, (available, consumed) -> available - consumed); + } finally { + this.throughputReadLock.unlock(); } - this.availableThroughput.getAndAccumulate(requestCharge, (available, consumed) -> available - consumed); + } public double getAvailableThroughput() { 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/config/GlobalThroughputControlGroup.java new file mode 100644 index 0000000000000..d7d19e16ba673 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/GlobalThroughputControlGroup.java @@ -0,0 +1,87 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl.config; + +import com.azure.cosmos.CosmosAsyncContainer; + +import java.time.Duration; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +public class GlobalThroughputControlGroup extends ThroughputControlGroupInternal { + private static final Duration DEFAULT_CONTROL_ITEM_RENEW_INTERVAL = Duration.ofSeconds(5); + + private final CosmosAsyncContainer globalControlContainer; + private final Duration controlItemRenewInterval; + private final Duration controlItemExpireInterval; + + public GlobalThroughputControlGroup( + String groupName, + CosmosAsyncContainer targetContainer, + Integer targetThroughput, + Double targetThroughputThreshold, + boolean isDefault, + CosmosAsyncContainer globalControlContainer, + Duration controlItemRenewInterval, + Duration controlItemExpireInterval) { + + super (groupName, targetContainer, targetThroughput, targetThroughputThreshold, isDefault); + + checkNotNull(globalControlContainer, "Global control container can not be null"); + + this.globalControlContainer = globalControlContainer; + this.controlItemRenewInterval = getDefaultControlItemRenewInterval(controlItemRenewInterval, controlItemRenewInterval); + this.controlItemExpireInterval = + controlItemExpireInterval != null ? controlItemExpireInterval : Duration.ofSeconds(2 * this.controlItemRenewInterval.getSeconds() + 1); + } + + private Duration getDefaultControlItemRenewInterval(Duration controlItemRenewInterval, Duration controlItemExpireInterval) { + if (controlItemRenewInterval != null) { + return controlItemRenewInterval; + } + + if (controlItemExpireInterval != null) { + return Duration.ofSeconds((controlItemExpireInterval.getSeconds() - 1) / 2); + } + + return DEFAULT_CONTROL_ITEM_RENEW_INTERVAL; + } + + /** + * Get the control container. + * This is the container to track all other clients throughput usage. + * + * @return The {@link CosmosAsyncContainer}. + */ + public CosmosAsyncContainer getGlobalControlContainer() { + return globalControlContainer; + } + + /** + * Get the control item renew interval. + * + * This controls how often the client is going to update the throughput usage of itself + * and adjust its own throughput share based on the throughput usage of other clients. + * + * In short words, it controls how quickly the shared throughput will reload balanced across different clients. + * + * @return The control item renew interval. + */ + public Duration getControlItemRenewInterval() { + return controlItemRenewInterval; + } + + + /** + * Get the control item expire interval. + * + * A client may be offline due to various reasons (being shutdown, network issue... ). + * This controls how quickly we will detect the client has been offline and hence allow its throughput share to be taken by other clients. + ** + * @return The control item renew interval. + */ + public Duration getControlItemExpireInterval() { + return controlItemExpireInterval; + } +} 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/config/LocalThroughputControlGroup.java new file mode 100644 index 0000000000000..5290c04ca59c8 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/LocalThroughputControlGroup.java @@ -0,0 +1,18 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl.config; + +import com.azure.cosmos.CosmosAsyncContainer; + +public class LocalThroughputControlGroup extends ThroughputControlGroupInternal { + + public LocalThroughputControlGroup( + String groupName, + CosmosAsyncContainer targetContainer, + Integer targetThroughput, + Double targetThroughputThreshold, + boolean isDefault) { + super (groupName, targetContainer, targetThroughput, targetThroughputThreshold, isDefault); + } +} 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/config/ThroughputControlGroupFactory.java new file mode 100644 index 0000000000000..fa3379487dfc3 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/ThroughputControlGroupFactory.java @@ -0,0 +1,47 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl.config; + +import com.azure.cosmos.BridgeInternal; +import com.azure.cosmos.CosmosAsyncContainer; +import com.azure.cosmos.ThroughputControlGroupConfig; +import com.azure.cosmos.GlobalThroughputControlConfig; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +public class ThroughputControlGroupFactory { + + public static LocalThroughputControlGroup createThroughputLocalControlGroup(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 LocalThroughputControlGroup( + groupConfig.getGroupName(), + targetContainer, + groupConfig.getTargetThroughput(), + groupConfig.getTargetThroughputThreshold(), + groupConfig.isDefault()); + } + + public static GlobalThroughputControlGroup createThroughputGlobalControlGroup( + ThroughputControlGroupConfig groupConfig, + GlobalThroughputControlConfig globalControlConfig, + CosmosAsyncContainer targetContainer) { + + checkNotNull(groupConfig, "Throughput control group config can not be null"); + checkNotNull(globalControlConfig, "Throughput global control config can not be null"); + checkNotNull(targetContainer, "Throughput target container can not be null"); + + return new GlobalThroughputControlGroup( + groupConfig.getGroupName(), + targetContainer, + groupConfig.getTargetThroughput(), + groupConfig.getTargetThroughputThreshold(), + groupConfig.isDefault(), + BridgeInternal.getControlContainerFromThroughputGlobalControlConfig(globalControlConfig), + globalControlConfig.getControlItemRenewInterval(), + globalControlConfig.getControlItemExpireInterval()); + + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroup.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/ThroughputControlGroupInternal.java similarity index 71% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroup.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/ThroughputControlGroupInternal.java index 5d89a37a96d14..f607188472da8 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroup.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/ThroughputControlGroupInternal.java @@ -1,21 +1,15 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos; +package com.azure.cosmos.implementation.throughputControl.config; +import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; -import com.azure.cosmos.implementation.throughputControl.ThroughputControlMode; -import com.azure.cosmos.util.Beta; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; -/** - * Group configuration which will be used in Throughput control. - */ -@Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) -public class ThroughputControlGroup { - private final ThroughputControlMode controlMode; +public abstract class ThroughputControlGroupInternal { private final String groupName; private final String id; private final boolean isDefault; @@ -23,12 +17,11 @@ public class ThroughputControlGroup { private final Integer targetThroughput; private final Double targetThroughputThreshold; - ThroughputControlGroup( + public ThroughputControlGroupInternal( String groupName, CosmosAsyncContainer targetContainer, Integer targetThroughput, Double targetThroughputThreshold, - ThroughputControlMode controlMode, boolean isDefault) { checkArgument(StringUtils.isNotEmpty(groupName), "Group name can not be null or empty"); @@ -42,10 +35,13 @@ public class ThroughputControlGroup { this.targetContainer = targetContainer; this.targetThroughput = targetThroughput; this.targetThroughputThreshold = targetThroughputThreshold; - this.controlMode = controlMode; this.isDefault = isDefault; - this.id = this.getId(); + this.id = String.format( + "%s/%s/%s", + this.targetContainer.getDatabase().getId(), + this.targetContainer.getId(), + this.groupName); } /** @@ -53,7 +49,6 @@ public class ThroughputControlGroup { * * @return the group name. */ - @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) public String getGroupName() { return this.groupName; } @@ -63,7 +58,7 @@ public String getGroupName() { * * @return the {@link CosmosAsyncContainer}. */ - CosmosAsyncContainer getTargetContainer() { + public CosmosAsyncContainer getTargetContainer() { return this.targetContainer; } @@ -75,7 +70,6 @@ CosmosAsyncContainer getTargetContainer() { * * @return the target throughput. */ - @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) public Integer getTargetThroughput() { return this.targetThroughput; } @@ -88,7 +82,6 @@ public Integer getTargetThroughput() { * * @return the target throughput threshold. */ - @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) public Double getTargetThroughputThreshold() { return this.targetThroughputThreshold; } @@ -97,20 +90,16 @@ public Double getTargetThroughputThreshold() { * Get whether this throughput control group will be used by default. * * By default, it is false. + * If it is true, requests without explicit override of the throughput control group will be routed to this group. * * @return {@code true} this throughput control group will be used by default unless being override. {@code false} otherwise. */ - @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) public boolean isDefault() { return this.isDefault; } - ThroughputControlMode getControlMode() { - return this.controlMode; - } - - private String getId() { - return this.targetContainer.getDatabase().getId() + "." + this.targetContainer.getId() + "." + this.groupName; + public String getId() { + return this.id; } @Override @@ -123,7 +112,7 @@ public boolean equals(Object other) { return false; } - ThroughputControlGroup that = (ThroughputControlGroup) other; + ThroughputControlGroupInternal that = (ThroughputControlGroupInternal) other; return StringUtils.equals(this.id, that.id); } 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/controller/IThroughputController.java index ad03cb7d68bfd..203479f8f9e47 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/controller/IThroughputController.java @@ -16,12 +16,6 @@ public interface IThroughputController { */ boolean canHandleRequest(RxDocumentServiceRequest request); - /** - * Close all the scheduled tasks and any other resources need to release. - * @return a representation of the deferred computation of this call. - */ - Mono close(); - /** * Initialize process. * Will create and initialize the lower level throughput controller and schedule tasks if needed. 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/controller/container/EmptyThroughputContainerController.java index ed807075db7b6..987d9f652bfe6 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/controller/container/EmptyThroughputContainerController.java @@ -16,11 +16,6 @@ public boolean canHandleRequest(RxDocumentServiceRequest request) { return true; } - @Override - public Mono close() { - return Mono.empty(); - } - @Override @SuppressWarnings("unchecked") public Mono init() { 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/controller/container/ThroughputContainerController.java index 4c87472fafe5c..388d97f82afbf 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/controller/container/ThroughputContainerController.java @@ -8,7 +8,6 @@ import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.CosmosBridgeInternal; import com.azure.cosmos.CosmosException; -import com.azure.cosmos.ThroughputControlGroup; import com.azure.cosmos.implementation.AsyncDocumentClient; import com.azure.cosmos.implementation.GlobalEndpointManager; import com.azure.cosmos.implementation.HttpConstants; @@ -16,10 +15,11 @@ 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.caches.RxCollectionCache; import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; -import com.azure.cosmos.implementation.changefeed.CancellationToken; -import com.azure.cosmos.implementation.changefeed.CancellationTokenSource; -import com.azure.cosmos.implementation.throughputControl.ThroughputResolveLevel; +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.models.CosmosQueryRequestOptions; @@ -35,8 +35,8 @@ import reactor.util.retry.RetrySpec; import java.time.Duration; -import java.util.Optional; import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicReference; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; @@ -53,66 +53,68 @@ public class ThroughputContainerController implements IThroughputContainerContro private static final Duration DEFAULT_THROUGHPUT_REFRESH_INTERVAL = Duration.ofMinutes(15); private static final int NO_OFFER_EXCEPTION_STATUS_CODE = HttpConstants.StatusCodes.BADREQUEST; - private static final int NO_OFFER_EXCEPTION_SUB_STATUS_CODE = HttpConstants.SubStatusCodes.UNKNOWN; + private static final int NO_OFFER_EXCEPTION_SUB_STATUS_CODE = HttpConstants.SubStatusCodes.OFFER_NOT_CONFIGURED; private final AsyncDocumentClient client; + private final RxCollectionCache collectionCache; private final ConnectionMode connectionMode; - private final GlobalEndpointManager globalEndpointManager; private final AsyncCache groupControllerCache; - private final Set groups; + private final Set groups; private final AtomicReference maxContainerThroughput; private final RxPartitionKeyRangeCache partitionKeyRangeCache; private final CosmosAsyncContainer targetContainer; - private final CancellationTokenSource cancellationTokenSource; + private final LinkedCancellationTokenSource cancellationTokenSource; + private final ConcurrentHashMap cancellationTokenMap; private ThroughputGroupControllerBase defaultGroupController; private String targetContainerRid; private String targetDatabaseRid; - private ThroughputResolveLevel throughputResolveLevel; + private ThroughputProvisioningScope throughputProvisioningScope; public ThroughputContainerController( + RxCollectionCache collectionCache, ConnectionMode connectionMode, - GlobalEndpointManager globalEndpointManager, - Set groups, - RxPartitionKeyRangeCache partitionKeyRangeCache) { + Set groups, + RxPartitionKeyRangeCache partitionKeyRangeCache, + LinkedCancellationToken parentToken) { - checkNotNull(globalEndpointManager, "GlobalEndpointManager can not be null"); + checkNotNull(collectionCache, "Collection cache can not be null"); checkArgument(groups != null && groups.size() > 0, "Throughput budget groups can not be null or empty"); checkNotNull(partitionKeyRangeCache, "RxPartitionKeyRangeCache can not be null"); + this.collectionCache = collectionCache; this.connectionMode = connectionMode; - this.globalEndpointManager = globalEndpointManager; this.groupControllerCache = new AsyncCache<>(); this.groups = groups; this.maxContainerThroughput = new AtomicReference<>(); this.partitionKeyRangeCache = partitionKeyRangeCache; - this.targetContainer = BridgeInternal.getTargetContainerFromThroughputControlGroup(groups.iterator().next()); + this.targetContainer = groups.iterator().next().getTargetContainer(); this.client = CosmosBridgeInternal.getContextClient(this.targetContainer); - this.throughputResolveLevel = this.getThroughputResolveLevel(groups); + this.throughputProvisioningScope = this.getThroughputResolveLevel(groups); - this.cancellationTokenSource = new CancellationTokenSource(); + this.cancellationTokenSource = new LinkedCancellationTokenSource(parentToken); + this.cancellationTokenMap = new ConcurrentHashMap<>(); } - private ThroughputResolveLevel 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 ThroughputResolveLevel.CONTAINER; + return ThroughputProvisioningScope.CONTAINER; } else { // There is no group configured with throughput threshold, so no need to query throughput - return ThroughputResolveLevel.NONE; + return ThroughputProvisioningScope.NONE; } } @Override @SuppressWarnings("unchecked") public Mono init() { - return this.resolveDatabaseResourceId() - .flatMap(controller -> this.resolveContainerResourceId()) - .flatMap(controller -> this.resolveContainerMaxThroughput()) + return this.resolveContainerResourceId() + .flatMap(containerRid -> this.resolveContainerMaxThroughput()) .flatMap(controller -> this.createAndInitializeGroupControllers()) .doOnSuccess(controller -> { Schedulers.parallel().schedule(() -> this.refreshContainerMaxThroughputTask(this.cancellationTokenSource.getToken()).subscribe()); @@ -120,39 +122,69 @@ public Mono init() { .thenReturn((T) this); } - private Mono resolveDatabaseResourceId() { + private Mono resolveDatabaseResourceId() { return this.targetContainer.getDatabase().read() .flatMap(response -> { this.targetDatabaseRid = response.getProperties().getResourceId(); - return Mono.just(this); + return Mono.just(this.targetDatabaseRid); }); } - private Mono resolveContainerResourceId() { + private Mono resolveContainerResourceId() { return this.targetContainer.read() .flatMap(response -> { this.targetContainerRid = response.getProperties().getResourceId(); - return Mono.just(this); + return Mono.just(this.targetContainerRid); }); } + private Mono resolveDatabaseThroughput() { + return Mono.justOrEmpty(this.targetDatabaseRid) + .switchIfEmpty(this.resolveDatabaseResourceId()) + .flatMap(databaseRid -> this.resolveThroughputByResourceId(databaseRid)); + } + + private Mono resolveContainerThroughput() { + if (StringUtils.isEmpty(this.targetContainerRid)) { + return this.resolveContainerResourceId() + .flatMap(containerRid -> this.resolveThroughputByResourceId(containerRid)) + .onErrorResume(throwable -> { + if (this.isOwnerResourceNotExistsException(throwable)) { + // During initialization time, the collection cache may contain staled info, + // refresh and retry one more time + this.collectionCache.refresh( + null, + BridgeInternal.getLink(this.targetContainer), + null + ); + } + + return Mono.error(throwable); + }) + .retryWhen(RetrySpec.max(1).filter(throwable -> this.isOwnerResourceNotExistsException(throwable))); + } else { + return Mono.just(this.targetContainerRid) + .flatMap(containerRid -> this.resolveThroughputByResourceId(containerRid)); + } + } + private Mono resolveContainerMaxThroughput() { - return Mono.just(this.throughputResolveLevel) // TODO: ---> test whether it works without defer - .flatMap(throughputResolveLevel -> { - if (throughputResolveLevel == ThroughputResolveLevel.CONTAINER) { - return this.resolveThroughputByResourceId(this.targetContainerRid) + return Mono.defer(() -> Mono.just(this.throughputProvisioningScope)) + .flatMap(throughputProvisioningScope -> { + if (throughputProvisioningScope == ThroughputProvisioningScope.CONTAINER) { + return this.resolveContainerThroughput() .onErrorResume(throwable -> { if (this.isOfferNotConfiguredException(throwable)) { - this.throughputResolveLevel = ThroughputResolveLevel.DATABASE; + this.throughputProvisioningScope = ThroughputProvisioningScope.DATABASE; } return Mono.error(throwable); }); - } else if (throughputResolveLevel == ThroughputResolveLevel.DATABASE) { - return this.resolveThroughputByResourceId(this.targetDatabaseRid) + } else if (throughputProvisioningScope == ThroughputProvisioningScope.DATABASE) { + return this.resolveDatabaseThroughput() .onErrorResume(throwable -> { if (this.isOfferNotConfiguredException(throwable)) { - this.throughputResolveLevel = ThroughputResolveLevel.CONTAINER; + this.throughputProvisioningScope = ThroughputProvisioningScope.CONTAINER; } return Mono.error(throwable); @@ -167,6 +199,13 @@ private Mono resolveContainerMaxThroughput() { this.updateMaxContainerThroughput(throughputResponse); return Mono.empty(); }) + .onErrorResume(throwable -> { + if (this.isOwnerResourceNotExistsException(throwable)) { + this.cancellationTokenSource.close(); + } + + return Mono.error(throwable); + }) .retryWhen( // Throughput can be configured on database level or container level // Retry at most 1 time so we can try on database and container both @@ -175,16 +214,21 @@ private Mono resolveContainerMaxThroughput() { } private Mono resolveThroughputByResourceId(String resourceId) { - // TODO: figure out how this work for serveless account - // TODO: work item: https://github.com/Azure/azure-sdk-for-java/issues/18776 + // Note: for serverless account, when we trying to query offers, + // we will get 400/0 with error message: Reading or replacing offers is not supported for serverless accounts. + // We are not supporting serverless account for throughput control for now. But the protocol may change in future, + // use https://github.com/Azure/azure-sdk-for-java/issues/18776 to keep track for possible future work. checkArgument(StringUtils.isNotEmpty(resourceId), "ResourceId can not be null or empty"); return this.client.queryOffers( BridgeInternal.getOfferQuerySpecFromResourceId(this.targetContainer, resourceId), new CosmosQueryRequestOptions()) .single() .flatMap(offerFeedResponse -> { if (offerFeedResponse.getResults().isEmpty()) { - return Mono.error( - BridgeInternal.createCosmosException(NO_OFFER_EXCEPTION_STATUS_CODE, "No offers found for the resource " + resourceId)); + CosmosException noOfferException = + BridgeInternal.createCosmosException(NO_OFFER_EXCEPTION_STATUS_CODE, "No offers found for the resource " + resourceId); + + BridgeInternal.setSubStatusCode(noOfferException, NO_OFFER_EXCEPTION_SUB_STATUS_CODE); + return Mono.error(noOfferException); } return this.client.readOffer(offerFeedResponse.getResults().get(0).getSelfLink()).single(); @@ -210,22 +254,21 @@ private boolean isOfferNotConfiguredException(Throwable throwable) { && cosmosException.getSubStatusCode() == NO_OFFER_EXCEPTION_SUB_STATUS_CODE; } + private boolean isOwnerResourceNotExistsException(Throwable throwable) { + checkNotNull(throwable, "Throwable should not be null"); + + CosmosException cosmosException = Utils.as(Exceptions.unwrap(throwable), CosmosException.class); + return cosmosException != null + && cosmosException.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND + && cosmosException.getSubStatusCode() == HttpConstants.SubStatusCodes.OWNER_RESOURCE_NOT_EXISTS; + } + @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 Mono.just(request) - .flatMap(request1 -> { - if (request1.getThroughputControlGroupName() == null) { - return Mono.just(new Utils.ValueHolder<>(this.defaultGroupController)); - } - else { - return this.getOrCreateThroughputGroupController(request.getThroughputControlGroupName()) - .defaultIfEmpty(this.defaultGroupController) - .map(Utils.ValueHolder::new); - } - }) + return this.getOrCreateThroughputGroupController(request.getThroughputControlGroupName()) .flatMap(groupController -> { if (groupController.v != null) { return groupController.v.processRequest(request, originalRequestMono); @@ -236,19 +279,22 @@ 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)) { - return Mono.empty(); + return Mono.just(new Utils.ValueHolder<>(this.defaultGroupController)); } - ThroughputControlGroup group = - this.groups.stream().filter(groupConfig -> StringUtils.equals(groupName, groupConfig.getGroupName())).findFirst().orElse(null); - if (group == null) { - return Mono.empty(); + for (ThroughputControlGroupInternal group : this.groups) { + if (StringUtils.equals(groupName, group.getGroupName())) { + return this.resolveThroughputGroupController(group) + .map(Utils.ValueHolder::new); + } } - return this.resolveThroughputGroupController(group); + // 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)); } public String getTargetContainerRid() { @@ -268,21 +314,26 @@ private Mono createAndInitializeGroupControllers( .then(Mono.just(this)); } - private Mono resolveThroughputGroupController(ThroughputControlGroup group) { + private Mono resolveThroughputGroupController(ThroughputControlGroupInternal group) { return this.groupControllerCache.getAsync( group.getGroupName(), null, () -> this.createAndInitializeGroupController(group)); } - private Mono createAndInitializeGroupController(ThroughputControlGroup group) { + private Mono createAndInitializeGroupController(ThroughputControlGroupInternal group) { + LinkedCancellationToken parentToken = + this.cancellationTokenMap.compute( + group.getGroupName(), + (key, cancellationToken) -> this.cancellationTokenSource.getToken()); + ThroughputGroupControllerBase groupController = ThroughputGroupControllerFactory.createController( this.connectionMode, - this.globalEndpointManager, group, this.maxContainerThroughput.get(), this.partitionKeyRangeCache, - this.targetContainerRid); + this.targetContainerRid, + parentToken); return groupController .init() @@ -295,38 +346,29 @@ private Mono createAndInitializeGroupController(T } - private Flux refreshContainerMaxThroughputTask(CancellationToken cancellationToken) { + private Flux refreshContainerMaxThroughputTask(LinkedCancellationToken cancellationToken) { checkNotNull(cancellationToken, "Cancellation token can not be null"); - if (this.throughputResolveLevel == ThroughputResolveLevel.NONE) { + if (this.throughputProvisioningScope == ThroughputProvisioningScope.NONE) { return Flux.empty(); } return Mono.delay(DEFAULT_THROUGHPUT_REFRESH_INTERVAL) - .flatMap(t -> this.resolveContainerMaxThroughput()) + .flatMap(t -> { + if (cancellationToken.isCancellationRequested()) { + return Mono.empty(); + } else { + return this.resolveContainerMaxThroughput(); + } + }) .flatMapIterable(controller -> this.groups) .flatMap(group -> this.resolveThroughputGroupController(group)) .doOnNext(groupController -> groupController.onContainerMaxThroughputRefresh(this.maxContainerThroughput.get())) .onErrorResume(throwable -> { - //TODO: Figure out how serverless work -// if (this.isOfferNotConfiguredException(throwable)) { -// // Throughput is not configured on container nor database, a good hint the resource does not exists any more -// this.close(); -// } - logger.warn("Refresh throughput failed with reason %s", throwable); return Mono.empty(); }) .then() .repeat(() -> !cancellationToken.isCancellationRequested()); } - - @Override - public Mono close() { - this.cancellationTokenSource.cancel(); - return Flux.fromIterable(this.groups) - .flatMap(group -> this.resolveThroughputGroupController(group)) - .flatMap(groupController -> groupController.close()) - .then(); - } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputResolveLevel.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/ThroughputProvisioningScope.java similarity index 53% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputResolveLevel.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/ThroughputProvisioningScope.java index 44d6c7d9650a4..a161cb8ed830a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputResolveLevel.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/ThroughputProvisioningScope.java @@ -1,9 +1,9 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl; +package com.azure.cosmos.implementation.throughputControl.controller.container; -public enum ThroughputResolveLevel { +public enum ThroughputProvisioningScope { NONE, CONTAINER, DATABASE 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/controller/group/ThroughputGroupControllerBase.java index 098aecddc50da..6f0cb59f6ff1b 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/controller/group/ThroughputGroupControllerBase.java @@ -5,15 +5,15 @@ import com.azure.cosmos.ConnectionMode; import com.azure.cosmos.CosmosException; -import com.azure.cosmos.ThroughputControlGroup; import com.azure.cosmos.implementation.GlobalEndpointManager; 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.caches.RxPartitionKeyRangeCache; -import com.azure.cosmos.implementation.changefeed.CancellationToken; -import com.azure.cosmos.implementation.changefeed.CancellationTokenSource; +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; @@ -23,7 +23,6 @@ import reactor.core.Exceptions; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; -import reactor.core.scheduler.Schedulers; import java.time.Duration; import java.util.concurrent.atomic.AtomicInteger; @@ -37,38 +36,35 @@ /** * Throughput group controller. Two common tasks across all group controller implementations: * 1. Create and initialize request controller based on connection mode - * 2. Schedule reset throughput uage every 1s. + * 2. Schedule reset throughput usage every 1s. */ public abstract class ThroughputGroupControllerBase implements IThroughputController { private final static Logger logger = LoggerFactory.getLogger(ThroughputGroupControllerBase.class); private final Duration DEFAULT_THROUGHPUT_USAGE_RESET_DURATION = Duration.ofSeconds(1); private final ConnectionMode connectionMode; - private final GlobalEndpointManager globalEndpointManager; - private final ThroughputControlGroup group; - private final AtomicReference groupThroughput; + private final ThroughputControlGroupInternal group; private final AtomicInteger maxContainerThroughput; private final RxPartitionKeyRangeCache partitionKeyRangeCache; private final AsyncCache requestControllerAsyncCache; private final String targetContainerRid; - private final CancellationTokenSource cancellationTokenSource; + protected final AtomicReference groupThroughput; + protected final LinkedCancellationTokenSource cancellationTokenSource; public ThroughputGroupControllerBase( ConnectionMode connectionMode, - GlobalEndpointManager globalEndpointManager, - ThroughputControlGroup group, + ThroughputControlGroupInternal group, Integer maxContainerThroughput, RxPartitionKeyRangeCache partitionKeyRangeCache, - String targetContainerRid) { + String targetContainerRid, + LinkedCancellationToken parentToken) { - checkNotNull(globalEndpointManager, "Global endpoint manager can not be null"); checkNotNull(group, "Throughput control group can not be null"); checkNotNull(partitionKeyRangeCache, "Partition key range cache can not be null or empty"); checkArgument(StringUtils.isNotEmpty(targetContainerRid), "Target container rid cannot be null nor empty"); this.connectionMode = connectionMode; - this.globalEndpointManager = globalEndpointManager; this.group = group; if (this.group.getTargetThroughputThreshold() != null) { @@ -85,10 +81,14 @@ public ThroughputGroupControllerBase( this.requestControllerAsyncCache = new AsyncCache<>(); this.targetContainerRid = targetContainerRid; - this.cancellationTokenSource = new CancellationTokenSource(); + this.cancellationTokenSource = new LinkedCancellationTokenSource(parentToken); } - private void calculateGroupThroughput() { + public abstract double getClientAllocatedThroughput(); + + public abstract void recordThroughputUsage(double loadFactor); + + protected void calculateGroupThroughput() { double allocatedThroughput = Double.MAX_VALUE; if (this.group.getTargetThroughputThreshold() != null) { allocatedThroughput = Math.min(allocatedThroughput, this.maxContainerThroughput.get() * this.group.getTargetThroughputThreshold()); @@ -101,23 +101,23 @@ private void calculateGroupThroughput() { this.groupThroughput.set(allocatedThroughput); } - @Override - @SuppressWarnings("unchecked") - public Mono init() { - return this.resolveRequestController() - .doOnSuccess(dummy -> { - this.throughputUsageCycleRenewTask(this.cancellationTokenSource.getToken()).subscribeOn(Schedulers.parallel()).subscribe(); - }) - .thenReturn((T)this); - } - - private Flux throughputUsageCycleRenewTask(CancellationToken cancellationToken) { + public Flux throughputUsageCycleRenewTask(LinkedCancellationToken cancellationToken) { checkNotNull(cancellationToken, "Cancellation token can not be null"); return Mono.delay(DEFAULT_THROUGHPUT_USAGE_RESET_DURATION) - .flatMap(t -> this.resolveRequestController()) - .doOnSuccess(requestController -> requestController.renewThroughputUsageCycle(this.groupThroughput.get())) + .flatMap(t -> { + if (cancellationToken.isCancellationRequested()) { + return Mono.empty(); + } else { + return this.resolveRequestController(); + } + }) + .doOnSuccess(requestController -> { + if (requestController != null) { + this.recordThroughputUsage(requestController.renewThroughputUsageCycle(this.getClientAllocatedThroughput())); + } + }) .onErrorResume(throwable -> { - logger.warn("Reset throughput usage failed with reason", throwable); + logger.warn("Reset throughput usage failed with reason ", throwable); return Mono.empty(); }) .then() @@ -126,15 +126,15 @@ private Flux throughputUsageCycleRenewTask(CancellationToken cancellationT private Mono createAndInitializeRequestController() { IThroughputRequestController requestController; + if (this.connectionMode == ConnectionMode.DIRECT) { requestController = new PkRangesThroughputRequestController( - this.globalEndpointManager, this.partitionKeyRangeCache, this.targetContainerRid, - this.groupThroughput.get()); + this.getClientAllocatedThroughput()); } else if (this.connectionMode == ConnectionMode.GATEWAY) { - requestController = new GlobalThroughputRequestController(this.globalEndpointManager, this.groupThroughput.get()); + requestController = new GlobalThroughputRequestController(this.getClientAllocatedThroughput()); } else { throw new IllegalArgumentException(String.format("Connection mode %s is not supported", this.connectionMode)); } @@ -153,64 +153,69 @@ public void onContainerMaxThroughputRefresh(int maxContainerThroughput) { } @Override - public Mono close() { - this.cancellationTokenSource.cancel(); - return this.resolveRequestController() - .flatMap(requestController -> requestController.close()); - } - - @Override - public Mono processRequest(RxDocumentServiceRequest request, Mono nextRequestMono) { + public Mono processRequest(RxDocumentServiceRequest request, Mono originalRequestMono) { return this.resolveRequestController() .flatMap(requestController -> { if (requestController.canHandleRequest(request)) { - return Mono.just(requestController); - } else { - // We can not find the matching pkRange, it is because either the group control is out of sync - // or the request has staled info. - // We will handle the first scenario by creating a new request controller, - // while for second scenario, we will go to the original request mono, which will eventually get exception from server - return this.shouldUpdateRequestController(request) - .flatMap(shouldUpdate -> { - if (shouldUpdate) { - requestController.close().subscribeOn(Schedulers.parallel()).subscribe(); - this.refreshRequestController(); - return this.resolveRequestController(); + return requestController.processRequest(request, originalRequestMono) + .doOnError(throwable -> this.handleException(throwable)); + } + + // We can not find the matching pkRange, it is because either the group control is out of sync + // or the request has staled info. + // We will handle the first scenario by creating a new request controller, + // while for second scenario, we will go to the original request mono, which will eventually get exception from server + return this.updateControllerAndRetry(request, originalRequestMono); + }); + } + + private Mono updateControllerAndRetry(RxDocumentServiceRequest request, Mono nextRequestMono) { + + return this.shouldUpdateRequestController(request) + .flatMap(shouldUpdate -> { + if (shouldUpdate) { + this.refreshRequestController(); + return this.resolveRequestController() + .flatMap(updatedController -> { + if (updatedController.canHandleRequest(request)) { + return updatedController.processRequest(request, nextRequestMono) + .doOnError(throwable -> this.handleException(throwable)); } else { - return Mono.just(requestController); + // If we reach here and still can not handle the request, it should mean the request has staled info + // and the request will fail by server + logger.warn( + "Can not find request controller to handle request {} with pkRangeId {}", + request.getActivityId(), + this.getResolvedPartitionKeyRangeId(request)); + return nextRequestMono; } }); - } - }) - .flatMap(requestController -> { - if (requestController.canHandleRequest(request)) { - return requestController.processRequest(request, nextRequestMono) - .doOnError(throwable -> this.handleException(throwable)); } else { - // If we reach here and still can not handle the request, it should mean the request has staled info - // and the request will fail by server - logger.warn( - "Can not find request controller to handle request {} with pkRangeId {}", - request.getActivityId(), - request.requestContext.resolvedPartitionKeyRange.getId()); return nextRequestMono; } }); } + private String getResolvedPartitionKeyRangeId(RxDocumentServiceRequest request) { + if (request.requestContext != null && request.requestContext.resolvedPartitionKeyRange != null) { + return request.requestContext.resolvedPartitionKeyRange.getId(); + } + + return StringUtils.EMPTY; + } + private Mono shouldUpdateRequestController(RxDocumentServiceRequest request) { return this.partitionKeyRangeCache.tryGetRangeByPartitionKeyRangeId( null, request.requestContext.resolvedCollectionRid, request.requestContext.resolvedPartitionKeyRange.getId(), null) - .map(pkRangeHolder -> pkRangeHolder.v) - .flatMap(pkRange -> { - if (pkRange == null) { + .flatMap(pkRangeHolder -> { + if (pkRangeHolder.v == null) { return Mono.just(Boolean.FALSE); } else { return Mono.just(Boolean.TRUE); }}); } - private Mono resolveRequestController() { + protected Mono resolveRequestController() { return this.requestControllerAsyncCache.getAsync( this.group.getGroupName(), null, 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/controller/group/ThroughputGroupControllerFactory.java index 43d11f162c3b0..01d1c9dbb040c 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/controller/group/ThroughputGroupControllerFactory.java @@ -3,36 +3,44 @@ package com.azure.cosmos.implementation.throughputControl.controller.group; -import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.ConnectionMode; -import com.azure.cosmos.ThroughputControlGroup; import com.azure.cosmos.implementation.GlobalEndpointManager; import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; -import com.azure.cosmos.implementation.throughputControl.ThroughputControlMode; +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; public class ThroughputGroupControllerFactory { public static ThroughputGroupControllerBase createController( ConnectionMode connectionMode, - GlobalEndpointManager globalEndpointManager, - ThroughputControlGroup group, + ThroughputControlGroupInternal group, Integer maxContainerThroughput, RxPartitionKeyRangeCache partitionKeyRangeCache, - String targetCollectionRid) { + String targetCollectionRid, + LinkedCancellationToken parentToken) { - ThroughputControlMode controlMode = BridgeInternal.getThroughputControlMode(group); - if (controlMode == ThroughputControlMode.LOCAL) { - return new ThroughputGroupLocalController( + if (group instanceof LocalThroughputControlGroup) { + return new LocalThroughputControlGroupController( connectionMode, - globalEndpointManager, - group, + (LocalThroughputControlGroup) group, maxContainerThroughput, partitionKeyRangeCache, - targetCollectionRid); + targetCollectionRid, + parentToken); + } else if (group instanceof GlobalThroughputControlGroup) { + return new GlobalThroughputControlGroupController( + connectionMode, + (GlobalThroughputControlGroup) group, + maxContainerThroughput, + partitionKeyRangeCache, + targetCollectionRid, + parentToken); } - // TODO: distributed mode support - - throw new IllegalArgumentException(String.format("Throughput group control mode %s is not supported", controlMode)); + throw new IllegalArgumentException(String.format("Throughput group control group %s is not supported", group.getClass())); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/ThroughputGroupLocalController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/ThroughputGroupLocalController.java deleted file mode 100644 index 120869d9d3430..0000000000000 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/ThroughputGroupLocalController.java +++ /dev/null @@ -1,23 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -package com.azure.cosmos.implementation.throughputControl.controller.group; - -import com.azure.cosmos.ConnectionMode; -import com.azure.cosmos.ThroughputControlGroup; -import com.azure.cosmos.implementation.GlobalEndpointManager; -import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; - -public class ThroughputGroupLocalController extends ThroughputGroupControllerBase { - - public ThroughputGroupLocalController( - ConnectionMode connectionMode, - GlobalEndpointManager globalEndpointManager, - ThroughputControlGroup group, - Integer maxContainerThroughput, - RxPartitionKeyRangeCache partitionKeyRangeCache, - String targetContainerRid) { - - super(connectionMode, globalEndpointManager, group, maxContainerThroughput, partitionKeyRangeCache, targetContainerRid); - } -} 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/controller/group/global/GlobalThroughputControlClientItem.java new file mode 100644 index 0000000000000..a20578e3f4e98 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlClientItem.java @@ -0,0 +1,62 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl.controller.group.global; + +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.time.Duration; +import java.time.Instant; + +public class GlobalThroughputControlClientItem extends GlobalThroughputControlItem { + + @JsonProperty(value = "initializeTime", required = true) + private String initializeTime; + + @JsonProperty(value = "loadFactor", required = true) + private double loadFactor; + + @JsonProperty(value = "allocatedThroughput", required = true) + private double allocatedThroughput; + + /** + * Constructor used for Json deserialization + */ + public GlobalThroughputControlClientItem() { + + } + + public GlobalThroughputControlClientItem( + String id, + String partitionKeyValue, + double loadFactor, + double allocatedThroughput, + Duration clientItemExpireInterval) { + super(id, partitionKeyValue); + + this.loadFactor = loadFactor; + this.allocatedThroughput = allocatedThroughput; + this.initializeTime = Instant.now().toString(); + this.setTtl((int)clientItemExpireInterval.getSeconds()); + } + + public String getInitializeTime() { + return initializeTime; + } + + public double getLoadFactor() { + return loadFactor; + } + + public void setLoadFactor(double loadFactor) { + this.loadFactor = loadFactor; + } + + public double getAllocatedThroughput() { + return allocatedThroughput; + } + + public void setAllocatedThroughput(double allocatedThroughput) { + this.allocatedThroughput = allocatedThroughput; + } +} 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/controller/group/global/GlobalThroughputControlConfigItem.java new file mode 100644 index 0000000000000..b1578ea9d4353 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlConfigItem.java @@ -0,0 +1,107 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl.controller.group.global; + +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.util.Objects; + +public class GlobalThroughputControlConfigItem extends GlobalThroughputControlItem { + + @JsonProperty(value = "targetThroughput") + @JsonInclude(JsonInclude.Include.NON_NULL) + private String targetThroughput; + + @JsonProperty(value = "targetThroughputThreshold") + @JsonInclude(JsonInclude.Include.NON_NULL) + private String targetThroughputThreshold; + + @JsonProperty(value = "isDefault", required = true) + private boolean isDefault; + + /** + * Constructor used for Json deserialization + */ + public GlobalThroughputControlConfigItem() { + + } + + public GlobalThroughputControlConfigItem( + String id, + String partitionKeyValue, + Integer targetThroughput, + Double targetThroughputThreshold, + boolean isDefault) { + + super(id, partitionKeyValue); + this.targetThroughput = targetThroughput != null ? targetThroughput.toString() : StringUtils.EMPTY; + this.targetThroughputThreshold = targetThroughputThreshold != null ? targetThroughputThreshold.toString() : StringUtils.EMPTY; + this.isDefault = isDefault; + } + + public String getTargetThroughput() { + return targetThroughput; + } + + public void setTargetThroughput(String targetThroughput) { + this.targetThroughput = targetThroughput; + } + + public String getTargetThroughputThreshold() { + return targetThroughputThreshold; + } + + public void setTargetThroughputThreshold(String targetThroughputThreshold) { + this.targetThroughputThreshold = targetThroughputThreshold; + } + + @JsonIgnore + public boolean isDefault() { + return isDefault; + } + + public void setDefault(boolean aDefault) { + isDefault = aDefault; + } + + @Override + public boolean equals(Object other) { + if (this == other) { + return true; + } + + if (other == null || getClass() != other.getClass()) { + return false; + } + + GlobalThroughputControlConfigItem that = (GlobalThroughputControlConfigItem) other; + + return StringUtils.equals(this.getId(), that.getId()) + && StringUtils.equals(this.getGroupId(), that.getGroupId()) + && StringUtils.equals(this.targetThroughput, that.targetThroughput) + && StringUtils.equals(this.targetThroughputThreshold, that.targetThroughputThreshold) + && this.isDefault == that.isDefault; + } + + @Override + public int hashCode() { + int result = Objects.hash(this.getId(), this.getGroupId(), targetThroughput, targetThroughputThreshold); + result = 31 * result + Boolean.hashCode(this.isDefault); + return result; + } + + @Override + public String toString() { + return "ThroughputGlobalControlConfigItem{" + + "id='" + this.getId() + '\'' + + "groupId='" + this.getGroupId() + '\'' + + "targetThroughput='" + this.targetThroughput + '\'' + + ", targetThroughputThreshold='" + this.targetThroughputThreshold + '\'' + + ", isDefault=" + this.isDefault + + '}'; + } +} 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/controller/group/global/GlobalThroughputControlGroupController.java new file mode 100644 index 0000000000000..47a2f26502c17 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlGroupController.java @@ -0,0 +1,126 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl.controller.group.global; + +import com.azure.cosmos.ConnectionMode; +import com.azure.cosmos.implementation.GlobalEndpointManager; +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 org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; + +import java.time.Duration; +import java.time.Instant; +import java.util.concurrent.atomic.AtomicReference; + +public class GlobalThroughputControlGroupController extends ThroughputGroupControllerBase { + 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; + private static final int DEFAULT_THROUGHPUT_USAGE_QUEUE_SIZE = 300; // 5 mins windows since we refresh ru usage every 1s + private static final double MIN_LOAD_FACTOR = 0.1; + + private final Duration controlItemRenewInterval; + private final ThroughputControlContainerManager containerManager; + private final EvictingQueue throughputUsageSnapshotQueue; + private final Object throughputUsageSnapshotQueueLock; + private AtomicReference clientThroughputShare; + + public GlobalThroughputControlGroupController( + ConnectionMode connectionMode, + GlobalThroughputControlGroup group, + Integer maxContainerThroughput, + RxPartitionKeyRangeCache partitionKeyRangeCache, + String targetContainerRid, + LinkedCancellationToken parentToken) { + super(connectionMode, group, maxContainerThroughput, partitionKeyRangeCache, targetContainerRid, parentToken); + + this.controlItemRenewInterval = group.getControlItemRenewInterval(); + this.containerManager = new ThroughputControlContainerManager(group); + + this.throughputUsageSnapshotQueue = EvictingQueue.create(DEFAULT_THROUGHPUT_USAGE_QUEUE_SIZE); + this.throughputUsageSnapshotQueue.add(new ThroughputUsageSnapshot(INITIAL_THROUGHPUT_USAGE)); + this.throughputUsageSnapshotQueueLock = new Object(); + this.clientThroughputShare = new AtomicReference<>(INITIAL_CLIENT_THROUGHPUT_RU_SHARE); + } + + @Override + @SuppressWarnings("unchecked") + public Mono init() { + return this.containerManager.validateControlContainer() + .flatMap(dummy -> this.containerManager.getOrCreateConfigItem()) + .flatMap(dummy -> { + double loadFactor = this.calculateLoadFactor(); + return this.calculateClientThroughputShare(loadFactor) + .flatMap(controller -> this.containerManager.createGroupClientItem(loadFactor, this.getClientAllocatedThroughput())); + }) + .flatMap(dummy -> this.resolveRequestController()) + .doOnSuccess(dummy -> { + this.throughputUsageCycleRenewTask(this.cancellationTokenSource.getToken()).publishOn(Schedulers.parallel()).subscribe(); + this.calculateClientThroughputShareTask(this.cancellationTokenSource.getToken()).publishOn(Schedulers.parallel()).subscribe(); + }) + .thenReturn((T)this); + } + + @Override + public double getClientAllocatedThroughput() { + return this.groupThroughput.get() * this.clientThroughputShare.get(); + } + + @Override + public void recordThroughputUsage(double throughputUsage) { + synchronized (this.throughputUsageSnapshotQueueLock) { + this.throughputUsageSnapshotQueue.add(new ThroughputUsageSnapshot(throughputUsage)); + } + } + + private Mono calculateClientThroughputShare(double loadFactor) { + return this.containerManager.queryLoadFactorsOfAllClients(loadFactor) + .doOnSuccess(totalLoads -> this.clientThroughputShare.set(loadFactor / totalLoads)) + .thenReturn(this); + } + + private double calculateLoadFactor() { + synchronized (this.throughputUsageSnapshotQueueLock) { + Instant startTime = this.throughputUsageSnapshotQueue.peek().getTime(); + + double totalWeight = 0.0; + for (ThroughputUsageSnapshot throughputUsageSnapshot : this.throughputUsageSnapshotQueue) { + totalWeight += throughputUsageSnapshot.calculateWeight(startTime); + } + + double loadFactor = 0.0; + for (ThroughputUsageSnapshot throughputUsageSnapshot : this.throughputUsageSnapshotQueue) { + loadFactor += (throughputUsageSnapshot.getWeight() / totalWeight) * throughputUsageSnapshot.getThroughputUsage(); + } + + return Math.max(MIN_LOAD_FACTOR, loadFactor); + } + } + + private Flux calculateClientThroughputShareTask(LinkedCancellationToken cancellationToken) { + return Mono.delay(controlItemRenewInterval) + .flatMap(t -> { + if (cancellationToken.isCancellationRequested()) { + return Mono.empty(); + } else { + double loadFactor = this.calculateLoadFactor(); + return this.calculateClientThroughputShare(loadFactor) + .flatMap(dummy -> this.containerManager.replaceOrCreateGroupClientItem(loadFactor, this.getClientAllocatedThroughput())); + } + }) + .onErrorResume(throwable -> { + logger.warn("Calculate throughput task failed ", throwable); + return Mono.empty(); + }) + .then() + .repeat(() -> !cancellationToken.isCancellationRequested()); + } +} 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/controller/group/global/GlobalThroughputControlItem.java new file mode 100644 index 0000000000000..bfaeb5e2ba04c --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlItem.java @@ -0,0 +1,64 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl.controller.group.global; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; + +public abstract class GlobalThroughputControlItem { + @JsonProperty(value = "id", required = true) + private String id; + + @JsonProperty(value = "groupId", required = true) + private String groupId; + + @JsonProperty(value = "_etag") + @JsonInclude(JsonInclude.Include.NON_NULL) + private String etag; + + @JsonProperty(value = "ttl") + @JsonInclude(JsonInclude.Include.NON_NULL) + private Integer ttl; + + public GlobalThroughputControlItem() { + + } + + public GlobalThroughputControlItem(String id, String partitionKeyValue) { + this.id = id; + this.groupId = partitionKeyValue; + } + + public String getId() { + return id; + } + + public void setId(String id) { + this.id = id; + } + + public String getGroupId() { + return groupId; + } + + public void setGroupId(String groupId) { + this.groupId = groupId; + } + + public String getEtag() { + return etag; + } + + public void setEtag(String etag) { + this.etag = etag; + } + + public Integer getTtl() { + return ttl; + } + + public void setTtl(Integer ttl) { + this.ttl = ttl; + } +} 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/controller/group/global/ThroughputControlContainerManager.java new file mode 100644 index 0000000000000..e226abaa97b5a --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/ThroughputControlContainerManager.java @@ -0,0 +1,213 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl.controller.group.global; + +import com.azure.cosmos.CosmosAsyncContainer; +import com.azure.cosmos.CosmosException; +import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.implementation.Utils; +import com.azure.cosmos.implementation.throughputControl.config.GlobalThroughputControlGroup; +import com.azure.cosmos.models.CosmosItemRequestOptions; +import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.models.SqlParameter; +import com.azure.cosmos.models.SqlQuerySpec; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import reactor.core.Exceptions; +import reactor.core.publisher.Mono; +import reactor.util.retry.RetrySpec; + +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; + +/** + * In Throughput global control mode, in order to coordinate with other clients and share the throughput defined in the group, + * there is a need to create/read/replace related items in the control container. This class contains all those related operations. + */ +public class ThroughputControlContainerManager { + private static final Logger logger = LoggerFactory.getLogger(ThroughputControlContainerManager.class); + + private static final String CLIENT_ITEM_PARTITION_KEY_VALUE_SUFFIX = ".client"; + private static final String CONFIG_ITEM_ID_SUFFIX = ".info"; + private static final String CONFIG_ITEM_PARTITION_KEY_VALUE_SUFFIX = ".config"; + private static final String PARTITION_KEY_PATH = "/groupId"; + + private final String clientItemId; + private final String clientItemPartitionKeyValue; + private final String configItemId; + private final String configItemPartitionKeyValue; + private final CosmosAsyncContainer globalControlContainer; + private final GlobalThroughputControlGroup group; + + private GlobalThroughputControlConfigItem configItem; + private GlobalThroughputControlClientItem clientItem; + + public ThroughputControlContainerManager(GlobalThroughputControlGroup group) { + checkNotNull(group, "Global control group config can not be null"); + + this.globalControlContainer = group.getGlobalControlContainer(); + this.group = group; + + String encodedGroupId = Utils.encodeUrlBase64String(this.group.getId().getBytes(StandardCharsets.UTF_8)); + + this.clientItemId = encodedGroupId + UUID.randomUUID(); + this.clientItemPartitionKeyValue = this.group.getId() + CLIENT_ITEM_PARTITION_KEY_VALUE_SUFFIX; + this.configItemId = encodedGroupId + CONFIG_ITEM_ID_SUFFIX; + this.configItemPartitionKeyValue = this.group.getId() + CONFIG_ITEM_PARTITION_KEY_VALUE_SUFFIX; + } + + public Mono createGroupClientItem(double loadFactor, double allocatedThroughput) { + CosmosItemRequestOptions requestOptions = new CosmosItemRequestOptions(); + requestOptions.setContentResponseOnWriteEnabled(true); + + return Mono.just( + new GlobalThroughputControlClientItem( + this.clientItemId, + this.clientItemPartitionKeyValue, + loadFactor, + allocatedThroughput, + this.group.getControlItemExpireInterval())) + .flatMap(groupClientItem -> this.globalControlContainer.createItem(groupClientItem, requestOptions)) + .flatMap(itemResponse -> { + this.clientItem = itemResponse.getItem(); + return Mono.just(this.clientItem); + }); + } + + /** + * Get or create the throughput global control config item. + * This is to make sure all the clients are using the same configuration for the group. + * + * The config item in the control container will be used as the source of truth. + * If the client has a different config, it will be overwritten by the one in the control container. + * + * @return A {@link GlobalThroughputControlClientItem}. + */ + public Mono getOrCreateConfigItem() { + CosmosItemRequestOptions requestOptions = new CosmosItemRequestOptions(); + requestOptions.setContentResponseOnWriteEnabled(true); + + GlobalThroughputControlConfigItem expectedConfigItem = + new GlobalThroughputControlConfigItem( + this.configItemId, + this.configItemPartitionKeyValue, + this.group.getTargetThroughput(), + this.group.getTargetThroughputThreshold(), + this.group.isDefault()); + + return this.globalControlContainer.readItem( + this.configItemId, + new PartitionKey(this.configItemPartitionKeyValue), + GlobalThroughputControlConfigItem.class) + .onErrorResume(throwable -> { + CosmosException cosmosException = Utils.as(Exceptions.unwrap(throwable), CosmosException.class); + if (cosmosException != null && cosmosException.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND) { + // hooray, you are the first one, needs to create the config file now + return this.globalControlContainer.createItem(expectedConfigItem, requestOptions); + } + + return Mono.error(throwable); + }) + .retryWhen(RetrySpec.max(10).filter(throwable -> { + CosmosException cosmosException = Utils.as(Exceptions.unwrap(throwable), CosmosException.class); + return cosmosException != null && cosmosException.getStatusCode() == HttpConstants.StatusCodes.CONFLICT; + })) + .flatMap(itemResponse -> { + this.configItem = itemResponse.getItem(); + + if (!expectedConfigItem.equals(configItem)) { + logger.warn( + "Group config using by this client is different than the one in control container, will be ignored. Using following config: {}", + this.configItem.toString()); + } + + return Mono.just(this.configItem); + }); + } + + /** + * Query the load factor of all other clients except itself for the group, and then add the client load factor to the final sum. + + * @param clientLoadFactor The load factor for the current client. + * @return The sum of load factor from all clients. + */ + public Mono queryLoadFactorsOfAllClients(double clientLoadFactor) { + // The current design is using ttl to expire client items, so there is no need to check whether the client item is expired. + + String sqlQueryTest = "SELECT VALUE SUM(c.loadFactor) FROM c WHERE c.groupId = @GROUPID AND c.id != @CLIENTITEMID"; + List parameters = new ArrayList<>(); + parameters.add(new SqlParameter("@GROUPID", this.clientItemPartitionKeyValue)); + parameters.add(new SqlParameter("@CLIENTITEMID", this.clientItemId)); + + SqlQuerySpec querySpec = new SqlQuerySpec(sqlQueryTest, parameters); + return this.globalControlContainer.queryItems(querySpec, Double.class) + .single() + .map(result -> result + clientLoadFactor); + } + + /** + * Update the existing group client item. + * + * If resource not found, then create a new client item. + * The client item may get deleted based on the ttl if the client can not keep updating the item due to unexpected failure (for example, network failure). + * + * @param loadFactor The new load factor of the client. + * @param clientAllocatedThroughput The new allocated throughput for the client. + * @return A {@link GlobalThroughputControlClientItem}; + */ + public Mono replaceOrCreateGroupClientItem(double loadFactor, double clientAllocatedThroughput) { + CosmosItemRequestOptions itemRequestOptions = new CosmosItemRequestOptions(); + itemRequestOptions.setContentResponseOnWriteEnabled(true); + + return Mono.just(this.clientItem) + .flatMap(groupClientItem -> { + groupClientItem.setLoadFactor(loadFactor); + groupClientItem.setAllocatedThroughput(clientAllocatedThroughput); + return this.globalControlContainer.replaceItem( + groupClientItem, groupClientItem.getId(), new PartitionKey(groupClientItem.getGroupId()), itemRequestOptions); + }) + .onErrorResume(throwable -> { + CosmosException cosmosException = Utils.as(Exceptions.unwrap(throwable), CosmosException.class); + if (cosmosException != null && cosmosException.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND) { + logger.warn("Can not find the expected client item {}, will recreate a new one", this.clientItem.getId()); + return this.globalControlContainer.createItem(this.clientItem, itemRequestOptions) + .retryWhen(RetrySpec.max(5)); + } + + return Mono.error(throwable); + }) + .flatMap(itemResponse -> { + this.clientItem = itemResponse.getItem(); + return Mono.just(this.clientItem); + }); + } + + /** + * Make sure the control container provided is partitioned as expected. + * + * @return A {@link ThroughputControlContainerManager}. + */ + public Mono validateControlContainer() { + return this.globalControlContainer.read() + .map(containerResponse -> containerResponse.getProperties()) + .flatMap(containerProperties -> { + boolean isPartitioned = + containerProperties.getPartitionKeyDefinition() != null && + containerProperties.getPartitionKeyDefinition().getPaths() != null && + containerProperties.getPartitionKeyDefinition().getPaths().size() > 0; + if (!isPartitioned + || (containerProperties.getPartitionKeyDefinition().getPaths().size() != 1 + || !containerProperties.getPartitionKeyDefinition().getPaths().get(0).equals(PARTITION_KEY_PATH))) { + return Mono.error(new IllegalArgumentException("The control container must have partition key equal to " + PARTITION_KEY_PATH)); + } + + return Mono.empty(); + }) + .thenReturn(this); + } +} 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/controller/group/global/ThroughputUsageSnapshot.java new file mode 100644 index 0000000000000..e26c038cd4e91 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/ThroughputUsageSnapshot.java @@ -0,0 +1,42 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl.controller.group.global; + +import java.time.Duration; +import java.time.Instant; + +public class ThroughputUsageSnapshot { + private final double throughputUsage; + private final Instant time; + private double weight; + + public ThroughputUsageSnapshot(double throughputUsage) { + this.throughputUsage = throughputUsage; + this.time = Instant.now(); + } + + public double getThroughputUsage() { + return throughputUsage; + } + + public Instant getTime() { + return time; + } + + public double getWeight() { + return this.weight; + } + + /** + * The most recent should have higher weight, + * so it makes higher impact on the final load factor of the client. + * + * @param startTime The start time. + * @return The weight of the throughput usage snapshot. + */ + public double calculateWeight(Instant startTime) { + this.weight = Math.exp(Duration.between(time, startTime).getSeconds()); + return this.weight; + } +} 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/controller/group/local/LocalThroughputControlGroupController.java new file mode 100644 index 0000000000000..1806770560121 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/local/LocalThroughputControlGroupController.java @@ -0,0 +1,47 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl.controller.group.local; + +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.LocalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.controller.group.ThroughputGroupControllerBase; +import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; + +public class LocalThroughputControlGroupController extends ThroughputGroupControllerBase { + + public LocalThroughputControlGroupController( + ConnectionMode connectionMode, + LocalThroughputControlGroup group, + Integer maxContainerThroughput, + RxPartitionKeyRangeCache partitionKeyRangeCache, + String targetContainerRid, + LinkedCancellationToken parentToken) { + + super(connectionMode, group, maxContainerThroughput, partitionKeyRangeCache, targetContainerRid, parentToken); + } + + @Override + @SuppressWarnings("unchecked") + public Mono init() { + return this.resolveRequestController() + .doOnSuccess(dummy -> { + this.throughputUsageCycleRenewTask(this.cancellationTokenSource.getToken()).publishOn(Schedulers.parallel()).subscribe(); + }) + .thenReturn((T)this); + } + + @Override + public double getClientAllocatedThroughput() { + return this.groupThroughput.get(); + } + + @Override + public void recordThroughputUsage(double loadFactor) { + return; + } +} 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/controller/request/GlobalThroughputRequestController.java index 0e8b68c21ed7b..c57f263fffdee 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/controller/request/GlobalThroughputRequestController.java @@ -3,40 +3,25 @@ package com.azure.cosmos.implementation.throughputControl.controller.request; -import com.azure.cosmos.implementation.GlobalEndpointManager; import com.azure.cosmos.implementation.RxDocumentServiceRequest; import com.azure.cosmos.implementation.throughputControl.ThroughputRequestThrottler; -import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; -import java.net.URI; -import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicReference; -import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; - public class GlobalThroughputRequestController implements IThroughputRequestController { - private final GlobalEndpointManager globalEndpointManager; private final AtomicReference scheduledThroughput; - private final ConcurrentHashMap requestThrottlerMapByRegion; - - public GlobalThroughputRequestController(GlobalEndpointManager globalEndpointManager, double initialScheduledThroughput) { - checkNotNull(globalEndpointManager, "Global endpoint manager can not be null"); + private final ThroughputRequestThrottler requestThrottler; - this.globalEndpointManager = globalEndpointManager; + public GlobalThroughputRequestController(double initialScheduledThroughput) { this.scheduledThroughput = new AtomicReference<>(initialScheduledThroughput); - this.requestThrottlerMapByRegion = new ConcurrentHashMap<>(); + this.requestThrottler = new ThroughputRequestThrottler(this.scheduledThroughput.get()); } @Override @SuppressWarnings("unchecked") public Mono init() { - return Flux.fromIterable(this.globalEndpointManager.getReadEndpoints()) - .flatMap(endpoint -> { - requestThrottlerMapByRegion.computeIfAbsent(endpoint, key -> new ThroughputRequestThrottler(this.scheduledThroughput.get())); - return Mono.empty(); - }) - .then(Mono.just((T)this)); + return Mono.just((T)requestThrottler); } @Override @@ -45,26 +30,13 @@ public boolean canHandleRequest(RxDocumentServiceRequest request) { } @Override - public Mono processRequest(RxDocumentServiceRequest request, Mono nextRequestMono) { - return Mono.defer( - () -> Mono.just( - this.requestThrottlerMapByRegion.computeIfAbsent( - this.globalEndpointManager.resolveServiceEndpoint(request), - key -> new ThroughputRequestThrottler(this.scheduledThroughput.get()))) - ) - .flatMap(requestThrottler -> requestThrottler.processRequest(request, nextRequestMono)); + public Mono processRequest(RxDocumentServiceRequest request, Mono originalRequestMono) { + return this.requestThrottler.processRequest(request, originalRequestMono); } @Override - public void renewThroughputUsageCycle(double throughput) { + public double renewThroughputUsageCycle(double throughput) { this.scheduledThroughput.set(throughput); - this.requestThrottlerMapByRegion.values() - .stream() - .forEach(requestThrottler -> requestThrottler.renewThroughputUsageCycle(throughput)); - } - - @Override - public Mono close() { - return Mono.empty(); + return this.requestThrottler.renewThroughputUsageCycle(throughput); } } 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/controller/request/IThroughputRequestController.java index 38ec5e14be8be..2e1fcb4437c3c 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/controller/request/IThroughputRequestController.java @@ -4,7 +4,6 @@ package com.azure.cosmos.implementation.throughputControl.controller.request; import com.azure.cosmos.implementation.throughputControl.controller.IThroughputController; -import reactor.core.publisher.Mono; /** * Represents a throughput request controller. @@ -15,9 +14,10 @@ public interface IThroughputRequestController extends IThroughputController { * * Each request controller will maintain one to many request throttlers. * By calling this method, it will also renew the throughput usage cycles for the request throttlers. + * And it will return the throughput usage for the current cycle. * - * @param throughput - * @return + * @param throughput The scheduled throughput for the new cycle. + * @return The throughput usage for previous cycle. */ - void renewThroughputUsageCycle(double throughput); + double renewThroughputUsageCycle(double throughput); } 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/controller/request/PkRangesThroughputRequestController.java index 87ea0cb4d2c60..61f800148b6d8 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/controller/request/PkRangesThroughputRequestController.java @@ -3,7 +3,6 @@ package com.azure.cosmos.implementation.throughputControl.controller.request; -import com.azure.cosmos.implementation.GlobalEndpointManager; import com.azure.cosmos.implementation.PartitionKeyRange; import com.azure.cosmos.implementation.RxDocumentServiceRequest; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; @@ -13,13 +12,11 @@ import com.azure.cosmos.implementation.throughputControl.ThroughputRequestThrottler; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; -import java.net.URI; +import java.util.Comparator; import java.util.List; 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; @@ -30,40 +27,35 @@ public class PkRangesThroughputRequestController implements IThroughputRequestCo PartitionKeyInternalHelper.MinimumInclusiveEffectivePartitionKey, PartitionKeyInternalHelper.MaximumExclusiveEffectivePartitionKey, true, false); - private final GlobalEndpointManager globalEndpointManager; private final RxPartitionKeyRangeCache partitionKeyRangeCache; - private final ConcurrentHashMap> requestThrottlerMapByRegion; + private final ConcurrentHashMap requestThrottlerMap; private final String targetContainerRid; private double scheduledThroughput; private List pkRanges; public PkRangesThroughputRequestController( - GlobalEndpointManager globalEndpointManager, RxPartitionKeyRangeCache partitionKeyRangeCache, String targetContainerRid, double initialScheduledThroughput) { - checkNotNull(globalEndpointManager, "Global endpoint manager can not be null"); checkNotNull(partitionKeyRangeCache, "RxPartitionKeyRangeCache can not be null"); checkArgument(StringUtils.isNotEmpty(targetContainerRid), "Target container rid can not be null nor empty"); - this.globalEndpointManager = globalEndpointManager; this.partitionKeyRangeCache = partitionKeyRangeCache; - this.requestThrottlerMapByRegion = new ConcurrentHashMap<>(); + this.requestThrottlerMap = new ConcurrentHashMap<>(); this.targetContainerRid = targetContainerRid; this.scheduledThroughput = initialScheduledThroughput; } @Override - public void renewThroughputUsageCycle(double scheduledThroughput) { + public double renewThroughputUsageCycle(double scheduledThroughput) { this.scheduledThroughput = scheduledThroughput; double throughputPerPkRange = this.calculateThroughputPerPkRange(); - this.requestThrottlerMapByRegion.values() - .forEach(requestThrottlerMapByRegion -> { - requestThrottlerMapByRegion - .values() - .forEach(requestThrottler -> requestThrottler.renewThroughputUsageCycle(throughputPerPkRange)); - }); + return this.requestThrottlerMap.values() + .stream() + .map(requestThrottler -> requestThrottler.renewThroughputUsageCycle(throughputPerPkRange)) + .max(Comparator.naturalOrder())// return the max throughput usage among all regions + .get(); } @Override @@ -78,41 +70,29 @@ public boolean canHandleRequest(RxDocumentServiceRequest request) { return false; } - @Override - public Mono close() { - return Mono.empty(); - } - @Override @SuppressWarnings("unchecked") public Mono init() { return this.getPartitionKeyRanges(RANGE_INCLUDING_ALL_PARTITION_KEY_RANGES) - .flatMap(pkRanges -> { + .doOnSuccess(pkRanges -> { this.pkRanges = pkRanges; - return this.createRequestThrottlers(); + this.createRequestThrottlers(); }) .then(Mono.just((T)this)); } - private Mono createRequestThrottlers() { - // create request throttlers by region - return Flux.fromIterable(this.globalEndpointManager.getReadEndpoints()) - .flatMap(endpoint -> this.getOrCreateRegionRequestThrottlers(endpoint)) - .then(); - } - - private Mono> getOrCreateRegionRequestThrottlers(URI endpoint) { + private void createRequestThrottlers() { double throughputPerPkRange = this.calculateThroughputPerPkRange(); - return Mono.just(this.requestThrottlerMapByRegion.computeIfAbsent(endpoint, key -> { - ConcurrentHashMap requestThrottlerPerPkRange = - new ConcurrentHashMap<>(); - for (PartitionKeyRange pkRange : pkRanges) { - requestThrottlerPerPkRange.put(pkRange.getId(), new ThroughputRequestThrottler(throughputPerPkRange)); - } + for (PartitionKeyRange pkRange : pkRanges) { + requestThrottlerMap.compute(pkRange.getId(), (pkRangeId, requestThrottler) -> { + if (requestThrottler == null) { + requestThrottler = new ThroughputRequestThrottler(throughputPerPkRange); + } - return requestThrottlerPerPkRange; - })); + return requestThrottler; + }); + } } @Override @@ -120,19 +100,17 @@ public Mono processRequest(RxDocumentServiceRequest request, Mono next PartitionKeyRange resolvedPkRange = request.requestContext.resolvedPartitionKeyRange; // If we reach here, it means we should find the mapping pkRange - return this.getOrCreateRegionRequestThrottlers(this.globalEndpointManager.resolveServiceEndpoint(request)) - .flatMap(regionRequestThrottlers -> Mono.just(regionRequestThrottlers.get(resolvedPkRange.getId()))) - .flatMap(requestThrottler -> { - if (requestThrottler != null) { - return requestThrottler.processRequest(request, nextRequestMono); - } else { - logger.warn( - "Can not find matching request throttler to process request {} with pkRangeId {}", - request.getActivityId(), - resolvedPkRange.getId()); - return nextRequestMono; - } - }); + ThroughputRequestThrottler requestThrottler = this.requestThrottlerMap.get(resolvedPkRange.getId()); + + if (requestThrottler != null) { + return requestThrottler.processRequest(request, nextRequestMono); + } else { + logger.warn( + "Can not find matching request throttler to process request {} with pkRangeId {}", + request.getActivityId(), + resolvedPkRange.getId()); + return nextRequestMono; + } } private double calculateThroughputPerPkRange() { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosChangeFeedRequestOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosChangeFeedRequestOptions.java index 76c48b945714a..dc2ac89112282 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosChangeFeedRequestOptions.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosChangeFeedRequestOptions.java @@ -32,6 +32,7 @@ public final class CosmosChangeFeedRequestOptions { private ChangeFeedMode mode; private ChangeFeedStartFromInternal startFromInternal; private boolean isSplitHandlingDisabled; + private String throughputControlGroupName; private CosmosChangeFeedRequestOptions( FeedRangeInternal feedRange, @@ -335,6 +336,7 @@ CosmosChangeFeedRequestOptions withCosmosPagedFluxOptions( CosmosChangeFeedRequestOptions.createForProcessingFromContinuation( pagedFluxOptions.getRequestContinuation()); effectiveRequestOptions.setMaxPrefetchPageCount(this.getMaxPrefetchPageCount()); + effectiveRequestOptions.setThroughputControlGroupName(this.getThroughputControlGroupName()); } if (pagedFluxOptions.getMaxItemCount() != null) { @@ -383,4 +385,26 @@ public CosmosChangeFeedRequestOptions fullFidelity() { this.mode = ChangeFeedMode.FULL_FIDELITY; return this; } + + /** + * Get the throughput control group name. + * + * @return The throughput control group name. + */ + @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public String getThroughputControlGroupName() { + return this.throughputControlGroupName; + } + + /** + * Set the throughput control group name. + * + * @param throughputControlGroupName The throughput control group name. + * @return A {@link CosmosChangeFeedRequestOptions}. + */ + @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public CosmosChangeFeedRequestOptions setThroughputControlGroupName(String throughputControlGroupName) { + this.throughputControlGroupName = throughputControlGroupName; + return this; + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosItemRequestOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosItemRequestOptions.java index 80c7ed6976970..d8b04ba09d076 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosItemRequestOptions.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosItemRequestOptions.java @@ -288,12 +288,12 @@ RequestOptions toRequestOptions() { return requestOptions; } - @Beta(value = Beta.SinceVersion.V4_12_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) public String getThroughputControlGroupName() { return this.throughputControlGroupName; } - @Beta(value = Beta.SinceVersion.V4_12_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) public void setThroughputControlGroupName(String throughputControlGroupName) { this.throughputControlGroupName = throughputControlGroupName; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosQueryRequestOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosQueryRequestOptions.java index 32fa1b62ce59b..2514b0b153f80 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosQueryRequestOptions.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosQueryRequestOptions.java @@ -4,6 +4,7 @@ package com.azure.cosmos.models; import com.azure.cosmos.ConsistencyLevel; +import com.azure.cosmos.util.Beta; import java.util.Map; @@ -27,6 +28,7 @@ public class CosmosQueryRequestOptions { private Map properties; private boolean emptyPagesAllowed; private FeedRange feedRange; + private String throughputControlGroupName; /** * Instantiates a new query request options. @@ -54,6 +56,7 @@ public CosmosQueryRequestOptions() { this.partitionkey = options.partitionkey; this.queryMetricsEnabled = options.queryMetricsEnabled; this.emptyPagesAllowed = options.emptyPagesAllowed; + this.throughputControlGroupName = options.throughputControlGroupName; } /** @@ -381,4 +384,25 @@ FeedRange getFeedRange() { void setFeedRange(FeedRange feedRange) { this.feedRange = feedRange; } + + /** + * Get throughput control group name. + * @return The throughput control group name. + */ + @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public String getThroughputControlGroupName() { + return this.throughputControlGroupName; + } + + /** + * Set the throughput control group name. + * + * @param throughputControlGroupName The throughput control group name. + * @return A {@link CosmosQueryRequestOptions}. + */ + @Beta(value = Beta.SinceVersion.V4_13_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public CosmosQueryRequestOptions setThroughputControlGroupName(String throughputControlGroupName) { + this.throughputControlGroupName = throughputControlGroupName; + return this; + } } 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 0019980690c35..0956a781916ad 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/module-info.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/module-info.java @@ -51,6 +51,7 @@ opens com.azure.cosmos.implementation.clientTelemetry to com.fasterxml.jackson.databind; opens com.azure.cosmos.models to com.fasterxml.jackson.databind; opens com.azure.cosmos.util 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; opens com.azure.cosmos.implementation.throughputControl.controller.request; 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 new file mode 100644 index 0000000000000..2f77edeefa0e1 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/samples/java/com/azure/cosmos/ThroughputControlCodeSnippet.java @@ -0,0 +1,56 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos; + +import com.azure.cosmos.implementation.TestConfigurations; + +import java.time.Duration; + +public class ThroughputControlCodeSnippet { + private CosmosAsyncClient client; + private CosmosAsyncDatabase database; + private CosmosAsyncContainer container; + + public ThroughputControlCodeSnippet() { + this.client = new CosmosClientBuilder() + .endpoint(TestConfigurations.HOST) + .key(TestConfigurations.MASTER_KEY) + .contentResponseOnWriteEnabled(true) + .consistencyLevel(ConsistencyLevel.SESSION) + .buildAsyncClient(); + + this.database = this.client.getDatabase("TestDB"); + this.container = this.database.getContainer("TestContainer"); + } + + public void codeSnippetForEnableLocalThroughputControl() { + // BEGIN: com.azure.cosmos.throughputControl.localControl + ThroughputControlGroupConfig groupConfig = + new ThroughputControlGroupConfigBuilder() + .setGroupName("localControlGroup") + .setTargetThroughputThreshold(0.1) + .build(); + + container.enableLocalThroughputControlGroup(groupConfig); + // END: com.azure.cosmos.throughputControl.localControl + } + + public void codeSnippetForEnableGlobalThroughputControl() { + // BEGIN: com.azure.cosmos.throughputControl.globalControl + ThroughputControlGroupConfig groupConfig = + new ThroughputControlGroupConfigBuilder() + .setGroupName("localControlGroup") + .setTargetThroughputThreshold(0.1) + .build(); + + GlobalThroughputControlConfig globalControlConfig = + this.client.createGlobalThroughputControlConfigBuilder(database.getId(), container.getId()) + .setControlItemRenewInterval(Duration.ofSeconds(5)) + .setControlItemExpireInterval(Duration.ofSeconds(10)) + .build(); + + container.enableGlobalThroughputControlGroup(groupConfig, globalControlConfig); + // END: com.azure.cosmos.throughputControl.globalControl + } +} diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/directconnectivity/ReflectionUtils.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/directconnectivity/ReflectionUtils.java index 817ff9c3f6f3f..8f4efc382a209 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/directconnectivity/ReflectionUtils.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/directconnectivity/ReflectionUtils.java @@ -220,14 +220,22 @@ public static void setTransportClient(ConsistencyWriter consistencyWriter, Trans } @SuppressWarnings("unchecked") - public static ConcurrentHashMap getRequestThrottlerMap(GlobalThroughputRequestController requestController) { - return get(ConcurrentHashMap.class, requestController, "requestThrottlerMapByRegion"); + public static ThroughputRequestThrottler getRequestThrottler(GlobalThroughputRequestController requestController) { + return get(ThroughputRequestThrottler.class, requestController, "requestThrottler"); } @SuppressWarnings("unchecked") - public static ConcurrentHashMap> getRequestThrottlerMap( + public static void setRequestThrottler( + GlobalThroughputRequestController requestController, + ThroughputRequestThrottler throughputRequestThrottler) { + + set(requestController, throughputRequestThrottler, "requestThrottler"); + } + + @SuppressWarnings("unchecked") + public static ConcurrentHashMap getRequestThrottler( PkRangesThroughputRequestController requestController) { - return get(ConcurrentHashMap.class, requestController, "requestThrottlerMapByRegion"); + return get(ConcurrentHashMap.class, requestController, "requestThrottlerMap"); } } diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationTokenSourceTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationTokenSourceTests.java new file mode 100644 index 0000000000000..e505391a7c181 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationTokenSourceTests.java @@ -0,0 +1,28 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl; + +import org.testng.annotations.Test; + +import static org.assertj.core.api.Assertions.assertThat; + +public class LinkedCancellationTokenSourceTests { + + @Test(groups = "unit") + public void linkedCancellationTokenSource() { + LinkedCancellationTokenSource source1 = new LinkedCancellationTokenSource(); + LinkedCancellationTokenSource source2 = new LinkedCancellationTokenSource(source1.getToken()); + LinkedCancellationToken source2Token = source2.getToken(); + LinkedCancellationTokenSource source3 = new LinkedCancellationTokenSource(source2.getToken()); + LinkedCancellationToken source3Token = source3.getToken(); + + source1.close(); + + assertThat(source2.isClosed()).isTrue(); + assertThat(source2Token.isCancellationRequested()).isTrue(); + assertThat(source3.isClosed()).isTrue(); + assertThat(source3Token.isCancellationRequested()).isTrue(); + + } +} diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupConfigConfigurationTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupConfigConfigurationTests.java new file mode 100644 index 0000000000000..67005d75a6bfe --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupConfigConfigurationTests.java @@ -0,0 +1,54 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl; + +import com.azure.cosmos.CosmosAsyncClient; +import com.azure.cosmos.CosmosAsyncContainer; +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.ThroughputControlGroupConfig; +import com.azure.cosmos.ThroughputControlGroupConfigBuilder; +import com.azure.cosmos.rx.TestSuiteBase; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Factory; +import org.testng.annotations.Test; + +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +public class ThroughputControlGroupConfigConfigurationTests extends TestSuiteBase { + private CosmosAsyncClient client; + private CosmosAsyncContainer container; + + @Factory(dataProvider = "clientBuildersWithSessionConsistency") + public ThroughputControlGroupConfigConfigurationTests(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + this.subscriberValidationTimeout = TIMEOUT; + } + + @Test(groups = { "emulator" }) + public void validateMultipleDefaultGroups() { + ThroughputControlGroupConfig groupConfig = + new ThroughputControlGroupConfigBuilder() + .setGroupName("group-1") + .setTargetThroughput(10) + .setDefault(true) + .build(); + container.enableLocalThroughputControlGroup(groupConfig); + + ThroughputControlGroupConfig groupConfig2 = + new ThroughputControlGroupConfigBuilder() + .setGroupName("group-2") + .setTargetThroughputThreshold(1.0) + .setDefault(true) + .build(); + assertThatThrownBy(() -> container.enableLocalThroughputControlGroup(groupConfig2)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("A default group already exists"); + } + + @BeforeClass(groups = { "emulator" }, timeOut = 4 * SETUP_TIMEOUT) + public void before_ThroughputControlGroupConfigurationTests() { + client = getClientBuilder().buildAsyncClient(); + container = getSharedMultiPartitionCosmosContainer(client); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupConfigurationTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupConfigurationTests.java deleted file mode 100644 index a6c02a6fa41c6..0000000000000 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupConfigurationTests.java +++ /dev/null @@ -1,41 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -package com.azure.cosmos.implementation.throughputControl; - -import com.azure.cosmos.CosmosAsyncClient; -import com.azure.cosmos.CosmosAsyncContainer; -import com.azure.cosmos.CosmosClientBuilder; -import com.azure.cosmos.rx.TestSuiteBase; -import org.testng.annotations.BeforeClass; -import org.testng.annotations.Factory; -import org.testng.annotations.Test; - -import static org.assertj.core.api.Assertions.assertThatThrownBy; - -public class ThroughputControlGroupConfigurationTests extends TestSuiteBase { - private CosmosAsyncClient client; - private CosmosAsyncContainer container; - - // TODO: reenable when enable enableThroughputLocalControlGroup public API -// @Factory(dataProvider = "clientBuildersWithSessionConsistency") -// public ThroughputControlGroupConfigurationTests(CosmosClientBuilder clientBuilder) { -// super(clientBuilder); -// this.subscriberValidationTimeout = TIMEOUT; -// } -// -// @Test(groups = { "emulator" }) -// public void validateMultipleDefaultGroups() { -// container.enableThroughputLocalControlGroup("group-1", 10, true); -// -// assertThatThrownBy(() -> container.enableThroughputLocalControlGroup("group-2", 10, true)) -// .isInstanceOf(IllegalArgumentException.class) -// .hasMessage("A default group already exists"); -// } -// -// @BeforeClass(groups = { "emulator" }, timeOut = 4 * SETUP_TIMEOUT) -// public void before_ThroughputControlGroupConfigurationTests() { -// client = getClientBuilder().buildAsyncClient(); -// container = getSharedMultiPartitionCosmosContainer(client); -// } -} diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java index 5328f41823624..11299512b8e33 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java @@ -7,109 +7,284 @@ import com.azure.cosmos.ConnectionMode; import com.azure.cosmos.CosmosAsyncClient; import com.azure.cosmos.CosmosAsyncContainer; +import com.azure.cosmos.CosmosAsyncDatabase; import com.azure.cosmos.CosmosClientBuilder; -import com.azure.cosmos.ThroughputControlGroup; +import com.azure.cosmos.CosmosDiagnostics; +import com.azure.cosmos.CosmosException; +import com.azure.cosmos.ThroughputControlGroupConfig; +import com.azure.cosmos.ThroughputControlGroupConfigBuilder; +import com.azure.cosmos.GlobalThroughputControlConfig; +import com.azure.cosmos.implementation.OperationType; +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; +import com.azure.cosmos.models.CosmosContainerProperties; +import com.azure.cosmos.models.CosmosContainerRequestOptions; import com.azure.cosmos.models.CosmosItemRequestOptions; import com.azure.cosmos.models.CosmosItemResponse; +import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.FeedRange; +import com.azure.cosmos.models.FeedResponse; import com.azure.cosmos.models.PartitionKey; import com.azure.cosmos.rx.TestSuiteBase; import org.testng.annotations.BeforeClass; +import org.testng.annotations.DataProvider; import org.testng.annotations.Factory; import org.testng.annotations.Test; +import java.time.Duration; import java.util.UUID; import static org.assertj.core.api.Assertions.assertThat; public class ThroughputControlTests extends TestSuiteBase { + // Delete collections in emulator is not instant, + // so to avoid get 500 back, we are adding delay for creating the collection with same name, since in this case we want to test 410/1000 + private final static int COLLECTION_RECREATION_TIME_DELAY = 5000; + private CosmosAsyncClient client; + private CosmosAsyncDatabase database; private CosmosAsyncContainer container; -// @Factory(dataProvider = "clientBuildersWithSessionConsistency") -// public ThroughputControlTests(CosmosClientBuilder clientBuilder) { -// super(clientBuilder); -// this.subscriberValidationTimeout = TIMEOUT; -// } -// -// @Test(groups = {"emulator"}, timeOut = TIMEOUT) -// public void readItem() throws Exception { -// container.enableThroughputLocalControlGroup("group-1", 5, true); -// ThroughputControlGroup group2 = container.enableThroughputLocalControlGroup("group-2", 0.9); -// -// TestItem docDefinition = getDocumentDefinition(); -// container.createItem(docDefinition).block(); // since no group is defined, this will fall into the default control group -// -// CosmosItemRequestOptions requestOptions = new CosmosItemRequestOptions(); -// -// // Test read operation which will fall into the default control group, which should be throttled -// // and be succeeded during retry -// CosmosItemResponse readItemResponse = container.readItem(docDefinition.getId(), -// new PartitionKey(docDefinition.getMypk()), -// requestOptions, -// TestItem.class).block(); -// this.validateRequestThrottled( -// readItemResponse.getDiagnostics().toString(), -// BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); -// -// // Test read operation which will use a different control group, so it will pass -// requestOptions.setThroughputControlGroupName(group2.getGroupName()); -// CosmosItemResponse readItemResponse2 = container.readItem(docDefinition.getId(), -// new PartitionKey(docDefinition.getMypk()), -// requestOptions, -// TestItem.class).block(); -// this.validateRequestNotThrottled( -// readItemResponse2.getDiagnostics().toString(), -// BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); -// -// // Test read operation which will use an undefined control group, it will fall back to default group -// // but since the throughput usage has been reset, this request will not be throttled -// -// requestOptions.setThroughputControlGroupName("Undefined"); -// CosmosItemResponse readItemResponse3 = container.readItem(docDefinition.getId(), -// new PartitionKey(docDefinition.getMypk()), -// requestOptions, -// TestItem.class).block(); -// this.validateRequestNotThrottled( -// readItemResponse3.getDiagnostics().toString(), -// BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); -// } -// -// // TODO: add tests for other operations -// // TODO: add tests for split and collection recreation -// -// @BeforeClass(groups = { "emulator" }, timeOut = 4 * SETUP_TIMEOUT) -// public void before_ThroughputBudgetControllerTest() { -// client = getClientBuilder().buildAsyncClient(); -// container = getSharedMultiPartitionCosmosContainer(client); -// } -// -// private static TestItem getDocumentDefinition() { -// return new TestItem( -// UUID.randomUUID().toString(), -// UUID.randomUUID().toString(), -// UUID.randomUUID().toString() -// ); -// } -// -// private void validateRequestThrottled(String cosmosDiagnostics, ConnectionMode connectionMode) { -// assertThat(cosmosDiagnostics).isNotEmpty(); -// -// if (connectionMode == ConnectionMode.DIRECT) { -// assertThat(cosmosDiagnostics).contains("\"statusCode\":429"); -// assertThat(cosmosDiagnostics).contains("\"subStatusCode\":10003"); -// } else if (connectionMode == ConnectionMode.GATEWAY) { -// assertThat(cosmosDiagnostics).contains("\"statusAndSubStatusCodes\":[[429,10003]"); -// } -// } -// -// private void validateRequestNotThrottled(String cosmosDiagnostics, ConnectionMode connectionMode) { -// assertThat(cosmosDiagnostics).isNotEmpty(); -// -// if (connectionMode == ConnectionMode.DIRECT) { -// assertThat(cosmosDiagnostics).doesNotContain("\"statusCode\":429"); -// assertThat(cosmosDiagnostics).doesNotContain("\"subStatusCode\":10003"); -// } else if (connectionMode == ConnectionMode.GATEWAY) { -// assertThat(cosmosDiagnostics).doesNotContain("\"statusAndSubStatusCodes\":[[429,10003]"); -// } -// } + @Factory(dataProvider = "simpleClientBuildersForDirectTcpWithoutRetryOnThrottledRequests") + public ThroughputControlTests(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + this.subscriberValidationTimeout = TIMEOUT; + } + + @DataProvider + public static Object[][] operationTypeProvider() { + return new Object[][]{ + { OperationType.Read }, + { OperationType.Replace }, + { OperationType.Create }, + { OperationType.Delete }, + { OperationType.Query }, + // { OperationType.ReadFeed } // changeFeed only go for gateway + }; + } + + @Test(groups = {"emulator"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT) + public void throughputLocalControl(OperationType operationType) { + // 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() + .setGroupName("group-" + UUID.randomUUID()) + .setTargetThroughput(6) + .build(); + container.enableLocalThroughputControlGroup(groupConfig); + + 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()); + } + + @Test(groups = {"emulator"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT) + public void throughputGlobalControl(OperationType operationType) { + String controlContainerId = "throughputControlContainer"; + CosmosAsyncContainer controlContainer = database.getContainer(controlContainerId); + database.createContainerIfNotExists(controlContainer.getId(), "/groupId").block(); + + // 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() + .setGroupName("group-" + UUID.randomUUID()) + .setTargetThroughput(6) + .build(); + + GlobalThroughputControlConfig globalControlConfig = this.client.createGlobalThroughputControlConfigBuilder(this.database.getId(), controlContainerId) + .setControlItemRenewInterval(Duration.ofSeconds(5)) + .setControlItemExpireInterval(Duration.ofSeconds(20)) + .build(); + container.enableGlobalThroughputControlGroup(groupConfig, globalControlConfig); + + 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 same group. which will get throttled + CosmosDiagnostics cosmosDiagnostics = performDocumentOperation(this.container, operationType, createdItem, groupConfig.getGroupName()); + this.validateRequestThrottled( + cosmosDiagnostics.toString(), + BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + } + + @Test(groups = {"emulator"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT) + public void throughputLocalControlForContainerCreateDeleteWithSameName(OperationType operationType) throws InterruptedException { + ConnectionMode connectionMode = BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode(); + if (connectionMode == ConnectionMode.GATEWAY) { + // for gateway connection mode, gateway will handle the 410/1000 and retry. Hence the collection cache and container controller will not be refreshed. + // There is no point for this tests for gateway mode. + return; + } + + // step1: create container + String testContainerId = UUID.randomUUID().toString(); + CosmosContainerProperties containerProperties = getCollectionDefinition(testContainerId); + CosmosAsyncContainer createdContainer = createCollection(this.database, containerProperties, new CosmosContainerRequestOptions()); + + // The create document in this test usually takes around 6.29RU, + // pick a RU super small here so we know it will throttle requests for several cycles/seconds + ThroughputControlGroupConfig groupConfig = + new ThroughputControlGroupConfigBuilder() + .setGroupName("group-" + UUID.randomUUID()) + .setTargetThroughput(1) + .build(); + container.enableLocalThroughputControlGroup(groupConfig); + createdContainer.enableLocalThroughputControlGroup(groupConfig); + + CosmosItemRequestOptions requestOptions = new CosmosItemRequestOptions(); + requestOptions.setContentResponseOnWriteEnabled(true); + requestOptions.setThroughputControlGroupName(groupConfig.getGroupName()); + + // 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()); + + // Step 3: delete the container + safeDeleteCollection(createdContainer); + Thread.sleep(COLLECTION_RECREATION_TIME_DELAY); + + // step 4: recreate the container with the same name + createdContainer = createCollection(this.database, containerProperties, new CosmosContainerRequestOptions()); + + // Step 5: operation which will trigger cache refresh and a new container controller to be built + 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()); + + // 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()); + } + + @BeforeClass(groups = { "emulator" }, timeOut = 4 * SETUP_TIMEOUT) + public void before_ThroughputBudgetControllerTest() { + client = getClientBuilder().buildAsyncClient(); + database = getSharedCosmosDatabase(client); + container = getSharedMultiPartitionCosmosContainer(client); + } + + private static TestItem getDocumentDefinition() { + return getDocumentDefinition(null); + } + + private static TestItem getDocumentDefinition(String partitionKey) { + return new TestItem( + UUID.randomUUID().toString(), + StringUtils.isEmpty(partitionKey) ? UUID.randomUUID().toString() : partitionKey, + UUID.randomUUID().toString() + ); + } + + private void validateRequestThrottled(String cosmosDiagnostics, ConnectionMode connectionMode) { + assertThat(cosmosDiagnostics).isNotEmpty(); + assertThat(cosmosDiagnostics).contains("\"statusCode\":429"); + assertThat(cosmosDiagnostics).contains("\"subStatusCode\":10003"); + } + + private void validateRequestNotThrottled(String cosmosDiagnostics, ConnectionMode connectionMode) { + assertThat(cosmosDiagnostics).isNotEmpty(); + assertThat(cosmosDiagnostics).doesNotContain("\"statusCode\":429"); + assertThat(cosmosDiagnostics).doesNotContain("\"subStatusCode\":10003"); + } + + private CosmosDiagnostics performDocumentOperation( + CosmosAsyncContainer cosmosAsyncContainer, + OperationType operationType, + TestItem createdItem, + String throughputControlGroup) { + try { + if (operationType == OperationType.Query) { + CosmosQueryRequestOptions queryRequestOptions = new CosmosQueryRequestOptions(); + if (!StringUtils.isEmpty(throughputControlGroup)) { + queryRequestOptions.setThroughputControlGroupName(throughputControlGroup); + } + + String query = String.format("SELECT * from c where c.mypk = '%s'", createdItem.getMypk()); + FeedResponse itemFeedResponse = + cosmosAsyncContainer.queryItems(query, queryRequestOptions, TestItem.class).byPage().blockFirst(); + + return itemFeedResponse.getCosmosDiagnostics(); + } + + if (operationType == OperationType.ReadFeed) { + CosmosChangeFeedRequestOptions changeFeedRequestOptions = CosmosChangeFeedRequestOptions + .createForProcessingFromBeginning(FeedRange.forFullRange()); + if (!StringUtils.isEmpty(throughputControlGroup)) { + changeFeedRequestOptions.setThroughputControlGroupName(throughputControlGroup); + } + + FeedResponse itemFeedResponse = cosmosAsyncContainer.queryChangeFeed(changeFeedRequestOptions, TestItem.class).byPage().blockFirst(); + return itemFeedResponse.getCosmosDiagnostics(); + } + + if (operationType == OperationType.Read + || operationType == OperationType.Delete + || operationType == OperationType.Replace + || operationType == OperationType.Create) { + CosmosItemRequestOptions itemRequestOptions = new CosmosItemRequestOptions(); + if (!StringUtils.isEmpty((throughputControlGroup))) { + itemRequestOptions.setThroughputControlGroupName(throughputControlGroup); + } + + if (operationType == OperationType.Read) { + return cosmosAsyncContainer.readItem( + createdItem.getId(), + new PartitionKey(createdItem.getMypk()), + itemRequestOptions, + TestItem.class).block().getDiagnostics(); + } + + if (operationType == OperationType.Replace) { + return cosmosAsyncContainer.replaceItem( + createdItem, + createdItem.getId(), + new PartitionKey(createdItem.getMypk()), + itemRequestOptions).block().getDiagnostics(); + } + + if (operationType == OperationType.Delete) { + return cosmosAsyncContainer.deleteItem(createdItem, itemRequestOptions).block().getDiagnostics(); + } + + if (operationType == OperationType.Create) { + TestItem newItem = getDocumentDefinition(createdItem.getMypk()); + return cosmosAsyncContainer.createItem(newItem, itemRequestOptions).block().getDiagnostics(); + } + } + + throw new IllegalArgumentException("The operation type is not supported"); + } catch (CosmosException cosmosException) { + return cosmosException.getDiagnostics(); + } + } + + // TODO: add tests split } diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/GlobalThroughputRequestControllerTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/GlobalThroughputRequestControllerTests.java index ec2d53165b3e3..a551e94f6af3a 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/GlobalThroughputRequestControllerTests.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/GlobalThroughputRequestControllerTests.java @@ -3,65 +3,34 @@ package com.azure.cosmos.implementation.throughputControl.controller; -import com.azure.cosmos.implementation.GlobalEndpointManager; import com.azure.cosmos.implementation.RxDocumentServiceRequest; -import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList; 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 org.mockito.Mockito; -import org.testng.annotations.BeforeClass; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; import reactor.core.publisher.Mono; import reactor.test.StepVerifier; import reactor.test.publisher.TestPublisher; -import java.net.URI; -import java.net.URISyntaxException; -import java.util.ArrayList; -import java.util.Collections; -import java.util.HashSet; -import java.util.List; -import java.util.Set; -import java.util.concurrent.ConcurrentHashMap; - import static org.assertj.core.api.Assertions.assertThat; public class GlobalThroughputRequestControllerTests { - private static double scheduledThroughput; - private static GlobalEndpointManager globalEndpointManager; - private static URI readLocation; - + private static double scheduledThroughput = 2.0; private GlobalThroughputRequestController requestController; - @BeforeClass(groups = "unit") - public static void beforeClass_GlobalThroughputRequestControllerTests() throws URISyntaxException { - scheduledThroughput = 2.0; - - readLocation = new URI("https://read-localtion1.documents.azure.com"); - - globalEndpointManager = Mockito.mock(GlobalEndpointManager.class); - Mockito.doReturn(new UnmodifiableList<>(Collections.singletonList(readLocation))).when(globalEndpointManager).getReadEndpoints(); - } - @BeforeMethod(groups = "unit") public void before_GlobalThroughputRequestControllerTest() { - requestController = new GlobalThroughputRequestController(globalEndpointManager, scheduledThroughput); + requestController = new GlobalThroughputRequestController(scheduledThroughput); } @Test(groups = "unit") public void init() { + // Test init can complete without error requestController.init().subscribe(); - - Set locations = new HashSet<>(); - locations.add(readLocation); - - ConcurrentHashMap requestThrottlerMapByRegion = ReflectionUtils.getRequestThrottlerMap(requestController); - assertThat(requestThrottlerMapByRegion).size().isEqualTo(locations.size()); - assertThat(Collections.list(requestThrottlerMapByRegion.keys())).containsAll(locations); } @Test(groups = "unit") @@ -71,16 +40,15 @@ public void canHandleRequest() { } @Test(groups = "unit") - public void processRequest() throws URISyntaxException { + public void processRequest() { requestController.init().subscribe(); - ConcurrentHashMap requestThrottlerMapByRegion = ReflectionUtils.getRequestThrottlerMap(requestController); - ThroughputRequestThrottler writeLocationThrottlerSpy = Mockito.spy(requestThrottlerMapByRegion.get(readLocation)); - requestThrottlerMapByRegion.put(readLocation, writeLocationThrottlerSpy); + ThroughputRequestThrottler requestThrottler = ReflectionUtils.getRequestThrottler(requestController); + ThroughputRequestThrottler requestThrottlerSpy = Mockito.spy(requestThrottler); + ReflectionUtils.setRequestThrottler(requestController, requestThrottlerSpy); // First request: Can find the matching region request throttler in request controller RxDocumentServiceRequest request1Mock = Mockito.mock(RxDocumentServiceRequest.class); - Mockito.doReturn(readLocation).when(globalEndpointManager).resolveServiceEndpoint(request1Mock); TestPublisher request1MonoPublisher = TestPublisher.create(); Mono request1Mono = request1MonoPublisher.mono(); @@ -90,40 +58,19 @@ public void processRequest() throws URISyntaxException { .then(() -> request1MonoPublisher.emit(storeResponse1Mock)) .expectNext(storeResponse1Mock) .verifyComplete(); - Mockito.verify(writeLocationThrottlerSpy, Mockito.times(1)).processRequest(request1Mock, request1Mono); - - // Second request: Cannot find the matching region request throttler in request controller, will create a new one - RxDocumentServiceRequest request2Mock = Mockito.mock(RxDocumentServiceRequest.class); - Mockito.doReturn(new URI("https://write-localtion2.documents.azure.com")).when(globalEndpointManager).resolveServiceEndpoint(request2Mock); - - TestPublisher request2MonoPublisher = TestPublisher.create(); - Mono request2Mono = request2MonoPublisher.mono(); - StoreResponse storeResponse2Mock = Mockito.mock(StoreResponse.class); - - StepVerifier.create(requestController.processRequest(request2Mock, request2Mono)) - .then(() -> request2MonoPublisher.emit(storeResponse2Mock)) - .expectNext(storeResponse2Mock) - .verifyComplete(); - - assertThat(requestThrottlerMapByRegion).size().isEqualTo(2); + Mockito.verify(requestThrottlerSpy, Mockito.times(1)).processRequest(request1Mock, request1Mono); } @Test(groups = "unit") public void renewThroughputUsageCycle() { requestController.init().subscribe(); - ConcurrentHashMap requestThrottlerMapByRegion = ReflectionUtils.getRequestThrottlerMap(requestController); - List requestThrottlerSpies = new ArrayList<>(); - for (URI location : Collections.list(requestThrottlerMapByRegion.keys())) { - ThroughputRequestThrottler requestThrottlerSpy = Mockito.spy(requestThrottlerMapByRegion.get(location)); - requestThrottlerSpies.add(requestThrottlerSpy); - requestThrottlerMapByRegion.put(location, requestThrottlerSpy); - } + ThroughputRequestThrottler requestThrottler = ReflectionUtils.getRequestThrottler(requestController); + ThroughputRequestThrottler requestThrottlerSpy = Mockito.spy(requestThrottler); + ReflectionUtils.setRequestThrottler(requestController, requestThrottlerSpy); double newScheduledThroughput = 3.0; requestController.renewThroughputUsageCycle(newScheduledThroughput); - for (ThroughputRequestThrottler requestThrottlerSpy : requestThrottlerSpies) { - Mockito.verify(requestThrottlerSpy, Mockito.times(1)).renewThroughputUsageCycle(newScheduledThroughput); - } + Mockito.verify(requestThrottlerSpy, Mockito.times(1)).renewThroughputUsageCycle(newScheduledThroughput); } } diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/PkRangesThroughputRequestControllerTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/PkRangesThroughputRequestControllerTests.java index c26914b0aa9d8..7b42becbbbbbf 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/PkRangesThroughputRequestControllerTests.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/PkRangesThroughputRequestControllerTests.java @@ -4,12 +4,10 @@ package com.azure.cosmos.implementation.throughputControl.controller; import com.azure.cosmos.implementation.DocumentServiceRequestContext; -import com.azure.cosmos.implementation.GlobalEndpointManager; import com.azure.cosmos.implementation.MetadataDiagnosticsContext; import com.azure.cosmos.implementation.PartitionKeyRange; import com.azure.cosmos.implementation.RxDocumentServiceRequest; import com.azure.cosmos.implementation.Utils; -import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList; import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils; import com.azure.cosmos.implementation.directconnectivity.StoreResponse; @@ -25,16 +23,12 @@ import reactor.test.StepVerifier; import reactor.test.publisher.TestPublisher; -import java.net.URI; import java.net.URISyntaxException; import java.util.ArrayList; import java.util.Collections; -import java.util.HashSet; import java.util.List; -import java.util.Set; import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; -import java.util.stream.Collectors; import static org.assertj.core.api.Assertions.assertThat; import static org.mockito.Matchers.any; @@ -49,10 +43,8 @@ public class PkRangesThroughputRequestControllerTests { private static String targetCollectionRid; private static double scheduledThroughput; private static PartitionKeyRange randomPkRange; - private static GlobalEndpointManager globalEndpointManager; private static RxPartitionKeyRangeCache pkRangeCache; private static List pkRanges; - private static URI readLocation; private PkRangesThroughputRequestController requestController; @@ -62,10 +54,6 @@ public static void beforeClass_PkRangesThroughputRequestControllerTests() throws scheduledThroughput = 2.0; randomPkRange = new PartitionKeyRange(UUID.randomUUID().toString(), "randomMin", "randomMax"); - readLocation = new URI("https://read-localtion1.documents.azure.com"); - globalEndpointManager = Mockito.mock(GlobalEndpointManager.class); - Mockito.doReturn(new UnmodifiableList<>(Collections.singletonList(readLocation))).when(globalEndpointManager).getReadEndpoints(); - PartitionKeyRange pkRange1 = new PartitionKeyRange(UUID.randomUUID().toString(), "AA", "BB"); PartitionKeyRange pkRange2 = new PartitionKeyRange(UUID.randomUUID().toString(), "BB", "CC"); pkRanges = new ArrayList<>(); @@ -84,26 +72,13 @@ public static void beforeClass_PkRangesThroughputRequestControllerTests() throws @BeforeMethod(groups = "unit") public void before_PkRangesThroughputRequestControllerTests() { - requestController = new PkRangesThroughputRequestController(globalEndpointManager, pkRangeCache, targetCollectionRid, scheduledThroughput); + requestController = new PkRangesThroughputRequestController(pkRangeCache, targetCollectionRid, scheduledThroughput); } @Test(groups = "unit") public void init() { + // Test init can complete without error requestController.init().subscribe(); - - ConcurrentHashMap> requestThrottlerMap = - ReflectionUtils.getRequestThrottlerMap(requestController); - - Set locations = new HashSet<>(); - locations.add(readLocation); - - assertThat(requestThrottlerMap).size().isEqualTo(locations.size()); - assertThat(Collections.list(requestThrottlerMap.keys())).containsAll(locations); - - for (URI location : locations) { - assertThat(Collections.list(requestThrottlerMap.get(location).keys())) - .containsAll(pkRanges.stream().map(PartitionKeyRange::getId).collect(Collectors.toList())); - } } @Test(groups = "unit") @@ -124,19 +99,17 @@ public void canHandleRequest() { } @Test(groups = "unit") - public void processRequest() throws URISyntaxException { + public void processRequest() { requestController.init().subscribe(); - ConcurrentHashMap> requestThrottlerByRegion = - ReflectionUtils.getRequestThrottlerMap(requestController); + ConcurrentHashMap requestThrottlerMap = + ReflectionUtils.getRequestThrottler(requestController); PartitionKeyRange pkRange = pkRanges.get(0); - ConcurrentHashMap requestThrottlerByPkRangeId = requestThrottlerByRegion.get(readLocation); - ThroughputRequestThrottler writeLocationThrottlerSpy = Mockito.spy(requestThrottlerByPkRangeId.get(pkRange.getId())); - requestThrottlerByPkRangeId.put(pkRange.getId(), writeLocationThrottlerSpy); + ThroughputRequestThrottler writeLocationThrottlerSpy = Mockito.spy(requestThrottlerMap.get(pkRange.getId())); + requestThrottlerMap.put(pkRange.getId(), writeLocationThrottlerSpy); - // First request: Can find the matching region request throttler in request controller - RxDocumentServiceRequest request1Mock = this.createMockRequest(pkRange, readLocation); + RxDocumentServiceRequest request1Mock = this.createMockRequest(pkRange); TestPublisher request1MonoPublisher = TestPublisher.create(); Mono request1Mono = request1MonoPublisher.mono(); @@ -147,36 +120,20 @@ public void processRequest() throws URISyntaxException { .expectNext(storeResponse1Mock) .verifyComplete(); Mockito.verify(writeLocationThrottlerSpy, Mockito.times(1)).processRequest(request1Mock, request1Mono); - - // Second request: Cannot find the matching region request throttler in request controller, will create a new one - RxDocumentServiceRequest request2Mock = this.createMockRequest(pkRange, new URI("https://write-localtion2.documents.azure.com")); - - TestPublisher request2MonoPublisher = TestPublisher.create(); - Mono request2Mono = request2MonoPublisher.mono(); - StoreResponse storeResponse2Mock = Mockito.mock(StoreResponse.class); - - StepVerifier.create(requestController.processRequest(request2Mock, request2Mono)) - .then(() -> request2MonoPublisher.emit(storeResponse2Mock)) - .expectNext(storeResponse2Mock) - .verifyComplete(); - - assertThat(requestThrottlerByRegion).size().isEqualTo(2); } @Test(groups = "unit") public void renewThroughputUsageCycle() { requestController.init().subscribe(); - ConcurrentHashMap> requestThrottlerByRegion = - ReflectionUtils.getRequestThrottlerMap(requestController); + ConcurrentHashMap requestThrottlerMap = + ReflectionUtils.getRequestThrottler(requestController); List requestThrottlerSpies = new ArrayList<>(); - for (ConcurrentHashMap requestThrottlerByPkRangeId : requestThrottlerByRegion.values()) { - for (String pkRangeId: Collections.list(requestThrottlerByPkRangeId.keys())) { - ThroughputRequestThrottler requestThrottlerSpy = Mockito.spy(requestThrottlerByPkRangeId.get(pkRangeId)); - requestThrottlerSpies.add(requestThrottlerSpy); - requestThrottlerByPkRangeId.put(pkRangeId, requestThrottlerSpy); - } + for (String pkRangeId : Collections.list(requestThrottlerMap.keys())) { + ThroughputRequestThrottler requestThrottlerSpy = Mockito.spy(requestThrottlerMap.get(pkRangeId)); + requestThrottlerSpies.add(requestThrottlerSpy); + requestThrottlerMap.put(pkRangeId, requestThrottlerSpy); } double newScheduledThroughput = 3.0; @@ -187,12 +144,11 @@ public void renewThroughputUsageCycle() { } } - private RxDocumentServiceRequest createMockRequest(PartitionKeyRange resolvedPkRange, URI serviceEndpoint) { + private RxDocumentServiceRequest createMockRequest(PartitionKeyRange resolvedPkRange) { RxDocumentServiceRequest requestMock = Mockito.mock(RxDocumentServiceRequest.class); DocumentServiceRequestContext requestContextMock = Mockito.mock(DocumentServiceRequestContext.class); requestContextMock.resolvedPartitionKeyRange = resolvedPkRange; requestMock.requestContext = requestContextMock; - Mockito.doReturn(serviceEndpoint).when(globalEndpointManager).resolveServiceEndpoint(requestMock); return requestMock; } diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java index 351c396e138ac..fc60cf95c087b 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java @@ -850,15 +850,15 @@ public static void validateQueryFailure(Flux> flowable, @DataProvider public static Object[][] clientBuilders() { - return new Object[][]{{createGatewayRxDocumentClient(ConsistencyLevel.SESSION, false, null, true)}}; + return new Object[][]{{createGatewayRxDocumentClient(ConsistencyLevel.SESSION, false, null, true, true)}}; } @DataProvider public static Object[][] clientBuildersWithSessionConsistency() { return new Object[][]{ - {createDirectRxDocumentClient(ConsistencyLevel.SESSION, Protocol.HTTPS, false, null, true)}, - {createDirectRxDocumentClient(ConsistencyLevel.SESSION, Protocol.TCP, false, null, true)}, - {createGatewayRxDocumentClient(ConsistencyLevel.SESSION, false, null, true)} + {createDirectRxDocumentClient(ConsistencyLevel.SESSION, Protocol.HTTPS, false, null, true, true)}, + {createDirectRxDocumentClient(ConsistencyLevel.SESSION, Protocol.TCP, false, null, true, true)}, + {createGatewayRxDocumentClient(ConsistencyLevel.SESSION, false, null, true, true)} }; } @@ -906,25 +906,33 @@ static List parseProtocols(String protocols) { @DataProvider public static Object[][] simpleClientBuildersWithDirect() { - return simpleClientBuildersWithDirect(true, toArray(protocols)); + return simpleClientBuildersWithDirect(true, true, toArray(protocols)); } @DataProvider public static Object[][] simpleClientBuildersWithDirectHttps() { - return simpleClientBuildersWithDirect(true, Protocol.HTTPS); + return simpleClientBuildersWithDirect(true, true, Protocol.HTTPS); } @DataProvider public static Object[][] simpleClientBuildersWithDirectTcp() { - return simpleClientBuildersWithDirect(true, Protocol.TCP); + return simpleClientBuildersWithDirect(true, true, Protocol.TCP); } @DataProvider public static Object[][] simpleClientBuildersWithDirectTcpWithContentResponseOnWriteDisabled() { - return simpleClientBuildersWithDirect(false, Protocol.TCP); + return simpleClientBuildersWithDirect(false, true, Protocol.TCP); } - private static Object[][] simpleClientBuildersWithDirect(boolean contentResponseOnWriteEnabled, Protocol... protocols) { + @DataProvider + public static Object[][] simpleClientBuildersForDirectTcpWithoutRetryOnThrottledRequests() { + return new Object[][]{ + {createDirectRxDocumentClient(ConsistencyLevel.SESSION, Protocol.HTTPS, false, null, true, false)}, + {createDirectRxDocumentClient(ConsistencyLevel.SESSION, Protocol.TCP, false, null, true, false)}, + }; + } + + private static Object[][] simpleClientBuildersWithDirect(boolean contentResponseOnWriteEnabled, boolean retryOnThrottledRequests, Protocol... protocols) { logger.info("Max test consistency to use is [{}]", accountConsistency); List testConsistencies = ImmutableList.of(ConsistencyLevel.EVENTUAL); @@ -938,7 +946,8 @@ private static Object[][] simpleClientBuildersWithDirect(boolean contentResponse protocol, isMultiMasterEnabled, preferredLocations, - contentResponseOnWriteEnabled))); + contentResponseOnWriteEnabled, + retryOnThrottledRequests))); } cosmosConfigurations.forEach(c -> { @@ -951,37 +960,37 @@ private static Object[][] simpleClientBuildersWithDirect(boolean contentResponse ); }); - cosmosConfigurations.add(createGatewayRxDocumentClient(ConsistencyLevel.SESSION, false, null, contentResponseOnWriteEnabled)); + cosmosConfigurations.add(createGatewayRxDocumentClient(ConsistencyLevel.SESSION, false, null, contentResponseOnWriteEnabled, retryOnThrottledRequests)); return cosmosConfigurations.stream().map(b -> new Object[]{b}).collect(Collectors.toList()).toArray(new Object[0][]); } @DataProvider public static Object[][] clientBuildersWithDirect() { - return clientBuildersWithDirectAllConsistencies(true, toArray(protocols)); + return clientBuildersWithDirectAllConsistencies(true, true, toArray(protocols)); } @DataProvider public static Object[][] clientBuildersWithDirectHttps() { - return clientBuildersWithDirectAllConsistencies(true, Protocol.HTTPS); + return clientBuildersWithDirectAllConsistencies(true, true, Protocol.HTTPS); } @DataProvider public static Object[][] clientBuildersWithDirectTcp() { - return clientBuildersWithDirectAllConsistencies(true, Protocol.TCP); + return clientBuildersWithDirectAllConsistencies(true, true, Protocol.TCP); } @DataProvider public static Object[][] clientBuildersWithDirectTcpWithContentResponseOnWriteDisabled() { - return clientBuildersWithDirectAllConsistencies(false, Protocol.TCP); + return clientBuildersWithDirectAllConsistencies(false, true, Protocol.TCP); } @DataProvider public static Object[][] clientBuildersWithContentResponseOnWriteEnabledAndDisabled() { Object[][] clientBuildersWithDisabledContentResponseOnWrite = - clientBuildersWithDirectSession(false, Protocol.TCP); + clientBuildersWithDirectSession(false, true, Protocol.TCP); Object[][] clientBuildersWithEnabledContentResponseOnWrite = - clientBuildersWithDirectSession(true, Protocol.TCP); + clientBuildersWithDirectSession(true, true, Protocol.TCP); int length = clientBuildersWithDisabledContentResponseOnWrite.length + clientBuildersWithEnabledContentResponseOnWrite.length; Object[][] clientBuilders = new Object[length][]; @@ -997,27 +1006,27 @@ public static Object[][] clientBuildersWithContentResponseOnWriteEnabledAndDisab @DataProvider public static Object[][] clientBuildersWithDirectSession() { - return clientBuildersWithDirectSession(true, toArray(protocols)); + return clientBuildersWithDirectSession(true, true, toArray(protocols)); } @DataProvider public static Object[][] simpleClientBuilderGatewaySession() { - return clientBuildersWithDirectSession(true); + return clientBuildersWithDirectSession(true, true); } static Protocol[] toArray(List protocols) { return protocols.toArray(new Protocol[protocols.size()]); } - private static Object[][] clientBuildersWithDirectSession(boolean contentResponseOnWriteEnabled, Protocol... protocols) { + private static Object[][] clientBuildersWithDirectSession(boolean contentResponseOnWriteEnabled, boolean retryOnThrottledRequests, Protocol... protocols) { return clientBuildersWithDirect(new ArrayList() {{ add(ConsistencyLevel.SESSION); - }}, contentResponseOnWriteEnabled, protocols); + }}, contentResponseOnWriteEnabled, retryOnThrottledRequests, protocols); } - private static Object[][] clientBuildersWithDirectAllConsistencies(boolean contentResponseOnWriteEnabled, Protocol... protocols) { + private static Object[][] clientBuildersWithDirectAllConsistencies(boolean contentResponseOnWriteEnabled, boolean retryOnThrottledRequests, Protocol... protocols) { logger.info("Max test consistency to use is [{}]", accountConsistency); - return clientBuildersWithDirect(desiredConsistencies, contentResponseOnWriteEnabled, protocols); + return clientBuildersWithDirect(desiredConsistencies, contentResponseOnWriteEnabled, retryOnThrottledRequests); } static List parseDesiredConsistencies(String consistencies) { @@ -1059,7 +1068,11 @@ static List allEqualOrLowerConsistencies(ConsistencyLevel acco return testConsistencies; } - private static Object[][] clientBuildersWithDirect(List testConsistencies, boolean contentResponseOnWriteEnabled, Protocol... protocols) { + private static Object[][] clientBuildersWithDirect( + List testConsistencies, + boolean contentResponseOnWriteEnabled, + boolean retryOnThrottledRequests, + Protocol... protocols) { boolean isMultiMasterEnabled = preferredLocations != null && accountConsistency == ConsistencyLevel.SESSION; List cosmosConfigurations = new ArrayList<>(); @@ -1069,7 +1082,8 @@ private static Object[][] clientBuildersWithDirect(List testCo protocol, isMultiMasterEnabled, preferredLocations, - contentResponseOnWriteEnabled))); + contentResponseOnWriteEnabled, + retryOnThrottledRequests))); } cosmosConfigurations.forEach(c -> { @@ -1082,7 +1096,13 @@ private static Object[][] clientBuildersWithDirect(List testCo ); }); - cosmosConfigurations.add(createGatewayRxDocumentClient(ConsistencyLevel.SESSION, isMultiMasterEnabled, preferredLocations, contentResponseOnWriteEnabled)); + cosmosConfigurations.add( + createGatewayRxDocumentClient( + ConsistencyLevel.SESSION, + isMultiMasterEnabled, + preferredLocations, + contentResponseOnWriteEnabled, + retryOnThrottledRequests)); return cosmosConfigurations.stream().map(c -> new Object[]{c}).collect(Collectors.toList()).toArray(new Object[0][]); } @@ -1099,32 +1119,45 @@ static protected CosmosClientBuilder createGatewayHouseKeepingDocumentClient(boo .consistencyLevel(ConsistencyLevel.SESSION); } - static protected CosmosClientBuilder createGatewayRxDocumentClient(ConsistencyLevel consistencyLevel, boolean multiMasterEnabled, - List preferredRegions, boolean contentResponseOnWriteEnabled) { + static protected CosmosClientBuilder createGatewayRxDocumentClient( + ConsistencyLevel consistencyLevel, + boolean multiMasterEnabled, + List preferredRegions, + boolean contentResponseOnWriteEnabled, + boolean retryOnThrottledRequests) { + GatewayConnectionConfig gatewayConnectionConfig = new GatewayConnectionConfig(); - return new CosmosClientBuilder().endpoint(TestConfigurations.HOST) - .credential(credential) - .gatewayMode(gatewayConnectionConfig) - .multipleWriteRegionsEnabled(multiMasterEnabled) - .preferredRegions(preferredRegions) - .contentResponseOnWriteEnabled(contentResponseOnWriteEnabled) - .consistencyLevel(consistencyLevel); + CosmosClientBuilder builder = new CosmosClientBuilder().endpoint(TestConfigurations.HOST) + .credential(credential) + .gatewayMode(gatewayConnectionConfig) + .multipleWriteRegionsEnabled(multiMasterEnabled) + .preferredRegions(preferredRegions) + .contentResponseOnWriteEnabled(contentResponseOnWriteEnabled) + .consistencyLevel(consistencyLevel); + + if (!retryOnThrottledRequests) { + builder.throttlingRetryOptions(new ThrottlingRetryOptions().setMaxRetryAttemptsOnThrottledRequests(0)); + } + + return builder; } static protected CosmosClientBuilder createGatewayRxDocumentClient() { - return createGatewayRxDocumentClient(ConsistencyLevel.SESSION, false, null, true); + return createGatewayRxDocumentClient(ConsistencyLevel.SESSION, false, null, true, true); } static protected CosmosClientBuilder createDirectRxDocumentClient(ConsistencyLevel consistencyLevel, Protocol protocol, boolean multiMasterEnabled, List preferredRegions, - boolean contentResponseOnWriteEnabled) { + boolean contentResponseOnWriteEnabled, + boolean retryOnThrottledRequests) { CosmosClientBuilder builder = new CosmosClientBuilder().endpoint(TestConfigurations.HOST) .credential(credential) .directMode(DirectConnectionConfig.getDefaultConfig()) .contentResponseOnWriteEnabled(contentResponseOnWriteEnabled) .consistencyLevel(consistencyLevel); + if (preferredRegions != null) { builder.preferredRegions(preferredRegions); } @@ -1133,6 +1166,10 @@ static protected CosmosClientBuilder createDirectRxDocumentClient(ConsistencyLev builder.multipleWriteRegionsEnabled(true); } + if (!retryOnThrottledRequests) { + builder.throttlingRetryOptions(new ThrottlingRetryOptions().setMaxRetryAttemptsOnThrottledRequests(0)); + } + Configs configs = spy(new Configs()); doAnswer((Answer)invocation -> protocol).when(configs).getProtocol();