diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ClosedClientTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ClosedClientTests.java index 3949f6cafcc3..ade9497e76b1 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ClosedClientTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ClosedClientTests.java @@ -9,7 +9,6 @@ import com.azure.cosmos.implementation.OperationType; import com.azure.cosmos.implementation.Utils; import com.azure.cosmos.implementation.directconnectivity.rntbd.ClosedClientTransportException; -import com.azure.cosmos.implementation.throughputControl.TestItem; import com.azure.cosmos.models.CosmosItemRequestOptions; import com.azure.cosmos.models.CosmosPatchOperations; import com.azure.cosmos.models.CosmosQueryRequestOptions; @@ -95,7 +94,7 @@ public void failFastWhenUsingClosedClient(OperationType operationType) { safeClose(asyncClient); - performDocumentOperation(cosmosAsyncContainer, TestItem.createNewItem(), operationType); + performDocumentOperation(cosmosAsyncContainer, TestObject.create(), operationType); fail("Operation is expected to fail!"); } catch (Exception ex) { @@ -121,7 +120,7 @@ public void failFastWhenUsingClosedClient(OperationType operationType) { private void performDocumentOperation( CosmosAsyncContainer asyncContainer, - TestItem testItem, + TestObject testItem, OperationType operationType) { if (operationType == OperationType.Query) { @@ -131,7 +130,7 @@ private void performDocumentOperation( SqlQuerySpec sqlQuerySpec = new SqlQuerySpec(query); asyncContainer - .queryItems(sqlQuerySpec, queryRequestOptions, TestItem.class) + .queryItems(sqlQuerySpec, queryRequestOptions, TestObject.class) .byPage() .blockFirst(); } @@ -139,7 +138,7 @@ private void performDocumentOperation( if (operationType == OperationType.Read) { asyncContainer - .readItem(testItem.getId(), new PartitionKey(testItem.getId()), TestItem.class) + .readItem(testItem.getId(), new PartitionKey(testItem.getId()), TestObject.class) .block(); } @@ -181,7 +180,7 @@ private void performDocumentOperation( CosmosPatchOperations patchOperations = CosmosPatchOperations.create().add("/" + "newProperty", "newVal"); asyncContainer - .patchItem(testItem.getId(), new PartitionKey(testItem.getId()), patchOperations, TestItem.class) + .patchItem(testItem.getId(), new PartitionKey(testItem.getId()), patchOperations, TestObject.class) .block(); } } diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/EndToEndTimeOutWithAvailabilityTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/EndToEndTimeOutWithAvailabilityTest.java index b73bf1d39759..5d82877afe68 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/EndToEndTimeOutWithAvailabilityTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/EndToEndTimeOutWithAvailabilityTest.java @@ -9,7 +9,6 @@ import com.azure.cosmos.implementation.OperationType; import com.azure.cosmos.implementation.RxDocumentClientImpl; import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils; -import com.azure.cosmos.implementation.throughputControl.TestItem; import com.azure.cosmos.models.CosmosItemIdentity; import com.azure.cosmos.models.CosmosItemRequestOptions; import com.azure.cosmos.models.CosmosItemResponse; @@ -107,7 +106,7 @@ public void testThresholdAvailabilityStrategy( FaultInjectionConnectionType.DIRECT : FaultInjectionConnectionType.GATEWAY; - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); CosmosItemRequestOptions options = new CosmosItemRequestOptions(); this.cosmosAsyncContainer.createItem(createdItem).block(); @@ -185,7 +184,7 @@ public void testThresholdAvailabilityStrategyForReadsDefaultEnablementWithPpaf( CosmosAsyncDatabase asyncDatabase = cosmosAsyncClient.getDatabase(this.cosmosAsyncContainer.getDatabase().getId()); CosmosAsyncContainer asyncContainer = asyncDatabase.getContainer(this.cosmosAsyncContainer.getId()); - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); CosmosItemRequestOptions options = new CosmosItemRequestOptions(); asyncContainer.createItem(createdItem).block(); @@ -306,7 +305,7 @@ public void afterClass() { private CosmosDiagnostics performDocumentOperation( CosmosAsyncContainer cosmosAsyncContainer, OperationType operationType, - TestItem createdItem, + TestObject createdItem, CosmosItemRequestOptions cosmosItemRequestOptions, boolean ignoreE2E2LatencyCfgOnRequestOptions, QueryFlavor queryFlavor) { @@ -330,8 +329,8 @@ private CosmosDiagnostics performDocumentOperation( logger.info("Running readAllItems..."); - FeedResponse response = cosmosAsyncContainer - .readAllItems(queryRequestOptions, TestItem.class) + FeedResponse response = cosmosAsyncContainer + .readAllItems(queryRequestOptions, TestObject.class) .byPage() .blockFirst(); @@ -347,11 +346,11 @@ private CosmosDiagnostics performDocumentOperation( logger.info("Running readMany..."); - FeedResponse response = cosmosAsyncContainer + FeedResponse response = cosmosAsyncContainer .readMany( Arrays.asList(new CosmosItemIdentity(new PartitionKey(createdItem.getMypk()), createdItem.getId())), readManyRequestOptions, - TestItem.class) + TestObject.class) .block(); assertThat(response).isNotNull(); @@ -361,8 +360,8 @@ private CosmosDiagnostics performDocumentOperation( logger.info("Running query ..."); String query = String.format("SELECT * from c where c.id = '%s'", createdItem.getId()); - FeedResponse itemFeedResponse = - cosmosAsyncContainer.queryItems(query, queryRequestOptions, TestItem.class).byPage().blockFirst(); + FeedResponse itemFeedResponse = + cosmosAsyncContainer.queryItems(query, queryRequestOptions, TestObject.class).byPage().blockFirst(); assertThat(itemFeedResponse).isNotNull(); @@ -378,11 +377,11 @@ private CosmosDiagnostics performDocumentOperation( if (operationType == OperationType.Read) { - CosmosItemResponse response = cosmosAsyncContainer.readItem( + CosmosItemResponse response = cosmosAsyncContainer.readItem( createdItem.getId(), new PartitionKey(createdItem.getMypk()), cosmosItemRequestOptions, - TestItem.class) + TestObject.class) .block(); assertThat(response).isNotNull(); @@ -391,7 +390,7 @@ private CosmosDiagnostics performDocumentOperation( } if (operationType == OperationType.Replace) { - CosmosItemResponse response = cosmosAsyncContainer.replaceItem( + CosmosItemResponse response = cosmosAsyncContainer.replaceItem( createdItem, createdItem.getId(), new PartitionKey(createdItem.getMypk()), @@ -404,7 +403,7 @@ private CosmosDiagnostics performDocumentOperation( } if (operationType == OperationType.Delete) { - TestItem toBeDeletedItem = TestItem.createNewItem(); + TestObject toBeDeletedItem = TestObject.create(); cosmosAsyncContainer.createItem(toBeDeletedItem, cosmosItemRequestOptions).block(); CosmosItemResponse response = cosmosAsyncContainer .deleteItem(toBeDeletedItem, cosmosItemRequestOptions) @@ -416,8 +415,8 @@ private CosmosDiagnostics performDocumentOperation( } if (operationType == OperationType.Create) { - CosmosItemResponse response = cosmosAsyncContainer - .createItem(TestItem.createNewItem(), cosmosItemRequestOptions) + CosmosItemResponse response = cosmosAsyncContainer + .createItem(TestObject.create(), cosmosItemRequestOptions) .block(); assertThat(response).isNotNull(); @@ -426,8 +425,8 @@ private CosmosDiagnostics performDocumentOperation( } if (operationType == OperationType.Upsert) { - CosmosItemResponse response = cosmosAsyncContainer - .upsertItem(TestItem.createNewItem(), cosmosItemRequestOptions) + CosmosItemResponse response = cosmosAsyncContainer + .upsertItem(TestObject.create(), cosmosItemRequestOptions) .block(); assertThat(response).isNotNull(); @@ -444,8 +443,8 @@ private CosmosDiagnostics performDocumentOperation( CosmosPatchItemRequestOptions patchItemRequestOptions = new CosmosPatchItemRequestOptions(); patchItemRequestOptions.setNonIdempotentWriteRetryPolicy(true, true); patchItemRequestOptions.setCosmosEndToEndOperationLatencyPolicyConfig(config); - CosmosItemResponse response = cosmosAsyncContainer - .patchItem(createdItem.getId(), new PartitionKey(createdItem.getMypk()), patchOperations, patchItemRequestOptions, TestItem.class) + CosmosItemResponse response = cosmosAsyncContainer + .patchItem(createdItem.getId(), new PartitionKey(createdItem.getMypk()), patchOperations, patchItemRequestOptions, TestObject.class) .block(); assertThat(response).isNotNull(); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ExcludeRegionTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ExcludeRegionTests.java index 03aa0deb98e4..e0e91464d82d 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ExcludeRegionTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ExcludeRegionTests.java @@ -10,7 +10,6 @@ import com.azure.cosmos.implementation.OperationType; import com.azure.cosmos.implementation.RxDocumentClientImpl; import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils; -import com.azure.cosmos.implementation.throughputControl.TestItem; import com.azure.cosmos.models.CosmosBatch; import com.azure.cosmos.models.CosmosBatchRequestOptions; import com.azure.cosmos.models.CosmosBatchResponse; @@ -122,7 +121,7 @@ public void excludeRegionTest_SkipFirstPreferredRegion(OperationType operationTy throw new SkipException("excludeRegionTest_SkipFirstPreferredRegion can only be tested for multi-master with multi-regions"); } - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); this.cosmosAsyncContainer.createItem(createdItem).block(); Thread.sleep(1000); @@ -157,7 +156,7 @@ public void excludeRegionTest_readSessionNotAvailable( throw new SkipException("excludeRegionTest_SkipFirstPreferredRegion can only be tested for multi-master with multi-regions"); } - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); this.cosmosAsyncContainer.createItem(createdItem).block(); FaultInjectionRule serverErrorRule = new FaultInjectionRuleBuilder("excludeRegionTest-" + operationType) @@ -229,7 +228,7 @@ private List getPreferredRegionList(CosmosAsyncClient client) { private CosmosDiagnosticsContext performDocumentOperation( CosmosAsyncContainer cosmosAsyncContainer, OperationType operationType, - TestItem createdItem, + TestObject createdItem, List excludeRegions, CosmosEndToEndOperationLatencyPolicyConfig cosmosEndToEndOperationLatencyPolicyConfig) { @@ -240,8 +239,8 @@ private CosmosDiagnosticsContext performDocumentOperation( String query = String.format("SELECT * from c where c.id = '%s'", createdItem.getId()); queryRequestOptions.setExcludedRegions(excludeRegions); - FeedResponse itemFeedResponse = - cosmosAsyncContainer.queryItems(query, queryRequestOptions, TestItem.class).byPage().blockFirst(); + FeedResponse itemFeedResponse = + cosmosAsyncContainer.queryItems(query, queryRequestOptions, TestObject.class).byPage().blockFirst(); assertThat(itemFeedResponse).isNotNull(); CosmosDiagnostics cosmosDiagnostics = itemFeedResponse.getCosmosDiagnostics(); @@ -271,11 +270,11 @@ private CosmosDiagnosticsContext performDocumentOperation( throw new RuntimeException(e); } - CosmosItemResponse itemResponse = cosmosAsyncContainer.readItem( + CosmosItemResponse itemResponse = cosmosAsyncContainer.readItem( createdItem.getId(), new PartitionKey(createdItem.getMypk()), cosmosItemRequestOptions, - TestItem.class).block(); + TestObject.class).block(); assertThat(itemResponse).isNotNull(); assertThat(itemResponse.getDiagnostics()).isNotNull(); @@ -295,7 +294,7 @@ private CosmosDiagnosticsContext performDocumentOperation( throw new RuntimeException(e); } - CosmosItemResponse itemResponse = cosmosAsyncContainer.replaceItem( + CosmosItemResponse itemResponse = cosmosAsyncContainer.replaceItem( createdItem, createdItem.getId(), new PartitionKey(createdItem.getMypk()), @@ -313,7 +312,7 @@ private CosmosDiagnosticsContext performDocumentOperation( if (operationType == OperationType.Delete) { - TestItem itemToBeDeleted = TestItem.createNewItem(); + TestObject itemToBeDeleted = TestObject.create(); cosmosAsyncContainer.createItem(itemToBeDeleted, cosmosItemRequestOptions).block(); @@ -337,8 +336,8 @@ private CosmosDiagnosticsContext performDocumentOperation( } if (operationType == OperationType.Create) { - CosmosItemResponse itemResponse = cosmosAsyncContainer - .createItem(TestItem.createNewItem(), cosmosItemRequestOptions).block(); + CosmosItemResponse itemResponse = cosmosAsyncContainer + .createItem(TestObject.create(), cosmosItemRequestOptions).block(); assertThat(itemResponse).isNotNull(); assertThat(itemResponse.getDiagnostics()).isNotNull(); @@ -351,8 +350,8 @@ private CosmosDiagnosticsContext performDocumentOperation( } if (operationType == OperationType.Upsert) { - CosmosItemResponse itemResponse - = cosmosAsyncContainer.upsertItem(TestItem.createNewItem(), cosmosItemRequestOptions).block(); + CosmosItemResponse itemResponse + = cosmosAsyncContainer.upsertItem(TestObject.create(), cosmosItemRequestOptions).block(); assertThat(itemResponse).isNotNull(); assertThat(itemResponse.getDiagnostics()).isNotNull(); @@ -375,8 +374,8 @@ private CosmosDiagnosticsContext performDocumentOperation( patchItemRequestOptions.setCosmosEndToEndOperationLatencyPolicyConfig(cosmosEndToEndOperationLatencyPolicyConfig); patchItemRequestOptions.setExcludedRegions(excludeRegions); - CosmosItemResponse itemResponse = cosmosAsyncContainer - .patchItem(createdItem.getId(), new PartitionKey(createdItem.getMypk()), patchOperations, patchItemRequestOptions, TestItem.class) + CosmosItemResponse itemResponse = cosmosAsyncContainer + .patchItem(createdItem.getId(), new PartitionKey(createdItem.getMypk()), patchOperations, patchItemRequestOptions, TestObject.class) .block(); assertThat(itemResponse).isNotNull(); @@ -396,7 +395,7 @@ private CosmosDiagnosticsContext performDocumentOperation( cosmosBatchRequestOptions.setExcludedRegions(excludeRegions); - TestItem testItem = TestItem.createNewItem(); + TestObject testItem = TestObject.create(); PartitionKey partitionKey = new PartitionKey(testItem.getMypk()); CosmosBatch cosmosBatch = CosmosBatch.createCosmosBatch(partitionKey); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PerPartitionAutomaticFailoverE2ETests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PerPartitionAutomaticFailoverE2ETests.java index 4bd5006e53ba..e53658ae0dde 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PerPartitionAutomaticFailoverE2ETests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PerPartitionAutomaticFailoverE2ETests.java @@ -37,7 +37,6 @@ import com.azure.cosmos.implementation.http.HttpResponse; import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper; import com.azure.cosmos.implementation.routing.RegionalRoutingContext; -import com.azure.cosmos.implementation.throughputControl.TestItem; import com.azure.cosmos.models.CosmosBatch; import com.azure.cosmos.models.CosmosBatchResponse; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; @@ -1015,7 +1014,7 @@ public void testPpafWithWriteFailoverWithEligibleErrorStatusCodes( regionalRoutingContextWithIssues, cosmosException); - TestItem testItem = TestItem.createNewItem(); + TestObject testItem = TestObject.create(); Function> dataPlaneOperation = resolveDataPlaneOperation(operationType); @@ -1109,7 +1108,7 @@ public void testPpafWithWriteFailoverWithEligibleErrorStatusCodes( shouldThrowReadTimeoutExceptionWhenNetworkError, shouldUseE2ETimeout); - TestItem testItem = TestItem.createNewItem(); + TestObject testItem = TestObject.create(); Function> dataPlaneOperation = resolveDataPlaneOperation(operationType); @@ -1378,7 +1377,7 @@ private Function> resolveDa return (paramsWrapper) -> { CosmosAsyncContainer asyncContainer = paramsWrapper.asyncContainer; - TestItem createdTestObject = paramsWrapper.createdTestItem; + TestObject createdTestObject = paramsWrapper.createdTestItem; CosmosItemRequestOptions itemRequestOptions = paramsWrapper.itemRequestOptions; try { @@ -1405,12 +1404,12 @@ private Function> resolveDa return (paramsWrapper) -> { CosmosAsyncContainer asyncContainer = paramsWrapper.asyncContainer; - TestItem createdTestObject = paramsWrapper.createdTestItem; + TestObject createdTestObject = paramsWrapper.createdTestItem; CosmosItemRequestOptions itemRequestOptions = paramsWrapper.itemRequestOptions; try { - CosmosItemResponse upsertItemResponse = asyncContainer.upsertItem( + CosmosItemResponse upsertItemResponse = asyncContainer.upsertItem( createdTestObject, new PartitionKey(createdTestObject.getId()), itemRequestOptions) @@ -1431,12 +1430,12 @@ private Function> resolveDa return (paramsWrapper) -> { CosmosAsyncContainer asyncContainer = paramsWrapper.asyncContainer; - TestItem createdTestObject = TestItem.createNewItem(); + TestObject createdTestObject = TestObject.create(); CosmosItemRequestOptions itemRequestOptions = paramsWrapper.itemRequestOptions; try { - CosmosItemResponse createItemResponse = asyncContainer.createItem( + CosmosItemResponse createItemResponse = asyncContainer.createItem( createdTestObject, new PartitionKey(createdTestObject.getId()), itemRequestOptions) @@ -1457,7 +1456,7 @@ private Function> resolveDa return (paramsWrapper) -> { CosmosAsyncContainer asyncContainer = paramsWrapper.asyncContainer; - TestItem createdTestObject = paramsWrapper.createdTestItem; + TestObject createdTestObject = paramsWrapper.createdTestItem; CosmosItemRequestOptions itemRequestOptions = paramsWrapper.itemRequestOptions; try { @@ -1483,19 +1482,19 @@ private Function> resolveDa return (paramsWrapper) -> { CosmosAsyncContainer asyncContainer = paramsWrapper.asyncContainer; - TestItem createdTestObject = paramsWrapper.createdTestItem; + TestObject createdTestObject = paramsWrapper.createdTestItem; CosmosPatchItemRequestOptions patchItemRequestOptions = (CosmosPatchItemRequestOptions) paramsWrapper.patchItemRequestOptions; CosmosPatchOperations patchOperations = CosmosPatchOperations.create().add("/number", 555); try { - CosmosItemResponse patchItemResponse = asyncContainer.patchItem( + CosmosItemResponse patchItemResponse = asyncContainer.patchItem( createdTestObject.getId(), new PartitionKey(createdTestObject.getId()), patchOperations, patchItemRequestOptions, - TestItem.class) + TestObject.class) .block(); return new ResponseWrapper<>(patchItemResponse); @@ -1540,12 +1539,12 @@ private Function> resolveDa return (paramsWrapper) -> { CosmosAsyncContainer asyncContainer = paramsWrapper.asyncContainer; - TestItem createdTestObject = paramsWrapper.createdTestItem; + TestObject createdTestObject = paramsWrapper.createdTestItem; CosmosItemRequestOptions itemRequestOptions = paramsWrapper.itemRequestOptions; try { - CosmosItemResponse deleteItemResponse = asyncContainer.replaceItem( + CosmosItemResponse deleteItemResponse = asyncContainer.replaceItem( createdTestObject, createdTestObject.getId(), new PartitionKey(createdTestObject.getId()), @@ -1566,7 +1565,7 @@ private Function> resolveDa case Batch: return (paramsWrapper) -> { - TestItem testObject = TestItem.createNewItem(); + TestObject testObject = TestObject.create(); CosmosBatch batch = CosmosBatch.createCosmosBatch(new PartitionKey(testObject.getId())); CosmosAsyncContainer asyncContainer = paramsWrapper.asyncContainer; @@ -1591,9 +1590,9 @@ private Function> resolveDa try { - FeedResponse feedResponseFromChangeFeed = asyncContainer.queryChangeFeed( + FeedResponse feedResponseFromChangeFeed = asyncContainer.queryChangeFeed( CosmosChangeFeedRequestOptions.createForProcessingFromBeginning(paramsWrapper.feedRangeToDrainForChangeFeed == null ? FeedRange.forFullRange() : paramsWrapper.feedRangeToDrainForChangeFeed), - TestItem.class) + TestObject.class) .byPage() .blockLast(); @@ -1651,7 +1650,7 @@ private static class ResponseWrapper { private static class OperationInvocationParamsWrapper { public CosmosAsyncContainer asyncContainer; - public TestItem createdTestItem; + public TestObject createdTestItem; public CosmosItemRequestOptions itemRequestOptions; public CosmosQueryRequestOptions queryRequestOptions; public CosmosItemRequestOptions patchItemRequestOptions; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/SessionConsistencyWithRegionScopingTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/SessionConsistencyWithRegionScopingTests.java index 1b913ccf6acd..0ac72022717b 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/SessionConsistencyWithRegionScopingTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/SessionConsistencyWithRegionScopingTests.java @@ -24,7 +24,6 @@ import com.azure.cosmos.implementation.guava25.hash.Funnels; import com.azure.cosmos.implementation.guava25.hash.PrimitiveSink; import com.azure.cosmos.implementation.routing.PartitionKeyInternal; -import com.azure.cosmos.implementation.throughputControl.TestItem; import com.azure.cosmos.models.CosmosBatch; import com.azure.cosmos.models.CosmosBatchOperationResult; import com.azure.cosmos.models.CosmosBatchRequestOptions; @@ -62,6 +61,7 @@ import reactor.core.publisher.Mono; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; @@ -322,7 +322,7 @@ public Object[][] readYouWriteWithNoExplicitRegionSwitchingTestContext() { Flux createOperationsFlux = Flux.range(0, createOperationCount).map(i -> { String documentId = UUID.randomUUID().toString(); - TestItem testItem = new TestItem(documentId, documentId, documentId); + TestObject testItem = new TestObject(documentId, documentId, Arrays.asList(), documentId); idsToCreate.add(documentId); @@ -353,7 +353,7 @@ public Object[][] readYouWriteWithNoExplicitRegionSwitchingTestContext() { bulkReadResponses.forEach(bulkReadResponse -> { assertThat(bulkReadResponse.getResponse()).isNotNull(); assertThat(bulkReadResponse.getResponse().getStatusCode()).isEqualTo(HttpConstants.StatusCodes.OK); - pksRead.add(bulkReadResponse.getResponse().getItem(TestItem.class).getMypk()); + pksRead.add(bulkReadResponse.getResponse().getItem(TestObject.class).getMypk()); }); return pksRead; @@ -366,7 +366,7 @@ public Object[][] readYouWriteWithNoExplicitRegionSwitchingTestContext() { Flux createOperationsFlux = Flux.range(0, createOperationCount).map(i -> { String documentId = UUID.randomUUID().toString(); - TestItem testItem = new TestItem(documentId, documentId, documentId); + TestObject testItem = new TestObject(documentId, documentId, Arrays.asList(), documentId); idsAddedByBulkCreate.add(documentId); return CosmosBulkOperations.getCreateItemOperation(testItem, new PartitionKey(documentId)); @@ -444,7 +444,7 @@ public Object[][] readYouWriteWithNoExplicitRegionSwitchingTestContext() { Flux createOperationsFlux = Flux.range(0, createOperationCount).map(i -> { String documentId = UUID.randomUUID().toString(); - TestItem testItem = new TestItem(documentId, documentId, documentId); + TestObject testItem = new TestObject(documentId, documentId, Arrays.asList(), documentId); idsAddedByBulkCreate.add(documentId); return CosmosBulkOperations.getCreateItemOperation(testItem, new PartitionKey(documentId)); @@ -1191,7 +1191,7 @@ public Object[][] readYouWriteWithExplicitRegionSwitchingTestContext() { Flux createOperationsFlux = Flux.range(0, createOperationCount).map(i -> { String documentId = UUID.randomUUID().toString(); - TestItem testItem = new TestItem(documentId, documentId, documentId); + TestObject testItem = new TestObject(documentId, documentId, Arrays.asList(), documentId); idsToCreate.add(documentId); @@ -1247,7 +1247,7 @@ public Object[][] readYouWriteWithExplicitRegionSwitchingTestContext() { Flux createOperationsFlux = Flux.range(0, createOperationCount).map(i -> { String documentId = UUID.randomUUID().toString(); - TestItem testItem = new TestItem(documentId, documentId, documentId); + TestObject testItem = new TestObject(documentId, documentId, Arrays.asList(), documentId); idsAddedByCreates.add(documentId); return CosmosBulkOperations.getCreateItemOperation(testItem, new PartitionKey(documentId)); @@ -1320,7 +1320,7 @@ public Object[][] readYouWriteWithExplicitRegionSwitchingTestContext() { Flux createOperationsFlux = Flux.range(0, createOperationCount).map(i -> { String documentId = UUID.randomUUID().toString(); - TestItem testItem = new TestItem(documentId, documentId, documentId); + TestObject testItem = new TestObject(documentId, documentId, Arrays.asList(), documentId); idsAddedByBulkCreate.add(documentId); return CosmosBulkOperations.getCreateItemOperation(testItem, new PartitionKey(documentId)); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/ExcludedRegionWithFaultInjectionTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/ExcludedRegionWithFaultInjectionTests.java index 6ddd63647243..79bc5818cc92 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/ExcludedRegionWithFaultInjectionTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/ExcludedRegionWithFaultInjectionTests.java @@ -12,6 +12,7 @@ import com.azure.cosmos.CosmosExcludedRegions; import com.azure.cosmos.CosmosRegionSwitchHint; import com.azure.cosmos.SessionRetryOptionsBuilder; +import com.azure.cosmos.TestObject; import com.azure.cosmos.implementation.AsyncDocumentClient; import com.azure.cosmos.implementation.DatabaseAccount; import com.azure.cosmos.implementation.DatabaseAccountLocation; @@ -19,7 +20,6 @@ import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.TestConfigurations; import com.azure.cosmos.implementation.Utils; -import com.azure.cosmos.implementation.throughputControl.TestItem; import com.azure.cosmos.models.CosmosBatch; import com.azure.cosmos.models.CosmosBatchRequestOptions; import com.azure.cosmos.models.CosmosBatchResponse; @@ -150,15 +150,15 @@ public Object[][] regionExclusionReadAfterCreateTestConfigs() { Function> readItemCallback = (params) -> { - TestItem alreadyCreatedItem = params.createdItem; + TestObject alreadyCreatedItem = params.createdItem; try { - CosmosItemResponse response = params.cosmosAsyncContainer + CosmosItemResponse response = params.cosmosAsyncContainer .readItem( alreadyCreatedItem.getId(), new PartitionKey(alreadyCreatedItem.getId()), - TestItem.class) + TestObject.class) .block(); return new OperationExecutionResult<>(response); @@ -481,7 +481,7 @@ public Object[][] regionExclusionQueryAfterCreateTestConfigs() { Function> queryItemCallback = (params) -> { - TestItem alreadyCreatedItem = params.createdItem; + TestObject alreadyCreatedItem = params.createdItem; String query = String.format("SELECT * FROM c WHERE c.id = '%s'", alreadyCreatedItem.getId()); CosmosQueryRequestOptions queryRequestOptions = params.queryRequestOptionsForCallbackAfterMutation != null @@ -490,8 +490,8 @@ public Object[][] regionExclusionQueryAfterCreateTestConfigs() { SqlQuerySpec sqlQuerySpec = new SqlQuerySpec(query); try { - FeedResponse feedResponse = params.cosmosAsyncContainer - .queryItems(sqlQuerySpec, queryRequestOptions, TestItem.class) + FeedResponse feedResponse = params.cosmosAsyncContainer + .queryItems(sqlQuerySpec, queryRequestOptions, TestObject.class) .byPage() .blockFirst(); @@ -767,9 +767,9 @@ public Object[][] regionExclusionWriteAfterCreateTestConfigs() { itemRequestOptions.setNonIdempotentWriteRetryPolicy(params.nonIdempotentWriteRetriesEnabled, true); - CosmosItemResponse response = params.cosmosAsyncContainer + CosmosItemResponse response = params.cosmosAsyncContainer .createItem( - new TestItem(newDocumentId, newDocumentId, newDocumentId), + new TestObject(newDocumentId, newDocumentId, Arrays.asList(), newDocumentId), new PartitionKey(newDocumentId), itemRequestOptions) .block(); @@ -792,8 +792,8 @@ public Object[][] regionExclusionWriteAfterCreateTestConfigs() { Function> replaceItemCallback = (params) -> { - TestItem alreadyCreatedItem = params.createdItem; - alreadyCreatedItem.setProp(UUID.randomUUID().toString()); + TestObject alreadyCreatedItem = params.createdItem; + alreadyCreatedItem.setStringProp(UUID.randomUUID().toString()); CosmosItemRequestOptions itemRequestOptions = params.itemRequestOptionsForCallbackAfterMutation != null ? params.itemRequestOptionsForCallbackAfterMutation : new CosmosItemRequestOptions(); @@ -802,7 +802,7 @@ public Object[][] regionExclusionWriteAfterCreateTestConfigs() { try { - CosmosItemResponse response = params.cosmosAsyncContainer + CosmosItemResponse response = params.cosmosAsyncContainer .replaceItem( alreadyCreatedItem, alreadyCreatedItem.getId(), @@ -827,7 +827,7 @@ public Object[][] regionExclusionWriteAfterCreateTestConfigs() { Function> deleteItemCallback = (params) -> { - TestItem alreadyCreatedItem = params.createdItem; + TestObject alreadyCreatedItem = params.createdItem; CosmosItemRequestOptions itemRequestOptions = params.itemRequestOptionsForCallbackAfterMutation != null ? params.itemRequestOptionsForCallbackAfterMutation : new CosmosItemRequestOptions(); @@ -860,8 +860,8 @@ public Object[][] regionExclusionWriteAfterCreateTestConfigs() { Function> upsertExistingItemCallback = (params) -> { - TestItem alreadyCreatedItem = params.createdItem; - alreadyCreatedItem.setProp(UUID.randomUUID().toString()); + TestObject alreadyCreatedItem = params.createdItem; + alreadyCreatedItem.setStringProp(UUID.randomUUID().toString()); CosmosItemRequestOptions itemRequestOptions = params.itemRequestOptionsForCallbackAfterMutation != null ? params.itemRequestOptionsForCallbackAfterMutation : new CosmosItemRequestOptions(); @@ -870,7 +870,7 @@ public Object[][] regionExclusionWriteAfterCreateTestConfigs() { try { - CosmosItemResponse response = params.cosmosAsyncContainer.upsertItem( + CosmosItemResponse response = params.cosmosAsyncContainer.upsertItem( alreadyCreatedItem, new PartitionKey(alreadyCreatedItem.getId()), itemRequestOptions).block(); @@ -893,7 +893,7 @@ public Object[][] regionExclusionWriteAfterCreateTestConfigs() { Function> upsertNonExistingItemCallback = (params) -> { - TestItem newItem = TestItem.createNewItem(); + TestObject newItem = TestObject.create(); CosmosItemRequestOptions itemRequestOptions = params.itemRequestOptionsForCallbackAfterMutation != null ? params.itemRequestOptionsForCallbackAfterMutation : new CosmosItemRequestOptions(); @@ -902,7 +902,7 @@ public Object[][] regionExclusionWriteAfterCreateTestConfigs() { try { - CosmosItemResponse response = params.cosmosAsyncContainer.upsertItem( + CosmosItemResponse response = params.cosmosAsyncContainer.upsertItem( newItem, new PartitionKey(newItem.getId()), itemRequestOptions).block(); @@ -933,12 +933,12 @@ public Object[][] regionExclusionWriteAfterCreateTestConfigs() { patchItemRequestOptions.setNonIdempotentWriteRetryPolicy(params.nonIdempotentWriteRetriesEnabled, true); try { - CosmosItemResponse response = params.cosmosAsyncContainer.patchItem( + CosmosItemResponse response = params.cosmosAsyncContainer.patchItem( params.createdItem.getId(), new PartitionKey(params.createdItem.getId()), patchOperations, (CosmosPatchItemRequestOptions) patchItemRequestOptions, - TestItem.class + TestObject.class ).block(); return new OperationExecutionResult<>(response); @@ -1558,7 +1558,7 @@ public Object[][] regionExclusionBatchTestConfigs() { Function> batchCreateAndReadCallback = (params) -> { String documentId = UUID.randomUUID().toString(); - TestItem testItem = new TestItem(documentId, documentId, documentId); + TestObject testItem = new TestObject(documentId, documentId, Arrays.asList(), documentId); CosmosBatch batch = CosmosBatch.createCosmosBatch(new PartitionKey(documentId)); @@ -1933,7 +1933,7 @@ public Object[][] regionExclusionBulkTestConfigs() { Flux createOperationsFlux = Flux.range(0, totalCreates).map(i -> { String documentId = UUID.randomUUID().toString(); - TestItem testItem = new TestItem(documentId, documentId, documentId); + TestObject testItem = new TestObject(documentId, documentId, Arrays.asList(), documentId); return CosmosBulkOperations.getCreateItemOperation(testItem, new PartitionKey(documentId)); @@ -2364,7 +2364,7 @@ private void execute(MutationTestConfig mutationTestConfig, boolean shouldInject .getDatabase(this.cosmosAsyncContainer.getDatabase().getId()) .getContainer(this.cosmosAsyncContainer.getId()); - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); Thread.sleep(5_000); @@ -2675,11 +2675,11 @@ private static AccountLevelLocationContext getAccountLevelLocationContext(Databa } private static String getRegionResolvedForDefaultEndpoint(CosmosAsyncContainer container, List preferredRegions) { - TestItem testItem = TestItem.createNewItem(); + TestObject testItem = TestObject.create(); CosmosItemRequestOptions itemRequestOptions = new CosmosItemRequestOptions(); itemRequestOptions.setExcludedRegions(preferredRegions); - CosmosItemResponse response = container.createItem(testItem, itemRequestOptions).block(); + CosmosItemResponse response = container.createItem(testItem, itemRequestOptions).block(); Set contactedRegion = response.getDiagnostics().getContactedRegionNames(); @@ -2902,7 +2902,7 @@ public MutationTestConfig withPerRegionDuplicateCount(int perRegionDuplicateCoun private static class ItemOperationInvocationParameters { public boolean nonIdempotentWriteRetriesEnabled; public CosmosAsyncContainer cosmosAsyncContainer; - public TestItem createdItem; + public TestObject createdItem; public CosmosItemRequestOptions itemRequestOptionsForCallbackAfterMutation; public CosmosItemRequestOptions patchItemRequestOptionsForCallbackAfterMutation; public CosmosQueryRequestOptions queryRequestOptionsForCallbackAfterMutation; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionConnectionErrorRuleTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionConnectionErrorRuleTests.java index c9694190c34f..b3eedd79ff17 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionConnectionErrorRuleTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionConnectionErrorRuleTests.java @@ -8,6 +8,7 @@ import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.CosmosClientBuilder; import com.azure.cosmos.CosmosContainerProactiveInitConfigBuilder; +import com.azure.cosmos.TestObject; import com.azure.cosmos.implementation.AsyncDocumentClient; import com.azure.cosmos.implementation.DatabaseAccount; import com.azure.cosmos.implementation.DatabaseAccountLocation; @@ -16,7 +17,6 @@ import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils; import com.azure.cosmos.implementation.directconnectivity.RntbdTransportClient; import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdEndpoint; -import com.azure.cosmos.implementation.throughputControl.TestItem; import com.azure.cosmos.models.CosmosContainerIdentity; import com.azure.cosmos.models.FeedRange; import com.azure.cosmos.models.PartitionKey; @@ -103,7 +103,7 @@ public void faultInjectionConnectionErrorRuleTestWithNoConnectionWarmup(FaultInj CosmosAsyncContainer singlePartitionContainer = getSharedSinglePartitionCosmosContainer(client); // validate one channel exists - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); singlePartitionContainer.createItem(createdItem).block(); RntbdTransportClient rntbdTransportClient = (RntbdTransportClient) ReflectionUtils.getTransportClient(this.client); @@ -146,7 +146,7 @@ public void faultInjectionConnectionErrorRuleTestWithNoConnectionWarmup(FaultInj assertThat(connectionErrorRule.getHitCountDetails()).isNull(); // do another request to open a new connection - singlePartitionContainer.createItem(TestItem.createNewItem()).block(); + singlePartitionContainer.createItem(TestObject.create()).block(); Thread.sleep(Duration.ofSeconds(2).toMillis()); // the configured connection rule should have disabled after 2s, so the connection will remain open @@ -200,7 +200,7 @@ public void faultInjectionConnectionErrorRuleTestWithConnectionWarmup(FaultInjec singlePartitionContainer = connectionWarmupClient.getDatabase(databaseId).getContainer(containerId); // validate one channel exists - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); singlePartitionContainer.createItem(createdItem).block(); RntbdTransportClient rntbdTransportClient = (RntbdTransportClient) ReflectionUtils.getTransportClient(connectionWarmupClient); @@ -269,7 +269,7 @@ public void faultInjectionConnectionErrorRuleTestWithConnectionWarmup(FaultInjec assertThat(connectionErrorRule.getHitCountDetails()).isNull(); // do another request to open a new connection - singlePartitionContainer.createItem(TestItem.createNewItem()).block(); + singlePartitionContainer.createItem(TestObject.create()).block(); Thread.sleep(Duration.ofSeconds(2).toMillis()); // the configured connection rule should have disabled after 2s, so the connection will remain open @@ -300,7 +300,7 @@ public void connectionCloseError_NoEndpoint_NoWarmup() { CosmosAsyncContainer singlePartitionContainer = getSharedSinglePartitionCosmosContainer(client); // validate one channel exists - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); singlePartitionContainer.createItem(createdItem).block(); RntbdTransportClient rntbdTransportClient = (RntbdTransportClient) ReflectionUtils.getTransportClient(this.client); @@ -339,7 +339,7 @@ public void connectionCloseError_NoEndpoint_NoWarmup() { assertThat(connectionErrorRule.getHitCountDetails()).isNull(); // do another request to open a new connection - singlePartitionContainer.createItem(TestItem.createNewItem()).block(); + singlePartitionContainer.createItem(TestObject.create()).block(); Thread.sleep(Duration.ofSeconds(2).toMillis()); // the configured connection rule should have disabled after 2s, so the connection will remain open diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionMetadataRequestRuleTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionMetadataRequestRuleTests.java index 505767816420..e7ed9ff15a69 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionMetadataRequestRuleTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionMetadataRequestRuleTests.java @@ -10,6 +10,7 @@ import com.azure.cosmos.CosmosDiagnostics; import com.azure.cosmos.CosmosEndToEndOperationLatencyPolicyConfigBuilder; import com.azure.cosmos.CosmosException; +import com.azure.cosmos.TestObject; import com.azure.cosmos.implementation.AsyncDocumentClient; import com.azure.cosmos.implementation.DatabaseAccount; import com.azure.cosmos.implementation.DatabaseAccountLocation; @@ -18,7 +19,6 @@ import com.azure.cosmos.implementation.MetadataDiagnosticsContext; import com.azure.cosmos.implementation.OperationType; import com.azure.cosmos.implementation.Utils; -import com.azure.cosmos.implementation.throughputControl.TestItem; import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.FeedRange; @@ -190,7 +190,7 @@ public void faultInjectionServerErrorRuleTests_AddressRefresh_ConnectionDelay(bo .duration(Duration.ofMinutes(5)) .build(); try { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); container.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules( @@ -295,7 +295,7 @@ public void faultInjectionServerErrorRuleTests_AddressRefresh_ResponseDelay( .build(); try { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); container.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules( @@ -358,7 +358,7 @@ public void faultInjectionServerErrorRuleTests_AddressRefresh_byPartition(boolea // first create few documents for (int i = 0; i < 10; i++) { - container.createItem(TestItem.createNewItem()).block(); + container.createItem(TestObject.create()).block(); } List feedRanges = container.getFeedRanges().block(); @@ -367,14 +367,14 @@ public void faultInjectionServerErrorRuleTests_AddressRefresh_byPartition(boolea CosmosQueryRequestOptions cosmosQueryRequestOptions = new CosmosQueryRequestOptions(); cosmosQueryRequestOptions.setFeedRange(feedRanges.get(0)); String query = "select * from c"; - TestItem itemOnFeedRange1 = container.queryItems(query, cosmosQueryRequestOptions, TestItem.class) + TestObject itemOnFeedRange1 = container.queryItems(query, cosmosQueryRequestOptions, TestObject.class) .byPage(1) .blockFirst() .getResults() .get(0); cosmosQueryRequestOptions.setFeedRange(feedRanges.get(1)); - TestItem itemOnFeedRange2 = container.queryItems(query, cosmosQueryRequestOptions, TestItem.class) + TestObject itemOnFeedRange2 = container.queryItems(query, cosmosQueryRequestOptions, TestObject.class) .byPage(1) .blockFirst() .getResults() @@ -509,7 +509,7 @@ public void faultInjectionServerErrorRuleTests_AddressRefresh_TooManyRequest(boo .duration(Duration.ofMinutes(5)) .build(); try { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); container.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules( @@ -589,7 +589,7 @@ public void faultInjectionServerErrorRuleTests_PartitionKeyRanges_DelayError( try { // create few items to first make sure the collection cache, pkRanges cache is being populated for (int i = 0; i < 10; i++) { - container.createItem(TestItem.createNewItem()).block(); + container.createItem(TestObject.create()).block(); } CosmosFaultInjectionHelper.configureFaultInjectionRules( @@ -598,7 +598,7 @@ public void faultInjectionServerErrorRuleTests_PartitionKeyRanges_DelayError( .block(); try { - CosmosDiagnostics cosmosDiagnostics = container.createItem(TestItem.createNewItem()).block().getDiagnostics(); + CosmosDiagnostics cosmosDiagnostics = container.createItem(TestObject.create()).block().getDiagnostics(); // The PkRanges requests may have retried in another region, // but the create request will only be retried locally for PARTITION_IS_SPLITTING assertThat(cosmosDiagnostics.getContactedRegionNames().size()).isEqualTo(1); @@ -695,7 +695,7 @@ public void faultInjectionServerErrorRuleTests_CollectionRead_ConnectionDelay(bo // issue few requests to first populate all the necessary caches // as connection delay will impact all other operations, and in this test, we want to limit the scope to only collection read for (int i = 0; i < 10; i++) { - container.createItem(TestItem.createNewItem()).block(); + container.createItem(TestObject.create()).block(); } CosmosFaultInjectionHelper.configureFaultInjectionRules( @@ -704,7 +704,7 @@ public void faultInjectionServerErrorRuleTests_CollectionRead_ConnectionDelay(bo .block(); try { - CosmosDiagnostics cosmosDiagnostics = container.createItem(TestItem.createNewItem()).block().getDiagnostics(); + CosmosDiagnostics cosmosDiagnostics = container.createItem(TestObject.create()).block().getDiagnostics(); // validate CONTAINER_LOOK_UP ObjectNode diagnosticsNode = (ObjectNode) Utils.getSimpleObjectMapper().readTree(cosmosDiagnostics.toString()); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java index ad92f2917932..079768788bfc 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java @@ -10,6 +10,7 @@ import com.azure.cosmos.CosmosClientBuilder; import com.azure.cosmos.CosmosDiagnostics; import com.azure.cosmos.DirectConnectionConfig; +import com.azure.cosmos.TestObject; import com.azure.cosmos.implementation.AsyncDocumentClient; import com.azure.cosmos.implementation.DatabaseAccount; import com.azure.cosmos.implementation.DatabaseAccountLocation; @@ -20,7 +21,6 @@ import com.azure.cosmos.implementation.ResourceType; import com.azure.cosmos.implementation.TestConfigurations; import com.azure.cosmos.implementation.Utils; -import com.azure.cosmos.implementation.throughputControl.TestItem; import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.FeedRange; @@ -234,7 +234,7 @@ public void faultInjectionServerErrorRuleTests_OperationType(OperationType opera .build(); try { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); cosmosAsyncContainer.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules(cosmosAsyncContainer, Arrays.asList(serverGoneErrorRule)).block(); @@ -282,7 +282,7 @@ public void faultInjectionServerErrorRuleTests_OperationType(OperationType opera @Test(groups = {"multi-region"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT) public void faultInjectionServerErrorRuleTests_OperationTypeImpactAddresses(OperationType operationType) throws JsonProcessingException { // Test the operation type can impact which region or replica the rule will be applicable - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); this.cosmosAsyncContainer.createItem(createdItem).block(); String writeRegionServerGoneRuleId = "serverErrorRule-writeRegionOnly-" + UUID.randomUUID(); @@ -441,7 +441,7 @@ public void faultInjectionServerErrorRuleTests_Region(boolean shouldInjectPrefer .getDatabase(this.cosmosAsyncContainer.getDatabase().getId()) .getContainer(this.cosmosAsyncContainer.getId()); - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); container.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules( @@ -485,7 +485,7 @@ public void faultInjectionServerErrorRuleTests_Region(boolean shouldInjectPrefer @Test(groups = {"multi-region", "long"}, timeOut = TIMEOUT) public void faultInjectionServerErrorRuleTests_Partition() throws JsonProcessingException { for (int i = 0; i < 10; i++) { - cosmosAsyncContainer.createItem(TestItem.createNewItem()).block(); + cosmosAsyncContainer.createItem(TestObject.create()).block(); } // getting one item from each feedRange @@ -495,10 +495,10 @@ public void faultInjectionServerErrorRuleTests_Partition() throws JsonProcessing String query = "select * from c"; CosmosQueryRequestOptions cosmosQueryRequestOptions = new CosmosQueryRequestOptions(); cosmosQueryRequestOptions.setFeedRange(feedRanges.get(0)); - TestItem itemOnFeedRange0 = cosmosAsyncContainer.queryItems(query, cosmosQueryRequestOptions, TestItem.class).blockFirst(); + TestObject itemOnFeedRange0 = cosmosAsyncContainer.queryItems(query, cosmosQueryRequestOptions, TestObject.class).blockFirst(); cosmosQueryRequestOptions.setFeedRange(feedRanges.get(1)); - TestItem itemOnFeedRange1 = cosmosAsyncContainer.queryItems(query, cosmosQueryRequestOptions, TestItem.class).blockFirst(); + TestObject itemOnFeedRange1 = cosmosAsyncContainer.queryItems(query, cosmosQueryRequestOptions, TestObject.class).blockFirst(); // set rule by feed range String feedRangeRuleId = "ServerErrorRule-FeedRange-" + UUID.randomUUID(); @@ -596,12 +596,12 @@ public void faultInjectionServerErrorRuleTests_ServerResponseDelay() throws Json .getContainer(cosmosAsyncContainer.getId()); // create a new item to be used by read operations - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); container.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules(container, Arrays.asList(timeoutRule)).block(); - CosmosItemResponse itemResponse = - container.readItem(createdItem.getId(), new PartitionKey(createdItem.getId()), TestItem.class).block(); + CosmosItemResponse itemResponse = + container.readItem(createdItem.getId(), new PartitionKey(createdItem.getId()), TestObject.class).block(); assertThat(timeoutRule.getHitCount()).isEqualTo(1); this.validateHitCount(timeoutRule, 1, OperationType.Read, ResourceType.Document); @@ -661,7 +661,7 @@ public void faultInjectionServerErrorRuleTests_ServerConnectionTimeout() throws .getContainer(cosmosAsyncContainer.getId()); CosmosFaultInjectionHelper.configureFaultInjectionRules(container, Arrays.asList(serverConnectionDelayRule)).block(); - CosmosItemResponse itemResponse = container.createItem(TestItem.createNewItem()).block(); + CosmosItemResponse itemResponse = container.createItem(TestObject.create()).block(); // Due to the replica validation, there could be an extra open connection call flow, while the rule will also be applied on. assertThat(serverConnectionDelayRule.getHitCount()).isBetween(1l, 2l); @@ -716,7 +716,7 @@ public void faultInjectionServerErrorRuleTests_ServerConnectionDelay() throws Js .getContainer(cosmosAsyncContainer.getId()); CosmosFaultInjectionHelper.configureFaultInjectionRules(container, Arrays.asList(serverConnectionDelayRule)).block(); - CosmosDiagnostics cosmosDiagnostics = container.createItem(TestItem.createNewItem()).block().getDiagnostics(); + CosmosDiagnostics cosmosDiagnostics = container.createItem(TestObject.create()).block().getDiagnostics(); // verify the request succeeded and the rule has applied List diagnosticsNode = new ArrayList<>(); @@ -868,7 +868,7 @@ public void faultInjectionServerErrorRuleTests_ServerErrorResponse( } // simulate high channel acquisition/connectionTimeout for read/query - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); cosmosAsyncContainer.createItem(createdItem).block(); String ruleId = "serverErrorRule-" + serverErrorType + "-" + UUID.randomUUID(); @@ -920,7 +920,7 @@ public void faultInjectionServerErrorRuleTests_LeaseNotFound(OperationType opera shouldRetryCrossRegion = true; } - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); String ruleId = "serverErrorRule-" + FaultInjectionServerErrorType.LEASE_NOT_FOUND + "-" + UUID.randomUUID(); FaultInjectionRule serverErrorRule = @@ -967,7 +967,7 @@ public void faultInjectionServerErrorRuleTests_LeaseNotFound(OperationType opera @Test(groups = {"multi-region", "long"}, timeOut = TIMEOUT) public void faultInjectionServerErrorRuleTests_HitLimit() throws JsonProcessingException { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); cosmosAsyncContainer.createItem(createdItem).block(); // set rule by feed range @@ -1026,7 +1026,7 @@ public void afterClass() { @Test(groups = {"long"}, timeOut = TIMEOUT) public void faultInjectionServerErrorRuleTests_includePrimary() throws JsonProcessingException { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); CosmosAsyncContainer singlePartitionContainer = getSharedSinglePartitionCosmosContainer(clientWithoutPreferredRegions); List feedRanges = singlePartitionContainer.getFeedRanges().block(); @@ -1100,7 +1100,7 @@ public void faultInjectionServerErrorRuleTests_StaledAddressesServerGone() throw .build(); try { - TestItem testItem = this.cosmosAsyncContainer.createItem(TestItem.createNewItem()).block().getItem(); + TestObject testItem = this.cosmosAsyncContainer.createItem(TestObject.create()).block().getItem(); CosmosFaultInjectionHelper.configureFaultInjectionRules(this.cosmosAsyncContainer, Arrays.asList(staledAddressesServerGoneRule)).block(); @@ -1163,7 +1163,7 @@ public void connectionAcquisitionTimeoutAlignConnectionTimeout() throws JsonProc // validate none of the channelAcquisition stage take more than 2s List results = new ArrayList<>(); Flux.range(1, 6) - .flatMap(t -> containerWithSinglePartition.createItem(TestItem.createNewItem())) + .flatMap(t -> containerWithSinglePartition.createItem(TestObject.create())) .doOnNext(response -> results.add(response.getDiagnostics())) .blockLast(); @@ -1202,7 +1202,7 @@ public void faultInjectionServerErrorRuleTests_InjectionRate50Percent() throws J .build(); try { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); cosmosAsyncContainer.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules( @@ -1244,7 +1244,7 @@ public void faultInjectionServerErrorRuleTests_InjectionRate100Percent() throws .build(); try { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); cosmosAsyncContainer.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules( @@ -1311,7 +1311,7 @@ public void faultInjectionServerErrorRuleTests_InjectionRate25Percent() throws J .build(); try { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); cosmosAsyncContainer.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules( @@ -1353,7 +1353,7 @@ public void faultInjectionServerErrorRuleTests_InjectionRate75Percent() throws J .build(); try { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); cosmosAsyncContainer.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules( diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnGatewayTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnGatewayTests.java index b8a552c1a505..884a63f329d3 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnGatewayTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnGatewayTests.java @@ -10,6 +10,7 @@ import com.azure.cosmos.CosmosDiagnostics; import com.azure.cosmos.CosmosException; import com.azure.cosmos.DirectConnectionConfig; +import com.azure.cosmos.TestObject; import com.azure.cosmos.implementation.AsyncDocumentClient; import com.azure.cosmos.implementation.DatabaseAccount; import com.azure.cosmos.implementation.DatabaseAccountLocation; @@ -19,7 +20,6 @@ import com.azure.cosmos.implementation.ResourceType; import com.azure.cosmos.implementation.TestConfigurations; import com.azure.cosmos.implementation.Utils; -import com.azure.cosmos.implementation.throughputControl.TestItem; import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.FeedRange; @@ -194,7 +194,7 @@ public void faultInjectionServerErrorRuleTests_Region(boolean shouldUsePreferred .getDatabase(this.cosmosAsyncContainer.getDatabase().getId()) .getContainer(this.cosmosAsyncContainer.getId()); - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); container.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules( @@ -251,7 +251,7 @@ public void faultInjectionServerErrorRuleTests_Partition() throws JsonProcessing .getContainer(cosmosAsyncContainer.getId()); for (int i = 0; i < 10; i++) { - testContainer.createItem(TestItem.createNewItem()).block(); + testContainer.createItem(TestObject.create()).block(); } // getting one item from each feedRange @@ -261,10 +261,10 @@ public void faultInjectionServerErrorRuleTests_Partition() throws JsonProcessing String query = "select * from c"; CosmosQueryRequestOptions cosmosQueryRequestOptions = new CosmosQueryRequestOptions(); cosmosQueryRequestOptions.setFeedRange(feedRanges.get(0)); - TestItem itemOnFeedRange0 = testContainer.queryItems(query, cosmosQueryRequestOptions, TestItem.class).blockFirst(); + TestObject itemOnFeedRange0 = testContainer.queryItems(query, cosmosQueryRequestOptions, TestObject.class).blockFirst(); cosmosQueryRequestOptions.setFeedRange(feedRanges.get(1)); - TestItem itemOnFeedRange1 = testContainer.queryItems(query, cosmosQueryRequestOptions, TestItem.class).blockFirst(); + TestObject itemOnFeedRange1 = testContainer.queryItems(query, cosmosQueryRequestOptions, TestObject.class).blockFirst(); // set rule by feed range String feedRangeRuleId = "ServerErrorRule-FeedRange-" + UUID.randomUUID(); @@ -350,12 +350,12 @@ public void faultInjectionServerErrorRuleTests_ServerResponseDelay() throws Json directConnectionConfig.setConnectTimeout(Duration.ofSeconds(1)); // create a new item to be used by read operations - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); this.cosmosAsyncContainer.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules(this.cosmosAsyncContainer, Arrays.asList(timeoutRule)).block(); - CosmosItemResponse itemResponse = - this.cosmosAsyncContainer.readItem(createdItem.getId(), new PartitionKey(createdItem.getId()), TestItem.class).block(); + CosmosItemResponse itemResponse = + this.cosmosAsyncContainer.readItem(createdItem.getId(), new PartitionKey(createdItem.getId()), TestObject.class).block(); assertThat(timeoutRule.getHitCount()).isEqualTo(1); this.validateHitCount(timeoutRule, 1, OperationType.Read, ResourceType.Document); @@ -398,7 +398,7 @@ public void faultInjectionServerErrorRuleTests_ServerConnectionDelay() throws Js try { CosmosFaultInjectionHelper.configureFaultInjectionRules(this.cosmosAsyncContainer, Arrays.asList(serverConnectionDelayRule)).block(); - CosmosItemResponse itemResponse = this.cosmosAsyncContainer.createItem(TestItem.createNewItem()).block(); + CosmosItemResponse itemResponse = this.cosmosAsyncContainer.createItem(TestObject.create()).block(); assertThat(serverConnectionDelayRule.getHitCount()).isEqualTo(1l); this.validateFaultInjectionRuleApplied( @@ -449,7 +449,7 @@ public void faultInjectionServerErrorRuleTests_ServerErrorResponse( .build(); try { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); cosmosAsyncContainer.createItem(createdItem).block(); CosmosFaultInjectionHelper.configureFaultInjectionRules(cosmosAsyncContainer, Arrays.asList(serverErrorRule)).block(); @@ -459,7 +459,7 @@ public void faultInjectionServerErrorRuleTests_ServerErrorResponse( try { cosmosDiagnostics = cosmosAsyncContainer - .readItem(createdItem.getId(), new PartitionKey(createdItem.getId()), TestItem.class) + .readItem(createdItem.getId(), new PartitionKey(createdItem.getId()), TestObject.class) .block() .getDiagnostics(); } catch (Exception exception) { @@ -469,7 +469,7 @@ public void faultInjectionServerErrorRuleTests_ServerErrorResponse( try { cosmosDiagnostics = cosmosAsyncContainer - .readItem(createdItem.getId(), new PartitionKey(createdItem.getId()), TestItem.class) + .readItem(createdItem.getId(), new PartitionKey(createdItem.getId()), TestObject.class) .block() .getDiagnostics(); fail("Request should fail, but succeeded"); @@ -499,7 +499,7 @@ public void faultInjectionServerErrorRuleTests_HitLimit( OperationType operationType, FaultInjectionOperationType faultInjectionOperationType) throws JsonProcessingException { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); cosmosAsyncContainer.createItem(createdItem).block(); // set rule by feed range diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionTestBase.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionTestBase.java index af01ed9dde83..33f14591ea69 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionTestBase.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionTestBase.java @@ -7,8 +7,8 @@ import com.azure.cosmos.CosmosClientBuilder; import com.azure.cosmos.CosmosDiagnostics; import com.azure.cosmos.CosmosException; +import com.azure.cosmos.TestObject; import com.azure.cosmos.implementation.OperationType; -import com.azure.cosmos.implementation.throughputControl.TestItem; import com.azure.cosmos.models.CosmosBatch; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; import com.azure.cosmos.models.CosmosItemIdentity; @@ -30,14 +30,14 @@ public FaultInjectionTestBase(CosmosClientBuilder cosmosClientBuilder) { protected CosmosDiagnostics performDocumentOperation( CosmosAsyncContainer cosmosAsyncContainer, OperationType operationType, - TestItem createdItem, + TestObject createdItem, boolean isReadMany) { try { if (operationType == OperationType.Query && !isReadMany) { CosmosQueryRequestOptions queryRequestOptions = new CosmosQueryRequestOptions(); String query = String.format("SELECT * from c where c.id = '%s'", createdItem.getId()); - FeedResponse itemFeedResponse = - cosmosAsyncContainer.queryItems(query, queryRequestOptions, TestItem.class).byPage().blockLast(); + FeedResponse itemFeedResponse = + cosmosAsyncContainer.queryItems(query, queryRequestOptions, TestObject.class).byPage().blockLast(); return itemFeedResponse.getCosmosDiagnostics(); } @@ -54,7 +54,7 @@ protected CosmosDiagnostics performDocumentOperation( return cosmosAsyncContainer.readItem( createdItem.getId(), new PartitionKey(createdItem.getId()), - TestItem.class).block().getDiagnostics(); + TestObject.class).block().getDiagnostics(); } if (operationType == OperationType.Replace) { @@ -69,11 +69,11 @@ protected CosmosDiagnostics performDocumentOperation( } if (operationType == OperationType.Create) { - return cosmosAsyncContainer.createItem(TestItem.createNewItem()).block().getDiagnostics(); + return cosmosAsyncContainer.createItem(TestObject.create()).block().getDiagnostics(); } if (operationType == OperationType.Upsert) { - return cosmosAsyncContainer.upsertItem(TestItem.createNewItem()).block().getDiagnostics(); + return cosmosAsyncContainer.upsertItem(TestObject.create()).block().getDiagnostics(); } if (operationType == OperationType.Patch) { @@ -82,7 +82,7 @@ protected CosmosDiagnostics performDocumentOperation( .create() .add("/newPath", "newPath"); return cosmosAsyncContainer - .patchItem(createdItem.getId(), new PartitionKey(createdItem.getId()), patchOperations, TestItem.class) + .patchItem(createdItem.getId(), new PartitionKey(createdItem.getId()), patchOperations, TestObject.class) .block().getDiagnostics(); } @@ -101,8 +101,8 @@ protected CosmosDiagnostics performDocumentOperation( CosmosChangeFeedRequestOptions changeFeedRequestOptions = CosmosChangeFeedRequestOptions.createForProcessingFromBeginning(feedRanges.get(0)); - FeedResponse firstPage = cosmosAsyncContainer - .queryChangeFeed(changeFeedRequestOptions, TestItem.class) + FeedResponse firstPage = cosmosAsyncContainer + .queryChangeFeed(changeFeedRequestOptions, TestObject.class) .byPage() .blockFirst(); return firstPage.getCosmosDiagnostics(); @@ -111,7 +111,7 @@ protected CosmosDiagnostics performDocumentOperation( if (operationType == OperationType.Query) { return cosmosAsyncContainer.readMany( Arrays.asList(new CosmosItemIdentity(new PartitionKey(createdItem.getId()), createdItem.getId()), new CosmosItemIdentity(new PartitionKey(createdItem.getId()), createdItem.getId())), - TestItem.class).block().getCosmosDiagnostics(); + TestObject.class).block().getCosmosDiagnostics(); } throw new IllegalArgumentException("The operation type is not supported"); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/SessionRetryOptionsTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/SessionRetryOptionsTests.java index cab6bfa64a90..261ae9ee0a27 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/SessionRetryOptionsTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/SessionRetryOptionsTests.java @@ -11,6 +11,7 @@ import com.azure.cosmos.CosmosRegionSwitchHint; import com.azure.cosmos.SessionRetryOptions; import com.azure.cosmos.SessionRetryOptionsBuilder; +import com.azure.cosmos.TestObject; import com.azure.cosmos.implementation.AsyncDocumentClient; import com.azure.cosmos.implementation.Configs; import com.azure.cosmos.implementation.DatabaseAccount; @@ -20,7 +21,6 @@ import com.azure.cosmos.implementation.ImplementationBridgeHelpers; import com.azure.cosmos.implementation.OperationType; import com.azure.cosmos.implementation.TestConfigurations; -import com.azure.cosmos.implementation.throughputControl.TestItem; import com.azure.cosmos.models.CosmosItemRequestOptions; import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.CosmosPatchOperations; @@ -55,7 +55,6 @@ import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicReference; -import java.util.stream.Collectors; import static org.assertj.core.api.AssertionsForClassTypes.assertThat; import static org.testng.AssertJUnit.fail; @@ -238,7 +237,7 @@ public void nonWriteOperation_WithReadSessionUnavailable_test( .getDatabase(this.cosmosAsyncContainer.getDatabase().getId()) .getContainer(this.cosmosAsyncContainer.getId()); - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); containerForClientWithPreferredRegions.createItem(createdItem).block(); FaultInjectionRuleBuilder badSessionTokenRuleBuilder = new FaultInjectionRuleBuilder("serverErrorRule-bad-session-token-" + UUID.randomUUID()); @@ -334,7 +333,7 @@ public void writeOperation_withReadSessionUnavailable_test( .configureFaultInjectionRules(asyncContainerFromClientWithPreferredRegions, Arrays.asList(badSessionTokenRule)) .block(); - TestItem testItem = TestItem.createNewItem(); + TestObject testItem = TestObject.create(); validateOperationExecutionResult( performDocumentOperation( @@ -379,7 +378,7 @@ private Map getRegionMap(DatabaseAccount databaseAccount, boolea private OperationExecutionResult performDocumentOperation( CosmosAsyncContainer faultInjectedContainer, - TestItem testItem, + TestObject testItem, OperationType operationType) { AtomicReference operationStart = new AtomicReference<>(Instant.now()); @@ -390,8 +389,8 @@ private OperationExecutionResult performDocumentOperation( CosmosQueryRequestOptions queryRequestOptions = new CosmosQueryRequestOptions(); SqlQuerySpec sqlQuerySpec = new SqlQuerySpec(query); - FeedResponse feedResponse = faultInjectedContainer - .queryItems(sqlQuerySpec, queryRequestOptions, TestItem.class) + FeedResponse feedResponse = faultInjectedContainer + .queryItems(sqlQuerySpec, queryRequestOptions, TestObject.class) .byPage() .doOnSubscribe(ignore -> operationStart.set(Instant.now())) .doOnCancel(() -> operationEnd.set(Instant.now())) @@ -410,8 +409,8 @@ private OperationExecutionResult performDocumentOperation( if (operationType == OperationType.Read) { - CosmosItemResponse itemResponse = faultInjectedContainer - .readItem(testItem.getId(), new PartitionKey(testItem.getId()), TestItem.class) + CosmosItemResponse itemResponse = faultInjectedContainer + .readItem(testItem.getId(), new PartitionKey(testItem.getId()), TestObject.class) .doOnSubscribe(ignore -> operationStart.set(Instant.now())) .doOnSuccess(ignore -> operationEnd.set(Instant.now())) .block(); @@ -427,7 +426,7 @@ private OperationExecutionResult performDocumentOperation( CosmosItemRequestOptions itemRequestOptions = new CosmosItemRequestOptions(); - CosmosItemResponse itemResponse = faultInjectedContainer + CosmosItemResponse itemResponse = faultInjectedContainer .createItem(testItem, new PartitionKey(testItem.getId()), itemRequestOptions) .doOnSubscribe(ignore -> operationStart.set(Instant.now())) .doOnSuccess(ignore -> operationEnd.set(Instant.now())) @@ -448,7 +447,7 @@ private OperationExecutionResult performDocumentOperation( .createItem(testItem, new PartitionKey(testItem.getId()), itemRequestOptions) .block(); - CosmosItemResponse itemResponse = faultInjectedContainer + CosmosItemResponse itemResponse = faultInjectedContainer .replaceItem(testItem, testItem.getId(), new PartitionKey(testItem.getId()), itemRequestOptions) .doOnSubscribe(ignore -> operationStart.set(Instant.now())) .doOnSuccess(ignore -> operationEnd.set(Instant.now())) @@ -484,7 +483,7 @@ private OperationExecutionResult performDocumentOperation( if (operationType == OperationType.Upsert) { CosmosItemRequestOptions itemRequestOptions = new CosmosItemRequestOptions(); - CosmosItemResponse itemResponse = faultInjectedContainer + CosmosItemResponse itemResponse = faultInjectedContainer .upsertItem(testItem, new PartitionKey(testItem.getId()), itemRequestOptions) .doOnSubscribe(ignore -> operationStart.set(Instant.now())) .doOnSuccess(ignore -> operationEnd.set(Instant.now())) @@ -505,8 +504,8 @@ private OperationExecutionResult performDocumentOperation( .createItem(testItem, new PartitionKey(testItem.getId()), itemRequestOptions) .block(); - CosmosItemResponse itemResponse = faultInjectedContainer - .patchItem(testItem.getId(), new PartitionKey(testItem.getId()), patchOperations, TestItem.class) + CosmosItemResponse itemResponse = faultInjectedContainer + .patchItem(testItem.getId(), new PartitionKey(testItem.getId()), patchOperations, TestObject.class) .doOnSubscribe(ignore -> operationStart.set(Instant.now())) .doOnSuccess(ignore -> operationEnd.set(Instant.now())) .block(); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/MetadataRequestRetryPolicyTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/MetadataRequestRetryPolicyTests.java index 13b8cb738289..210112149959 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/MetadataRequestRetryPolicyTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/MetadataRequestRetryPolicyTests.java @@ -13,6 +13,7 @@ import com.azure.cosmos.CosmosEndToEndOperationLatencyPolicyConfigBuilder; import com.azure.cosmos.CosmosException; import com.azure.cosmos.DirectConnectionConfig; +import com.azure.cosmos.TestObject; import com.azure.cosmos.implementation.AsyncDocumentClient; import com.azure.cosmos.implementation.Configs; import com.azure.cosmos.implementation.DatabaseAccount; @@ -33,7 +34,6 @@ import com.azure.cosmos.implementation.http.HttpClientConfig; import com.azure.cosmos.implementation.http.HttpTimeoutPolicyControlPlaneHotPath; import com.azure.cosmos.implementation.routing.RegionalRoutingContext; -import com.azure.cosmos.implementation.throughputControl.TestItem; import com.azure.cosmos.models.CosmosItemRequestOptions; import com.azure.cosmos.models.CosmosPatchOperations; import com.azure.cosmos.models.CosmosQueryRequestOptions; @@ -291,7 +291,7 @@ public void forceBackgroundAddressRefresh_onConnectionTimeoutAndRequestCancellat CosmosEndToEndOperationLatencyPolicyConfig cosmosEndToEndOperationLatencyPolicyConfigForFaultyOperation = new CosmosEndToEndOperationLatencyPolicyConfigBuilder(Duration.ofMillis(500)).build(); - TestItem testItem = new TestItem(UUID.randomUUID().toString(), UUID.randomUUID().toString(), UUID.randomUUID().toString()); + TestObject testItem = new TestObject(UUID.randomUUID().toString(), UUID.randomUUID().toString(), Arrays.asList(), UUID.randomUUID().toString()); performDocumentOperation( container, @@ -378,7 +378,7 @@ public void metadataRetryPolicyTest( private void performDocumentOperation( CosmosAsyncContainer faultInjectedContainer, - TestItem testItem, + TestObject testItem, OperationType faultInjectedOperationType, FaultInjectionRule connectionDelayFault, HttpClientUnderTestWrapper httpClientUnderTestWrapper, @@ -409,7 +409,7 @@ private void performDocumentOperation( SqlQuerySpec sqlQuerySpec = new SqlQuerySpec(query); faultInjectedContainer - .queryItems(sqlQuerySpec, queryRequestOptions, TestItem.class) + .queryItems(sqlQuerySpec, queryRequestOptions, TestObject.class) .byPage() .subscribe(); } else if (faultInjectedOperationType == OperationType.Read) { @@ -427,7 +427,7 @@ private void performDocumentOperation( .setCosmosEndToEndOperationLatencyPolicyConfig(endToEndOperationLatencyPolicyConfigForFaultyOperation); faultInjectedContainer - .readItem(testItem.getId(), new PartitionKey(testItem.getMypk()), requestOptionsForRead, TestItem.class) + .readItem(testItem.getId(), new PartitionKey(testItem.getMypk()), requestOptionsForRead, TestObject.class) .subscribe(); } else if (faultInjectedOperationType == OperationType.Create) { @@ -496,7 +496,7 @@ private void performDocumentOperation( Thread.sleep(idleTimeInMillis); faultInjectedContainer - .patchItem(testItem.getId(), new PartitionKey(testItem.getMypk()), patchOperations, TestItem.class) + .patchItem(testItem.getId(), new PartitionKey(testItem.getMypk()), patchOperations, TestObject.class) .subscribe(); } } diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/ReflectionUtils.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/ReflectionUtils.java index 2fe1cef5a869..62831a86e9fb 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/ReflectionUtils.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/ReflectionUtils.java @@ -40,10 +40,10 @@ import com.azure.cosmos.implementation.http.HttpRequest; import com.azure.cosmos.implementation.routing.CollectionRoutingMap; import com.azure.cosmos.implementation.routing.LocationCache; -import com.azure.cosmos.implementation.throughputControl.ThroughputControlTrackingUnit; -import com.azure.cosmos.implementation.throughputControl.ThroughputRequestThrottler; -import com.azure.cosmos.implementation.throughputControl.controller.request.GlobalThroughputRequestController; -import com.azure.cosmos.implementation.throughputControl.controller.request.PkRangesThroughputRequestController; +import com.azure.cosmos.implementation.throughputControl.sdk.ThroughputControlTrackingUnit; +import com.azure.cosmos.implementation.throughputControl.sdk.ThroughputRequestThrottler; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.request.GlobalThroughputRequestController; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.request.PkRangesThroughputRequestController; import com.azure.cosmos.models.CosmosClientTelemetryConfig; import io.netty.handler.ssl.SslContext; import org.apache.commons.lang3.reflect.FieldUtils; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupConfigTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupConfigTests.java new file mode 100644 index 000000000000..0700f857d8dc --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupConfigTests.java @@ -0,0 +1,45 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl; + +import com.azure.cosmos.ThroughputControlGroupConfig; +import com.azure.cosmos.ThroughputControlGroupConfigBuilder; +import org.testng.annotations.Test; + +import java.util.UUID; + +import static org.assertj.core.api.AssertionsForClassTypes.assertThat; +import static org.assertj.core.api.AssertionsForClassTypes.assertThatThrownBy; + +public class ThroughputControlGroupConfigTests { + + @Test(groups = "unit") + public void throughputControlGroup_throughputBucket_priorityLevel() { + // validate throughputBucket >= 0 + assertThatThrownBy( + () -> new ThroughputControlGroupConfigBuilder() + .groupName("throughputBucket-" + UUID.randomUUID()) + .throughputBucket(-1) + .build()) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Throughput bucket should be no smaller than 0"); + + // valid config + ThroughputControlGroupConfig throughputControlGroupConfig = + new ThroughputControlGroupConfigBuilder() + .groupName("throughputBucket-" + UUID.randomUUID()) + .throughputBucket(1) + .build(); + + assertThat(throughputControlGroupConfig.getThroughputBucket()).isEqualTo(1); + + // neither priorityLevel nor throughput bucket configured + assertThatThrownBy( + () -> new ThroughputControlGroupConfigBuilder() + .groupName("throughputBucket-" + UUID.randomUUID()) + .build()) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("All targetThroughput, targetThroughputThreshold, priorityLevel and throughput bucket cannot be null or empty."); + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java index be0d563e54fc..ca97bcab6645 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTests.java @@ -21,7 +21,7 @@ import com.azure.cosmos.implementation.TestConfigurations; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.implementation.apachecommons.lang.tuple.Pair; -import com.azure.cosmos.implementation.throughputControl.controller.group.global.GlobalThroughputControlClientItem; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.group.global.GlobalThroughputControlClientItem; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; import com.azure.cosmos.models.CosmosContainerProperties; import com.azure.cosmos.models.CosmosContainerRequestOptions; @@ -52,6 +52,7 @@ import java.util.concurrent.atomic.AtomicInteger; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.AssertionsForClassTypes.assertThatThrownBy; import static org.testng.Assert.fail; public class ThroughputControlTests extends TestSuiteBase { @@ -106,15 +107,24 @@ public void throughputLocalControl_requestOptions(OperationType operationType) { CosmosItemResponse createItemResponse = container.createItem(getDocumentDefinition(), requestOptions).block(); TestItem createdItem = createItemResponse.getItem(); - this.validateRequestNotThrottled( - createItemResponse.getDiagnostics().toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, 6) + ); // second request to group-1. which will get throttled - CosmosDiagnostics cosmosDiagnostics = performDocumentOperation(this.container, operationType, createdItem, groupConfig.getGroupName()); - this.validateRequestThrottled( - cosmosDiagnostics.toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + cosmosDiagnosticsString = performDocumentOperation(this.container, operationType, createdItem, groupConfig.getGroupName()).toString(); + this.validateRequestThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, 6) + ); } @Test(groups = {"emulator"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT) @@ -141,15 +151,24 @@ public void throughputLocalControl_default(OperationType operationType) { CosmosItemResponse createItemResponse = testContainer.createItem(getDocumentDefinition()).block(); TestItem createdItem = createItemResponse.getItem(); - this.validateRequestNotThrottled( - createItemResponse.getDiagnostics().toString(), - BridgeInternal.getContextClient(cosmosAsyncClient).getConnectionPolicy().getConnectionMode()); + + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), true, 6) + ); // second request to group-1. which will get throttled - CosmosDiagnostics cosmosDiagnostics = performDocumentOperation(testContainer, operationType, createdItem, groupConfig.getGroupName()); - this.validateRequestThrottled( - cosmosDiagnostics.toString(), - BridgeInternal.getContextClient(cosmosAsyncClient).getConnectionPolicy().getConnectionMode()); + cosmosDiagnosticsString = + performDocumentOperation(testContainer, operationType, createdItem, groupConfig.getGroupName()).toString(); + this.validateRequestThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), true, 6) + ); } finally { safeClose(cosmosAsyncClient); } @@ -189,15 +208,23 @@ public void throughputLocalControlWithThroughputQuery() { CosmosItemResponse createItemResponse = testContainer.createItem(getDocumentDefinition(), requestOptions).block(); TestItem createdItem = createItemResponse.getItem(); - this.validateRequestNotThrottled( - createItemResponse.getDiagnostics().toString(), - BridgeInternal.getContextClient(cosmosAsyncClient).getConnectionPolicy().getConnectionMode()); + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRUThreshold=%s)", groupConfig.getGroupName(), false, 0.9) + ); // second request to group-1. which will get throttled - CosmosDiagnostics cosmosDiagnostics = performDocumentOperation(testContainer, OperationType.Read, createdItem, groupConfig.getGroupName()); - this.validateRequestThrottled( - cosmosDiagnostics.toString(), - BridgeInternal.getContextClient(cosmosAsyncClient).getConnectionPolicy().getConnectionMode()); + cosmosDiagnosticsString = + performDocumentOperation(testContainer, OperationType.Read, createdItem, groupConfig.getGroupName()).toString(); + this.validateRequestThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRUThreshold=%s)", groupConfig.getGroupName(), false, 0.9) + ); assertThat(throughputQueryMonoCalledCount.get()).isGreaterThanOrEqualTo(1); } finally { @@ -221,9 +248,19 @@ public void throughputLocalControlPriorityLevel(OperationType operationType) { CosmosItemResponse createItemResponse = container.createItem(getDocumentDefinition(), requestOptions).block(); TestItem createdItem = createItemResponse.getItem(); + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, priorityLevel=%s, targetRU=%s)", groupConfig.getGroupName(), false, PriorityLevel.LOW, Integer.MAX_VALUE) + ); - performDocumentOperation(container, operationType, createdItem, groupConfig.getGroupName()); - + cosmosDiagnosticsString = performDocumentOperation(container, operationType, createdItem, groupConfig.getGroupName()).toString(); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, priorityLevel=%s, targetRU=%s)", groupConfig.getGroupName(), false, PriorityLevel.LOW, Integer.MAX_VALUE) + ); assertThat(groupConfig.getTargetThroughput()).isEqualTo(Integer.MAX_VALUE); assertThat(createItemResponse.getStatusCode()).isEqualTo(201); } @@ -256,15 +293,24 @@ public void throughputGlobalControl(OperationType operationType) { CosmosItemResponse createItemResponse = container.createItem(getDocumentDefinition(), requestOptions).block(); TestItem createdItem = createItemResponse.getItem(); - this.validateRequestNotThrottled( - createItemResponse.getDiagnostics().toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, 6) + ); // second request to same group. which will get throttled - CosmosDiagnostics cosmosDiagnostics = performDocumentOperation(this.container, operationType, createdItem, groupConfig.getGroupName()); - this.validateRequestThrottled( - cosmosDiagnostics.toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + cosmosDiagnosticsString = + performDocumentOperation(this.container, operationType, createdItem, groupConfig.getGroupName()).toString(); + this.validateRequestThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, 6) + ); } finally { controlContainer .delete() @@ -313,15 +359,23 @@ public void throughputGlobalControlWithThroughputQuery() { CosmosItemResponse createItemResponse = testContainer.createItem(getDocumentDefinition(), requestOptions).block(); TestItem createdItem = createItemResponse.getItem(); - this.validateRequestNotThrottled( - createItemResponse.getDiagnostics().toString(), - BridgeInternal.getContextClient(cosmosAsyncClient).getConnectionPolicy().getConnectionMode()); + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRUThreshold=%s)", groupConfig.getGroupName(), false, 0.9) + ); // second request to same group. which will get throttled - CosmosDiagnostics cosmosDiagnostics = performDocumentOperation(testContainer, OperationType.Create, createdItem, groupConfig.getGroupName()); - this.validateRequestThrottled( - cosmosDiagnostics.toString(), - BridgeInternal.getContextClient(cosmosAsyncClient).getConnectionPolicy().getConnectionMode()); + cosmosDiagnosticsString = + performDocumentOperation(testContainer, OperationType.Create, createdItem, groupConfig.getGroupName()).toString(); + this.validateRequestThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRUThreshold=%s)", groupConfig.getGroupName(), false, 0.9) + ); assertThat(throughputQueryMonoCalledCount.get()).isGreaterThanOrEqualTo(1); } finally { @@ -377,21 +431,29 @@ public void throughputGlobalControlCanUpdateConfig(OperationType operationType) .createItem(getDocumentDefinition(), requestOptions) .block(); TestItem createdItem = createItemResponse.getItem(); - this.validateRequestNotThrottled( - createItemResponse.getDiagnostics().toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); - CosmosDiagnostics cosmosDiagnostics = performDocumentOperation(this.container, operationType, createdItem, groupConfig.getGroupName()); + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, targetThroughput) + ); + + cosmosDiagnosticsString = + performDocumentOperation(this.container, operationType, createdItem, groupConfig.getGroupName()).toString(); if (shouldThrottleSecondRequest) { - this.validateRequestThrottled( - cosmosDiagnostics.toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + this.validateRequestThrottled(cosmosDiagnosticsString); } else { - this.validateRequestNotThrottled( - cosmosDiagnostics.toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + this.validateRequestNotThrottled(cosmosDiagnosticsString); } + + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, targetThroughput) + ); } } finally { safeDeleteCollection(database.getContainer(controlContainerId)); @@ -435,9 +497,14 @@ public void throughputLocalControlForContainerCreateDeleteWithSameName(Operation // Step2: first request to group-1, which will not get throttled, but will consume all the rus of the throughput control group. CosmosItemResponse createItemResponse = createdContainer.createItem(getDocumentDefinition(), requestOptions).block(); TestItem createdItem = createItemResponse.getItem(); - this.validateRequestNotThrottled( - createItemResponse.getDiagnostics().toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, 1) + ); // Step 3: delete the container safeDeleteCollection(createdContainer); @@ -451,16 +518,24 @@ public void throughputLocalControlForContainerCreateDeleteWithSameName(Operation createdItem = createdContainer.createItem(getDocumentDefinition()).block().getItem(); // Step 6: second request to group-1. which will not get throttled because new container controller will be built. - CosmosDiagnostics cosmosDiagnostics = performDocumentOperation(createdContainer, operationType, createdItem, groupConfig.getGroupName()); - this.validateRequestNotThrottled( - cosmosDiagnostics.toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + cosmosDiagnosticsString = + performDocumentOperation(createdContainer, operationType, createdItem, groupConfig.getGroupName()).toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, 1) + ); // Step 7: third request to group-1, which will get throttled - cosmosDiagnostics = performDocumentOperation(createdContainer, operationType, createdItem, groupConfig.getGroupName()); - this.validateRequestThrottled( - cosmosDiagnostics.toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + cosmosDiagnosticsString = + performDocumentOperation(createdContainer, operationType, createdItem, groupConfig.getGroupName()).toString(); + this.validateRequestThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, 1) + ); } finally { createdContainer.delete().block(); } @@ -483,9 +558,14 @@ public void throughputLocalControl_createItem() throws InterruptedException { requestOptions.setThroughputControlGroupName(groupConfig.getGroupName()); CosmosItemResponse createItemResponse = container.createItem(getDocumentDefinition(), requestOptions).block(); - this.validateRequestNotThrottled( - createItemResponse.getDiagnostics().toString(), - BridgeInternal.getContextClient(client).getConnectionPolicy().getConnectionMode()); + + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, 6) + ); // second request to same group. which will get throttled TestItem itemGetThrottled = getDocumentDefinition(createItemResponse.getItem().getMypk()); @@ -543,7 +623,7 @@ public void throughputControlContinueOnInitError(boolean continueOnInitError) { } @Test(groups = {"emulator"}, timeOut = TIMEOUT * 4) - public void throughputGlobalControlMultipleClients() throws InterruptedException { + public void throughputGlobalControlMultipleClients() { this.ensureContainer(); List cosmosAsyncClients = new ArrayList<>(); // and do not enable ttl on the container so to test how many items are created. @@ -670,8 +750,242 @@ public void enableSameGroupMultipleTimes() { } } + @Test(groups = {"emulator"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT) + public void serverThroughputControl_throughputBucket(OperationType operationType) { + // TODO: currently there is no easy way to do e2e testing, so the testing here is to just verify that + // server throughput can be enabled on the container with throughput bucket + this.ensureContainer(); + + ThroughputControlGroupConfig serverThroughputControlGroup = + new ThroughputControlGroupConfigBuilder() + .groupName("serverThroughputControl") + .throughputBucket(2) + .build(); + container.enableServerThroughputControlGroup(serverThroughputControlGroup); + + CosmosItemRequestOptions requestOptions = new CosmosItemRequestOptions(); + requestOptions.setContentResponseOnWriteEnabled(true); + requestOptions.setThroughputControlGroupName(serverThroughputControlGroup.getGroupName()); + + CosmosItemResponse createItemResponse = container.createItem(getDocumentDefinition(), requestOptions).block(); + TestItem createdItem = createItemResponse.getItem(); + + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + serverThroughputControlGroup.getGroupName(), + String.format("(name=%s, default=%s, throughputBucket=%s)", serverThroughputControlGroup.getGroupName(), false, 2) + ); + + cosmosDiagnosticsString = performDocumentOperation( + this.container, + operationType, + createdItem, + serverThroughputControlGroup.getGroupName()).toString(); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + serverThroughputControlGroup.getGroupName(), + String.format("(name=%s, default=%s, throughputBucket=%s)", serverThroughputControlGroup.getGroupName(), false, 2) + ); + } + + @Test(groups = {"emulator"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT) + public void serverThroughputControl_priorityLevel(OperationType operationType) { + // TODO: currently there is no easy way to do e2e testing, so the testing here is to just verify that + // server throughput can be enabled on the container with priority level + this.ensureContainer(); + + ThroughputControlGroupConfig serverThroughputControlGroup = + new ThroughputControlGroupConfigBuilder() + .groupName("serverThroughputControl") + .priorityLevel(PriorityLevel.LOW) + .build(); + container.enableServerThroughputControlGroup(serverThroughputControlGroup); + + CosmosItemRequestOptions requestOptions = new CosmosItemRequestOptions(); + requestOptions.setContentResponseOnWriteEnabled(true); + requestOptions.setThroughputControlGroupName(serverThroughputControlGroup.getGroupName()); + + CosmosItemResponse createItemResponse = container.createItem(getDocumentDefinition(), requestOptions).block(); + TestItem createdItem = createItemResponse.getItem(); + + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + serverThroughputControlGroup.getGroupName(), + String.format("(name=%s, default=%s, priorityLevel=%s)", serverThroughputControlGroup.getGroupName(), false, PriorityLevel.LOW) + ); + + cosmosDiagnosticsString = performDocumentOperation( + this.container, + operationType, + createdItem, + serverThroughputControlGroup.getGroupName()).toString(); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + serverThroughputControlGroup.getGroupName(), + String.format("(name=%s, default=%s, priorityLevel=%s)", serverThroughputControlGroup.getGroupName(), false, PriorityLevel.LOW) + ); + } + + @Test(groups = {"emulator"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT) + public void serverThroughputControl_priorityLevel_throughputBucket(OperationType operationType) { + // TODO: currently there is no easy way to do e2e testing, so the testing here is to just verify that + // server throughput can be enabled on the container with priority level and throughput bucket + this.ensureContainer(); + + ThroughputControlGroupConfig serverThroughputControlGroup = + new ThroughputControlGroupConfigBuilder() + .groupName("serverThroughputControl") + .priorityLevel(PriorityLevel.LOW) + .throughputBucket(3) + .build(); + container.enableServerThroughputControlGroup(serverThroughputControlGroup); + + CosmosItemRequestOptions requestOptions = new CosmosItemRequestOptions(); + requestOptions.setContentResponseOnWriteEnabled(true); + requestOptions.setThroughputControlGroupName(serverThroughputControlGroup.getGroupName()); + + CosmosItemResponse createItemResponse = container.createItem(getDocumentDefinition(), requestOptions).block(); + TestItem createdItem = createItemResponse.getItem(); + + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + serverThroughputControlGroup.getGroupName(), + String.format( + "(name=%s, default=%s, priorityLevel=%s, throughputBucket=%s)", + serverThroughputControlGroup.getGroupName(), + false, + PriorityLevel.LOW, + 3) + ); + + cosmosDiagnosticsString = performDocumentOperation( + this.container, + operationType, + createdItem, + serverThroughputControlGroup.getGroupName()).toString(); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + serverThroughputControlGroup.getGroupName(), + String.format( + "(name=%s, default=%s, priorityLevel=%s, throughputBucket=%s)", + serverThroughputControlGroup.getGroupName(), + false, + PriorityLevel.LOW, + 3) + ); + } + + @Test(groups = {"emulator"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT) + public void throughputControl_LocalAndServer_requestOptions(OperationType operationType) { + this.ensureContainer(); + + // This test is verify that SDK throughput control group and server throughput control group can be enabled at the same time + + // The create document in this test usually takes around 6.29RU, pick a RU here relatively close, so to test throttled scenario + ThroughputControlGroupConfig groupConfig = + new ThroughputControlGroupConfigBuilder() + .groupName("group-sdk" + UUID.randomUUID()) + .targetThroughput(6) + .build(); + container.enableLocalThroughputControlGroup(groupConfig); + + ThroughputControlGroupConfig serverGroupConfig = + new ThroughputControlGroupConfigBuilder() + .groupName("group-server" + UUID.randomUUID()) + .throughputBucket(3) + .build(); + container.enableServerThroughputControlGroup(serverGroupConfig); + + CosmosItemRequestOptions requestOptions = new CosmosItemRequestOptions(); + requestOptions.setContentResponseOnWriteEnabled(true); + requestOptions.setThroughputControlGroupName(groupConfig.getGroupName()); + + CosmosItemResponse createItemResponse = container.createItem(getDocumentDefinition(), requestOptions).block(); + TestItem createdItem = createItemResponse.getItem(); + + String cosmosDiagnosticsString = createItemResponse.getDiagnostics().toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, 6) + ); + + // second request to group-1. which will get throttled + cosmosDiagnosticsString = + performDocumentOperation(this.container, operationType, createdItem, groupConfig.getGroupName()).toString(); + this.validateRequestThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + groupConfig.getGroupName(), + String.format("(name=%s, default=%s, targetRU=%s)", groupConfig.getGroupName(), false, 6) + ); + + // third request to server group, which will not get throttled + cosmosDiagnosticsString = + performDocumentOperation(this.container, operationType, createdItem, serverGroupConfig.getGroupName()).toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + this.validateThroughputControlDiagnostics( + cosmosDiagnosticsString, + serverGroupConfig.getGroupName(), + String.format("(name=%s, default=%s, throughputBucket=%s)", serverGroupConfig.getGroupName(), false, 3) + ); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void throughputControlDefaultGroup_LocalAndServer_requestOptions() { + this.ensureContainer(); + + // This test is verify that only one default throughput control group can be defined across sdk and server control group + ThroughputControlGroupConfig groupConfig = + new ThroughputControlGroupConfigBuilder() + .groupName("group-sdk" + UUID.randomUUID()) + .targetThroughput(6) + .defaultControlGroup(true) + .build(); + container.enableLocalThroughputControlGroup(groupConfig); + + ThroughputControlGroupConfig serverGroupConfig = + new ThroughputControlGroupConfigBuilder() + .groupName("group-server" + UUID.randomUUID()) + .throughputBucket(3) + .defaultControlGroup(true) + .build(); + + assertThatThrownBy( + () ->container.enableServerThroughputControlGroup(serverGroupConfig)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("A default group already exists"); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void throughputControl_noThroughputControlGroupEnabled() { + this.ensureContainer(); + + // This test is verify that request will succeed if no throughput control group defined + String throughputControlGroupName = "throughputControlGroup-" + UUID.randomUUID(); + CosmosItemRequestOptions itemRequestOptions = new CosmosItemRequestOptions(); + itemRequestOptions.setThroughputControlGroupName(throughputControlGroupName); + String cosmosDiagnosticsString = + container + .createItem(getDocumentDefinition(), itemRequestOptions) + .block() + .getDiagnostics() + .toString(); + this.validateRequestNotThrottled(cosmosDiagnosticsString); + assertThat(cosmosDiagnosticsString).contains("requestTCG"); + assertThat(cosmosDiagnosticsString).contains(throughputControlGroupName); + assertThat(cosmosDiagnosticsString).doesNotContain("requestTCGConfig"); + } + @BeforeClass(groups = { "emulator" }, timeOut = 4 * SETUP_TIMEOUT) - public void before_ThroughputBudgetControllerTest() { + public void before_ThroughputControllerTest() { this.ensureContainer(); } @@ -698,7 +1012,7 @@ private void ensureContainer() { } @AfterClass(groups = {"emulator"}, timeOut = TIMEOUT, alwaysRun = true) - public void after_ThroughputBudgetControllerTest() { + public void after_ThroughputControllerTest() { safeClose(this.client); } @@ -714,18 +1028,30 @@ private static TestItem getDocumentDefinition(String partitionKey) { ); } - private void validateRequestThrottled(String cosmosDiagnostics, ConnectionMode connectionMode) { + private void validateRequestThrottled(String cosmosDiagnostics) { assertThat(cosmosDiagnostics).isNotEmpty(); assertThat(cosmosDiagnostics).contains("\"statusCode\":429"); assertThat(cosmosDiagnostics).contains("\"subStatusCode\":10003"); } - private void validateRequestNotThrottled(String cosmosDiagnostics, ConnectionMode connectionMode) { + private void validateRequestNotThrottled(String cosmosDiagnostics) { assertThat(cosmosDiagnostics).isNotEmpty(); assertThat(cosmosDiagnostics).doesNotContain("\"statusCode\":429"); assertThat(cosmosDiagnostics).doesNotContain("\"subStatusCode\":10003"); } + private void validateThroughputControlDiagnostics( + String cosmosDiagnostics, + String groupName, + String groupConfig) { + + assertThat(cosmosDiagnostics).isNotEmpty(); + assertThat(cosmosDiagnostics).contains("requestTCG"); + assertThat(cosmosDiagnostics).contains(groupName); + assertThat(cosmosDiagnostics).contains("requestTCGConfig"); + assertThat(cosmosDiagnostics).contains(groupConfig); + } + private CosmosDiagnostics performDocumentOperation( CosmosAsyncContainer cosmosAsyncContainer, OperationType operationType, diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ContainerThroughputControlGroupPropertiesTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerSDKThroughputControlGroupPropertiesTests.java similarity index 92% rename from sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ContainerThroughputControlGroupPropertiesTests.java rename to sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerSDKThroughputControlGroupPropertiesTests.java index 9c8965c0610d..a4a3e94fd25e 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ContainerThroughputControlGroupPropertiesTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerSDKThroughputControlGroupPropertiesTests.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl; +package com.azure.cosmos.implementation.throughputControl.sdk; import com.azure.cosmos.CosmosAsyncClient; import com.azure.cosmos.CosmosAsyncContainer; @@ -9,7 +9,7 @@ import com.azure.cosmos.implementation.RxDocumentServiceRequest; import com.azure.cosmos.implementation.TestConfigurations; import com.azure.cosmos.implementation.apachecommons.lang.tuple.Pair; -import com.azure.cosmos.implementation.throughputControl.config.LocalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.sdk.config.LocalThroughputControlGroup; import com.azure.cosmos.models.PriorityLevel; import org.mockito.Mockito; import org.testng.annotations.Test; @@ -19,7 +19,7 @@ import static org.assertj.core.api.AssertionsForClassTypes.assertThat; import static org.assertj.core.api.AssertionsForClassTypes.assertThatThrownBy; -public class ContainerThroughputControlGroupPropertiesTests { +public class ContainerSDKThroughputControlGroupPropertiesTests { @Test(groups = "emulator") public void enableThroughputControlGroup() { @@ -30,8 +30,8 @@ public void enableThroughputControlGroup() { .key(TestConfigurations.MASTER_KEY) .buildAsyncClient(); - ContainerThroughputControlGroupProperties throughputControlContainerProperties = - new ContainerThroughputControlGroupProperties("/testDB/testContainer"); + ContainerSDKThroughputControlGroupProperties throughputControlContainerProperties = + new ContainerSDKThroughputControlGroupProperties("/testDB/testContainer"); CosmosAsyncContainer container = testClient.getDatabase("fakeDatabase").getContainer("fakeContainer"); @@ -160,8 +160,8 @@ public void enableThroughputControlGroupWithoutDefault() { .key(TestConfigurations.MASTER_KEY) .buildAsyncClient(); - ContainerThroughputControlGroupProperties throughputControlContainerProperties = - new ContainerThroughputControlGroupProperties("/testDB/testContainer"); + ContainerSDKThroughputControlGroupProperties throughputControlContainerProperties = + new ContainerSDKThroughputControlGroupProperties("/testDB/testContainer"); // Test: Without default group and request not having the group name, allowRequestToContinueOnInitError // should not throw NPE diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationTokenSourceTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/LinkedCancellationTokenSourceTests.java similarity index 94% rename from sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationTokenSourceTests.java rename to sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/LinkedCancellationTokenSourceTests.java index e505391a7c18..8a1cc2686a0c 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationTokenSourceTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/LinkedCancellationTokenSourceTests.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl; +package com.azure.cosmos.implementation.throughputControl.sdk; import org.testng.annotations.Test; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupConfigConfigurationTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputControlGroupConfigConfigurationTests.java similarity index 96% rename from sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupConfigConfigurationTests.java rename to sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputControlGroupConfigConfigurationTests.java index 597286a9883f..61488d7f9d66 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupConfigConfigurationTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputControlGroupConfigConfigurationTests.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl; +package com.azure.cosmos.implementation.throughputControl.sdk; import com.azure.cosmos.CosmosAsyncClient; import com.azure.cosmos.CosmosAsyncContainer; @@ -19,7 +19,6 @@ import java.time.Duration; import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.assertThatExceptionOfType; import static org.assertj.core.api.Assertions.assertThatThrownBy; public class ThroughputControlGroupConfigConfigurationTests extends TestSuiteBase { @@ -97,7 +96,7 @@ public void validatePriorityThroughputControlGroupWithThreshold() { public void validateThroughputControlGroupMalformed() { assertThatThrownBy(() -> new ThroughputControlGroupConfigBuilder().groupName("test").build()) .isInstanceOf(IllegalArgumentException.class) - .hasMessage("All targetThroughput, targetThroughputThreshold and priorityLevel cannot be null or empty."); + .hasMessage("All targetThroughput, targetThroughputThreshold, priorityLevel and throughput bucket cannot be null or empty."); } @BeforeClass(groups = { "emulator" }, timeOut = 4 * SETUP_TIMEOUT) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputRequestThrottlerTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputRequestThrottlerTests.java similarity index 93% rename from sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputRequestThrottlerTests.java rename to sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputRequestThrottlerTests.java index 299801ef74fc..6b93d901d6b8 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/ThroughputRequestThrottlerTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputRequestThrottlerTests.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl; +package com.azure.cosmos.implementation.throughputControl.sdk; import com.azure.cosmos.implementation.DocumentServiceRequestContext; import com.azure.cosmos.implementation.HttpConstants; @@ -53,7 +53,7 @@ public void processRequest() { this.assertRequestThrottlerState(requestThrottler, availableThroughput, scheduledThroughput); // Request2: will get throttled since there is no available throughput - requestMock.requestContext.throughputControlCycleId = StringUtils.EMPTY; + requestMock.requestContext.throughputControlRequestContext.setThroughputControlCycleId(StringUtils.EMPTY); TestPublisher requestPublisher2 = TestPublisher.create(); StepVerifier.create(requestThrottler.processRequest(requestMock, requestPublisher2.mono())) .verifyError(RequestRateTooLargeException.class); @@ -66,7 +66,7 @@ public void processRequest() { assertThat(requestThrottler.getAvailableThroughput()).isEqualTo(availableThroughput); // Request 3: will get throttled since there is no available throughput - requestMock.requestContext.throughputControlCycleId = StringUtils.EMPTY; + requestMock.requestContext.throughputControlRequestContext.setThroughputControlCycleId(StringUtils.EMPTY); TestPublisher requestPublisher3 = TestPublisher.create(); StepVerifier.create(requestThrottler.processRequest(requestMock, requestPublisher3.mono())) .verifyErrorSatisfies((t) -> { @@ -87,7 +87,7 @@ public void processRequest() { Mockito.doReturn(mockHeaders).when(bulkRequestMock).getHeaders(); bulkRequestMock.requestContext = new DocumentServiceRequestContext(); - bulkRequestMock.requestContext.throughputControlCycleId = StringUtils.EMPTY; + bulkRequestMock.requestContext.throughputControlRequestContext.setThroughputControlCycleId(StringUtils.EMPTY); TestPublisher requestPublisher4 = TestPublisher.create(); StepVerifier.create(requestThrottler.processRequest(bulkRequestMock, requestPublisher4.mono())) .verifyErrorSatisfies((t) -> { @@ -105,7 +105,7 @@ public void processRequest() { assertThat(requestThrottler.getAvailableThroughput()).isEqualTo(availableThroughput); // Request 5: will pass the request, and record the charge from exception - requestMock.requestContext.throughputControlCycleId = StringUtils.EMPTY; + requestMock.requestContext.throughputControlRequestContext.setThroughputControlCycleId(StringUtils.EMPTY); NotFoundException notFoundException = Mockito.mock(NotFoundException.class); Mockito.doReturn(requestChargePerRequest).when(notFoundException).getRequestCharge(); TestPublisher requestPublisher5 = TestPublisher.create(); @@ -137,7 +137,7 @@ public void responseOutOfCycle() { TestPublisher requestPublisher1 = TestPublisher.create(); StepVerifier.create(requestThrottler.processRequest(requestMock, requestPublisher1.mono())) .then(() -> { - requestMock.requestContext.throughputControlCycleId = UUID.randomUUID().toString(); + requestMock.requestContext.throughputControlRequestContext.setThroughputControlCycleId(UUID.randomUUID().toString()); requestPublisher1.emit(responseMock); }) .expectNext(responseMock) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/GlobalThroughputRequestControllerTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/GlobalThroughputRequestControllerTests.java similarity index 92% rename from sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/GlobalThroughputRequestControllerTests.java rename to sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/GlobalThroughputRequestControllerTests.java index 527ab40f6211..364c83499dd3 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/GlobalThroughputRequestControllerTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/GlobalThroughputRequestControllerTests.java @@ -1,15 +1,15 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller; +package com.azure.cosmos.implementation.throughputControl.sdk.controller; import com.azure.cosmos.implementation.DocumentServiceRequestContext; import com.azure.cosmos.implementation.OperationType; import com.azure.cosmos.implementation.RxDocumentServiceRequest; import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils; import com.azure.cosmos.implementation.directconnectivity.StoreResponse; -import com.azure.cosmos.implementation.throughputControl.ThroughputRequestThrottler; -import com.azure.cosmos.implementation.throughputControl.controller.request.GlobalThroughputRequestController; +import com.azure.cosmos.implementation.throughputControl.sdk.ThroughputRequestThrottler; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.request.GlobalThroughputRequestController; import org.mockito.Mockito; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/LocalThroughputControllerTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/LocalThroughputControllerTests.java similarity index 95% rename from sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/LocalThroughputControllerTests.java rename to sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/LocalThroughputControllerTests.java index 79680b6891b1..60202d2622a8 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/LocalThroughputControllerTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/LocalThroughputControllerTests.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller; +package com.azure.cosmos.implementation.throughputControl.sdk.controller; import com.azure.cosmos.ConnectionMode; import com.azure.cosmos.CosmosAsyncContainer; @@ -11,8 +11,8 @@ import com.azure.cosmos.implementation.directconnectivity.StoreResponse; import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper; import com.azure.cosmos.implementation.routing.Range; -import com.azure.cosmos.implementation.throughputControl.config.LocalThroughputControlGroup; -import com.azure.cosmos.implementation.throughputControl.controller.group.local.LocalThroughputControlGroupController; +import com.azure.cosmos.implementation.throughputControl.sdk.config.LocalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.group.local.LocalThroughputControlGroupController; import com.azure.cosmos.models.PriorityLevel; import org.mockito.Mockito; import org.testng.annotations.BeforeClass; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/PkRangesThroughputRequestControllerTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/PkRangesThroughputRequestControllerTests.java similarity index 95% rename from sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/PkRangesThroughputRequestControllerTests.java rename to sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/PkRangesThroughputRequestControllerTests.java index ba454eaf7372..fcc8f6c6bba2 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/controller/PkRangesThroughputRequestControllerTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/PkRangesThroughputRequestControllerTests.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller; +package com.azure.cosmos.implementation.throughputControl.sdk.controller; import com.azure.cosmos.implementation.DocumentServiceRequestContext; import com.azure.cosmos.implementation.OperationType; @@ -13,8 +13,8 @@ import com.azure.cosmos.implementation.directconnectivity.StoreResponse; import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper; import com.azure.cosmos.implementation.routing.Range; -import com.azure.cosmos.implementation.throughputControl.ThroughputRequestThrottler; -import com.azure.cosmos.implementation.throughputControl.controller.request.PkRangesThroughputRequestController; +import com.azure.cosmos.implementation.throughputControl.sdk.ThroughputRequestThrottler; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.request.PkRangesThroughputRequestController; import org.mockito.Mockito; import org.testng.annotations.BeforeClass; import org.testng.annotations.BeforeMethod; @@ -31,7 +31,6 @@ import java.util.concurrent.ConcurrentHashMap; import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.InstanceOfAssertFactories.OPTIONAL; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/server/ContainerServerThroughputControlGroupPropertiesTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/server/ContainerServerThroughputControlGroupPropertiesTests.java new file mode 100644 index 000000000000..e52f18d16598 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/server/ContainerServerThroughputControlGroupPropertiesTests.java @@ -0,0 +1,106 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl.server; + +import com.azure.cosmos.CosmosAsyncClient; +import com.azure.cosmos.CosmosAsyncContainer; +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.implementation.TestConfigurations; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroup; +import com.azure.cosmos.models.PriorityLevel; +import org.testng.annotations.Test; + +import java.util.UUID; + +import static org.assertj.core.api.AssertionsForClassTypes.assertThat; +import static org.assertj.core.api.AssertionsForClassTypes.assertThatThrownBy; + +public class ContainerServerThroughputControlGroupPropertiesTests { + @Test(groups = "emulator") + public void enableThroughputControlGroup() { + CosmosAsyncClient testClient = null; + try { + testClient = new CosmosClientBuilder() + .endpoint(TestConfigurations.HOST) + .key(TestConfigurations.MASTER_KEY) + .buildAsyncClient(); + + ContainerServerThroughputControlGroupProperties throughputControlContainerProperties = + new ContainerServerThroughputControlGroupProperties("/testDB/testContainer"); + + CosmosAsyncContainer container = testClient.getDatabase("fakeDatabase").getContainer("fakeContainer"); + + // Test 1: add default throughput control group successfully + ServerThroughputControlGroup throughputControlDefaultGroup = new ServerThroughputControlGroup( + "test-" + UUID.randomUUID(), + true, + PriorityLevel.HIGH, + 1, + container); + + Integer currentGroupSize = + throughputControlContainerProperties.enableThroughputControlGroup(throughputControlDefaultGroup); + assertThat(currentGroupSize).isEqualTo(1); + assertThat(throughputControlContainerProperties.hasDefaultGroup()).isTrue(); + assertThat(throughputControlContainerProperties.hasGroup(throughputControlDefaultGroup.getGroupName())).isTrue(); + + // Test 2: add another default throughput control group + ServerThroughputControlGroup throughputControlDefaultGroup2 = new ServerThroughputControlGroup( + "test-" + UUID.randomUUID(), + true, + PriorityLevel.HIGH, + 2, + container); + + assertThatThrownBy( + () -> throughputControlContainerProperties + .enableThroughputControlGroup(throughputControlDefaultGroup2)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("A default group already exists"); + assertThat(throughputControlContainerProperties.hasGroup(throughputControlDefaultGroup2.getGroupName())).isFalse(); + + // Test 3: add a new non-default group + ServerThroughputControlGroup throughputControlGroup1 = new ServerThroughputControlGroup( + "test-" + UUID.randomUUID(), + false, + PriorityLevel.HIGH, + 1, + container); + + currentGroupSize = throughputControlContainerProperties.enableThroughputControlGroup(throughputControlGroup1); + assertThat(currentGroupSize).isEqualTo(2); + assertThat(throughputControlContainerProperties.hasGroup(throughputControlGroup1.getGroupName())).isFalse(); + + // Test 4: add a throughput control group with name as throughputControlGroup1 but with different config + ServerThroughputControlGroup throughputControlGroup2 = new ServerThroughputControlGroup( + throughputControlGroup1.getGroupName(), + false, + PriorityLevel.HIGH, + 2, + container); + assertThatThrownBy( + () -> throughputControlContainerProperties + .enableThroughputControlGroup(throughputControlGroup2)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("A group with same name already exists, name: " + throughputControlGroup1.getGroupName()); + assertThat(throughputControlContainerProperties.hasGroup(throughputControlGroup2.getGroupName())).isFalse(); + + // Test 5: add a throughput control group with same config as throughputControlGroup1, to verify no errors will be thrown + ServerThroughputControlGroup throughputControlGroup3 = new ServerThroughputControlGroup( + throughputControlGroup1.getGroupName(), + throughputControlGroup1.isDefault(), + throughputControlGroup1.getPriorityLevel(), + throughputControlGroup1.getThroughputBucket(), + container); + + currentGroupSize = throughputControlContainerProperties.enableThroughputControlGroup(throughputControlGroup3); + assertThat(currentGroupSize).isEqualTo(2); + + } finally { + if (testClient != null) { + testClient.close(); + } + } + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ClientRetryPolicyE2ETests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ClientRetryPolicyE2ETests.java index 6869aeda23e6..5144d7bbc47d 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ClientRetryPolicyE2ETests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ClientRetryPolicyE2ETests.java @@ -14,6 +14,7 @@ import com.azure.cosmos.CosmosException; import com.azure.cosmos.DirectConnectionConfig; import com.azure.cosmos.FlakyTestRetryAnalyzer; +import com.azure.cosmos.TestObject; import com.azure.cosmos.implementation.AsyncDocumentClient; import com.azure.cosmos.implementation.DatabaseAccount; import com.azure.cosmos.implementation.DatabaseAccountLocation; @@ -21,7 +22,6 @@ import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.OperationType; import com.azure.cosmos.implementation.Utils; -import com.azure.cosmos.implementation.throughputControl.TestItem; import com.azure.cosmos.models.CosmosBatch; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; import com.azure.cosmos.models.CosmosItemIdentity; @@ -196,7 +196,7 @@ public static Object[] preferredRegionsConfigProvider() { @Test(groups = { "multi-master" }, dataProvider = "preferredRegionsConfigProvider", timeOut = TIMEOUT) public void queryPlanHttpTimeoutWillNotMarkRegionUnavailable(boolean shouldUsePreferredRegionsOnClient) { - TestItem newItem = TestItem.createNewItem(); + TestObject newItem = TestObject.create(); CosmosAsyncContainer resultantCosmosAsyncContainer; CosmosAsyncClient resultantCosmosAsyncClient; @@ -239,7 +239,7 @@ public void queryPlanHttpTimeoutWillNotMarkRegionUnavailable(boolean shouldUsePr try { // validate the query plan will be retried in a different region and the final requests will be succeeded // TODO: Also capture all retries for metadata requests in the diagnostics - FeedResponse firstPage = cosmosAsyncContainerFromClientWithPreferredRegions.queryItems(query, queryRequestOptions, TestItem.class) + FeedResponse firstPage = cosmosAsyncContainerFromClientWithPreferredRegions.queryItems(query, queryRequestOptions, TestObject.class) .byPage() .blockFirst(); @@ -274,7 +274,7 @@ public void addressRefreshHttpTimeoutWillDoCrossRegionRetryForReads(boolean shou throw new SkipException("queryPlanHttpTimeoutWillNotMarkRegionUnavailable() is only meant for DIRECT mode"); } - TestItem newItem = TestItem.createNewItem(); + TestObject newItem = TestObject.create(); resultantCosmosAsyncContainer.createItem(newItem).block(); // create fault injection rules for address refresh @@ -309,8 +309,8 @@ public void addressRefreshHttpTimeoutWillDoCrossRegionRetryForReads(boolean shou resultantCosmosAsyncContainer, Arrays.asList(addressRefreshDelayRule, serverGoneRule)).block(); try { - CosmosItemResponse itemResponse = resultantCosmosAsyncContainer - .readItem(newItem.getId(), new PartitionKey(newItem.getId()), TestItem.class) + CosmosItemResponse itemResponse = resultantCosmosAsyncContainer + .readItem(newItem.getId(), new PartitionKey(newItem.getId()), TestObject.class) .block(); assertThat(itemResponse).isNotNull(); @@ -379,7 +379,7 @@ public void addressRefreshHttpTimeoutWillNotDoCrossRegionRetryForWrites(boolean resultantCosmosAsyncContainer, Arrays.asList(addressRefreshDelayRule, serverGoneRule)).block(); try { - TestItem newItem = TestItem.createNewItem(); + TestObject newItem = TestObject.create(); resultantCosmosAsyncContainer.createItem(newItem).block(); } catch (CosmosException e) { assertThat(e.getDiagnostics().getContactedRegionNames().size()).isEqualTo(1); @@ -417,7 +417,7 @@ public void dataPlaneRequestHttpTimeout( throw new SkipException("queryPlanHttpTimeoutWillNotMarkRegionUnavailable() is only meant for GATEWAY mode"); } - TestItem newItem = TestItem.createNewItem(); + TestObject newItem = TestObject.create(); resultantCosmosAsyncContainer.createItem(newItem).block(); FaultInjectionRule requestHttpTimeoutRule = new FaultInjectionRuleBuilder("requestHttpTimeoutRule" + UUID.randomUUID()) .condition( @@ -506,7 +506,7 @@ public void dataPlaneRequestHitsLeaseNotFoundInFirstPreferredRegion( throw new SkipException("leaseNotFound is only meant for Direct mode"); } - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); FaultInjectionRule leaseNotFoundFaultRule = new FaultInjectionRuleBuilder("leaseNotFound-" + UUID.randomUUID()) .condition( @@ -619,7 +619,7 @@ public void channelAcquisitionExceptionOnWrites( CosmosFaultInjectionHelper.configureFaultInjectionRules(testContainer, Arrays.asList(channelAcquisitionExceptionRule)).block(); try { - TestItem createdItem = TestItem.createNewItem(); + TestObject createdItem = TestObject.create(); testContainer.createItem(createdItem).block(); // using a higher concurrency to force channelAcquisitionException to happen @@ -683,21 +683,21 @@ private boolean isChannelAcquisitionExceptionTriggeredRegionRetryExists(String c private Mono performDocumentOperation( CosmosAsyncContainer cosmosAsyncContainer, OperationType operationType, - TestItem createdItem, - Function extractPartitionKeyFunc, + TestObject createdItem, + Function extractPartitionKeyFunc, boolean isReadMany) { try { if (operationType == OperationType.Query && isReadMany) { CosmosQueryRequestOptions queryRequestOptions = new CosmosQueryRequestOptions(); String query = String.format("SELECT * from c where c.id = '%s'", createdItem.getId()); - FeedResponse itemFeedResponse = - cosmosAsyncContainer.queryItems(query, queryRequestOptions, TestItem.class).byPage().blockFirst(); + FeedResponse itemFeedResponse = + cosmosAsyncContainer.queryItems(query, queryRequestOptions, TestObject.class).byPage().blockFirst(); return Mono.just(itemFeedResponse.getCosmosDiagnostics()) .onErrorResume(throwable -> { if (throwable instanceof CosmosException) { CosmosException cosmosException = (CosmosException) throwable; - + return Mono.just(cosmosException.getDiagnostics()); } return Mono.error(throwable); @@ -717,7 +717,7 @@ private Mono performDocumentOperation( .readItem( createdItem.getId(), extractPartitionKeyFunc.apply(createdItem), - TestItem.class + TestObject.class ) .map(itemResponse -> itemResponse.getDiagnostics()) .onErrorResume(throwable -> { @@ -759,7 +759,7 @@ private Mono performDocumentOperation( } if (operationType == OperationType.Create) { - return cosmosAsyncContainer.createItem(TestItem.createNewItem()).map(itemResponse -> itemResponse.getDiagnostics()) .onErrorResume(throwable -> { + return cosmosAsyncContainer.createItem(TestObject.create()).map(itemResponse -> itemResponse.getDiagnostics()) .onErrorResume(throwable -> { if (throwable instanceof CosmosException) { CosmosException cosmosException = (CosmosException) throwable; @@ -770,7 +770,7 @@ private Mono performDocumentOperation( } if (operationType == OperationType.Upsert) { - return cosmosAsyncContainer.upsertItem(TestItem.createNewItem()).map(itemResponse -> itemResponse.getDiagnostics()) .onErrorResume(throwable -> { + return cosmosAsyncContainer.upsertItem(TestObject.create()).map(itemResponse -> itemResponse.getDiagnostics()) .onErrorResume(throwable -> { if (throwable instanceof CosmosException) { CosmosException cosmosException = (CosmosException) throwable; @@ -788,7 +788,7 @@ private Mono performDocumentOperation( createdItem.getId(), extractPartitionKeyFunc.apply(createdItem), patchOperations, - TestItem.class) + TestObject.class) .map(itemResponse -> itemResponse.getDiagnostics()) .onErrorResume(throwable -> { if (throwable instanceof CosmosException) { @@ -823,8 +823,8 @@ private Mono performDocumentOperation( CosmosChangeFeedRequestOptions changeFeedRequestOptions = CosmosChangeFeedRequestOptions.createForProcessingFromBeginning(feedRanges.get(0)); - FeedResponse firstPage = cosmosAsyncContainer - .queryChangeFeed(changeFeedRequestOptions, TestItem.class) + FeedResponse firstPage = cosmosAsyncContainer + .queryChangeFeed(changeFeedRequestOptions, TestObject.class) .byPage() .blockFirst(); return Mono.just(firstPage.getCosmosDiagnostics()); @@ -834,7 +834,7 @@ private Mono performDocumentOperation( return cosmosAsyncContainer .readMany( Arrays.asList(new CosmosItemIdentity(extractPartitionKeyFunc.apply(createdItem), createdItem.getId()), new CosmosItemIdentity(extractPartitionKeyFunc.apply(createdItem), createdItem.getId())), - TestItem.class) + TestObject.class) .map(itemResponse -> itemResponse.getCosmosDiagnostics()) .onErrorResume(throwable -> { if (throwable instanceof CosmosException) { diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/FullFidelityChangeFeedTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/FullFidelityChangeFeedTest.java index 2e43dcdfbb26..e26765752927 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/FullFidelityChangeFeedTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/FullFidelityChangeFeedTest.java @@ -6,7 +6,7 @@ import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.CosmosAsyncDatabase; import com.azure.cosmos.CosmosClientBuilder; -import com.azure.cosmos.implementation.throughputControl.TestItem; +import com.azure.cosmos.TestObject; import com.azure.cosmos.models.ChangeFeedPolicy; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; import com.azure.cosmos.models.CosmosContainerProperties; @@ -21,6 +21,7 @@ import org.testng.annotations.Test; import java.time.Duration; +import java.util.Arrays; import java.util.Iterator; import java.util.List; import java.util.UUID; @@ -60,22 +61,22 @@ public void fullFidelityChangeFeed_FromNowForLogicalPartition() throws Exception continuationToken1 = response.getContinuationToken(); } - TestItem item1 = new TestItem( + TestObject item1 = new TestObject( UUID.randomUUID().toString(), - "mypk-1", "Johnson"); - TestItem item2 = new TestItem( + "mypk-1", Arrays.asList(), "Johnson"); + TestObject item2 = new TestObject( UUID.randomUUID().toString(), - "mypk-1", "Smith"); - TestItem item3 = new TestItem( + "mypk-1", Arrays.asList(), "Smith"); + TestObject item3 = new TestObject( UUID.randomUUID().toString(), - "mypk-2", "John"); + "mypk-2", Arrays.asList(), "John"); cosmosContainer.createItem(item1).block(); cosmosContainer.createItem(item2).block(); - String originalLastNameItem1 = item1.getProp(); - item1.setProp("Gates"); + String originalLastNameItem1 = item1.getStringProp(); + item1.setStringProp("Gates"); cosmosContainer.upsertItem(item1).block(); - String originalLastNameItem2 = item2.getProp(); - item2.setProp("Doe"); + String originalLastNameItem2 = item2.getStringProp(); + item2.setStringProp("Doe"); cosmosContainer.upsertItem(item2).block(); cosmosContainer.deleteItem(item1, new CosmosItemRequestOptions()).block(); @@ -102,14 +103,14 @@ public void fullFidelityChangeFeed_FromNowForLogicalPartition() throws Exception assertThat(itemChanges.get(1).get("metadata").get("operationType").asText()).isEqualTo("create"); // Assert replace of item1 assertThat(itemChanges.get(2).get("current").get("id").asText()).isEqualTo(item1.getId()); - assertThat(itemChanges.get(2).get("current").get("prop").asText()).isEqualTo(item1.getProp()); + assertThat(itemChanges.get(2).get("current").get("prop").asText()).isEqualTo(item1.getStringProp()); assertThat(itemChanges.get(2).get("metadata").get("operationType").asText()).isEqualTo("replace"); if (itemChanges.get(2).get("previous") != null) { assertThat(itemChanges.get(2).get("previous")).isEqualTo(itemChanges.get(0).get("current")); } // Assert replace of item2 assertThat(itemChanges.get(3).get("current").get("id").asText()).isEqualTo(item2.getId()); - assertThat(itemChanges.get(3).get("current").get("prop").asText()).isEqualTo(item2.getProp()); + assertThat(itemChanges.get(3).get("current").get("prop").asText()).isEqualTo(item2.getStringProp()); assertThat(itemChanges.get(3).get("metadata").get("operationType").asText()).isEqualTo("replace"); if (itemChanges.get(3).get("previous") != null) { assertThat(itemChanges.get(3).get("previous")).isEqualTo(itemChanges.get(1).get("current")); @@ -141,8 +142,8 @@ public void fullFidelityChangeFeed_FromNowForLogicalPartition() throws Exception } cosmosContainer.createItem(item3).block(); - String originalLastNameItem3 = item3.getProp(); - item3.setProp("Potter"); + String originalLastNameItem3 = item3.getStringProp(); + item3.setStringProp("Potter"); cosmosContainer.upsertItem(item3).block(); cosmosContainer.deleteItem(item3, new CosmosItemRequestOptions()).block(); @@ -166,7 +167,7 @@ public void fullFidelityChangeFeed_FromNowForLogicalPartition() throws Exception assertThat(itemChanges.get(0).get("metadata").get("operationType").asText()).isEqualTo("create"); // Assert replace of item3 assertThat(itemChanges.get(1).get("current").get("id").asText()).isEqualTo(item3.getId()); - assertThat(itemChanges.get(1).get("current").get("prop").asText()).isEqualTo(item3.getProp()); + assertThat(itemChanges.get(1).get("current").get("prop").asText()).isEqualTo(item3.getStringProp()); assertThat(itemChanges.get(1).get("metadata").get("operationType").asText()).isEqualTo("replace"); if (itemChanges.get(1).get("previous") != null) { assertThat(itemChanges.get(1).get("previous")).isEqualTo(itemChanges.get(0).get("current")); @@ -209,16 +210,16 @@ public void fullFidelityChangeFeed_FromContinuationToken() throws Exception { .createForProcessingFromContinuation(continuationToken); options.allVersionsAndDeletes(); - TestItem item1 = new TestItem( + TestObject item1 = new TestObject( UUID.randomUUID().toString(), - "mypk", "Johnson"); - TestItem item2 = new TestItem( + "mypk", Arrays.asList(), "Johnson"); + TestObject item2 = new TestObject( UUID.randomUUID().toString(), - "mypk", "Smith"); + "mypk", Arrays.asList(), "Smith"); cosmosContainer.upsertItem(item1).block(); cosmosContainer.upsertItem(item2).block(); - String originalLastName = item1.getProp(); - item1.setProp("Gates"); + String originalLastName = item1.getStringProp(); + item1.setStringProp("Gates"); cosmosContainer.upsertItem(item1).block(); cosmosContainer.deleteItem(item1, new CosmosItemRequestOptions()).block(); @@ -237,11 +238,11 @@ public void fullFidelityChangeFeed_FromContinuationToken() throws Exception { assertThat(itemChanges.get(0).get("current").get("prop").asText()).isEqualTo(originalLastName); assertThat(itemChanges.get(0).get("metadata").get("operationType").asText()).isEqualTo("create"); assertThat(itemChanges.get(1).get("current").get("id").asText()).isEqualTo(item2.getId()); - assertThat(itemChanges.get(1).get("current").get("prop").asText()).isEqualTo(item2.getProp()); + assertThat(itemChanges.get(1).get("current").get("prop").asText()).isEqualTo(item2.getStringProp()); assertThat(itemChanges.get(1).get("metadata").get("operationType").asText()).isEqualTo("create"); // Assert replace of item1 assertThat(itemChanges.get(2).get("current").get("id").asText()).isEqualTo(item1.getId()); - assertThat(itemChanges.get(2).get("current").get("prop").asText()).isEqualTo(item1.getProp()); + assertThat(itemChanges.get(2).get("current").get("prop").asText()).isEqualTo(item1.getStringProp()); assertThat(itemChanges.get(2).get("metadata").get("operationType").asText()).isEqualTo("replace"); if (itemChanges.get(2).get("previous") != null) { assertThat(itemChanges.get(2).get("previous")).isEqualTo(itemChanges.get(0).get("current")); @@ -286,9 +287,9 @@ public void fullFidelityChangeFeed_FromContinuationTokenOperationsOrder() throws // Create, replace, and delete 50 objects for 150 total operations for (int i = 0; i < 50; i++) { - TestItem currentItem = new TestItem("item"+ i, "mypk", "Smith"); + TestObject currentItem = new TestObject("item"+ i, "mypk", Arrays.asList(), "Smith"); cosmosContainer.upsertItem(currentItem).block(); - currentItem.setProp("Jefferson"); + currentItem.setStringProp("Jefferson"); cosmosContainer.upsertItem(currentItem).block(); cosmosContainer.deleteItem(currentItem, new CosmosItemRequestOptions()).block(); } @@ -342,18 +343,18 @@ public void fullFidelityChangeFeed_VerifyPreviousPresentOnReplace() throws Excep .createForProcessingFromContinuation(continuationToken); options.allVersionsAndDeletes(); - TestItem item1 = new TestItem( + TestObject item1 = new TestObject( UUID.randomUUID().toString(), - "mypk", "Johnson"); + "mypk", Arrays.asList(), "Johnson"); cosmosContainer.upsertItem(item1).block(); - String originalLastName = item1.getProp(); - item1.setProp("Gates"); + String originalLastName = item1.getStringProp(); + item1.setStringProp("Gates"); cosmosContainer.upsertItem(item1).block(); - String secondLastName = item1.getProp(); - item1.setProp("DiCaprio"); - String thirdLastName = item1.getProp(); + String secondLastName = item1.getStringProp(); + item1.setStringProp("DiCaprio"); + String thirdLastName = item1.getStringProp(); cosmosContainer.upsertItem(item1).block(); - item1.setProp(originalLastName); + item1.setStringProp(originalLastName); cosmosContainer.upsertItem(item1).block(); cosmosContainer.deleteItem(item1, new CosmosItemRequestOptions()).block(); @@ -387,7 +388,7 @@ public void fullFidelityChangeFeed_VerifyPreviousPresentOnReplace() throws Excep assertThat(itemChanges.get(2).get("previous")).isEqualTo(itemChanges.get(1).get("current")); assertThat(itemChanges.get(3).get("previous").get("id").asText()).isEqualTo(item1.getId()); - assertThat(itemChanges.get(3).get("current").get("prop").asText()).isEqualTo(item1.getProp()); + assertThat(itemChanges.get(3).get("current").get("prop").asText()).isEqualTo(item1.getStringProp()); assertThat(itemChanges.get(3).get("metadata").get("operationType").asText()).isEqualTo("replace"); assertThat(itemChanges.get(3).get("metadata").get("previousImageLSN").asText() ).isEqualTo(itemChanges.get(2).get("metadata").get("lsn").asText()); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index 6736cd186e18..680b98a6a915 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java @@ -29,7 +29,8 @@ import com.azure.cosmos.implementation.clienttelemetry.TagName; import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdMetrics; import com.azure.cosmos.implementation.faultinjection.IFaultInjectorProvider; -import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroup; import com.azure.cosmos.models.CosmosAuthorizationTokenResolver; import com.azure.cosmos.models.CosmosClientTelemetryConfig; import com.azure.cosmos.models.CosmosContainerIdentity; @@ -572,9 +573,19 @@ DiagnosticsProvider getDiagnosticsProvider() { * @param group Throughput control group going to be enabled. * @param throughputQueryMono The throughput query mono. */ - void enableThroughputControlGroup(ThroughputControlGroupInternal group, Mono throughputQueryMono) { + void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono throughputQueryMono) { checkNotNull(group, "Throughput control group cannot be null"); - this.asyncDocumentClient.enableThroughputControlGroup(group, throughputQueryMono); + this.asyncDocumentClient.enableSDKThroughputControlGroup(group, throughputQueryMono); + } + + /*** + * Enable server throughput control group. + * + * @param group the server throughput control group. + */ + void enableServerThroughputControlGroup(ServerThroughputControlGroup group) { + checkNotNull(group, "Argument 'group' can not be null"); + this.asyncDocumentClient.enableServerThroughputControlGroup(group); } /*** diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index fdbaa4b8af38..2fcf9bfe50cf 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -37,9 +37,10 @@ import com.azure.cosmos.implementation.feedranges.FeedRangeInternal; import com.azure.cosmos.implementation.routing.PartitionKeyInternal; import com.azure.cosmos.implementation.routing.Range; -import com.azure.cosmos.implementation.throughputControl.config.GlobalThroughputControlGroup; -import com.azure.cosmos.implementation.throughputControl.config.LocalThroughputControlGroup; -import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupFactory; +import com.azure.cosmos.implementation.throughputControl.sdk.config.GlobalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.sdk.config.LocalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.ThroughputControlGroupFactory; import com.azure.cosmos.models.CosmosBatch; import com.azure.cosmos.models.CosmosBatchOperationResult; import com.azure.cosmos.models.CosmosBatchRequestOptions; @@ -70,6 +71,7 @@ import com.azure.cosmos.models.SqlQuerySpec; import com.azure.cosmos.models.ThroughputProperties; import com.azure.cosmos.models.ThroughputResponse; +import com.azure.cosmos.util.Beta; import com.azure.cosmos.util.CosmosPagedFlux; import com.azure.cosmos.util.UtilBridgeInternal; import org.slf4j.Logger; @@ -2769,7 +2771,7 @@ void enableLocalThroughputControlGroup( LocalThroughputControlGroup localControlGroup = ThroughputControlGroupFactory.createThroughputLocalControlGroup(groupConfig, this); - this.database.getClient().enableThroughputControlGroup(localControlGroup, throughputQueryMono); + this.database.getClient().enableSDKThroughputControlGroup(localControlGroup, throughputQueryMono); } /** @@ -2794,7 +2796,7 @@ void enableLocalThroughputControlGroup( * * * - * @param groupConfig The throughput control group configuration, see {@link GlobalThroughputControlGroup}. + * @param groupConfig The throughput control group configuration, see {@link ThroughputControlGroupConfig}. * @param globalControlConfig The global throughput control configuration, see {@link GlobalThroughputControlConfig}. */ public void enableGlobalThroughputControlGroup( @@ -2807,7 +2809,7 @@ public void enableGlobalThroughputControlGroup( /*** * Only used internally. *
- * @param groupConfig The throughput control group configuration, see {@link GlobalThroughputControlGroup}. + * @param groupConfig The throughput control group configuration, see {@link ThroughputControlGroupConfig}. * @param globalControlConfig The global throughput control configuration, see {@link GlobalThroughputControlConfig}. * @param throughputQueryMono The throughput query mono. */ @@ -2819,7 +2821,36 @@ void enableGlobalThroughputControlGroup( GlobalThroughputControlGroup globalControlGroup = ThroughputControlGroupFactory.createThroughputGlobalControlGroup(groupConfig, globalControlConfig, this); - this.database.getClient().enableThroughputControlGroup(globalControlGroup, throughputQueryMono); + this.database.getClient().enableSDKThroughputControlGroup(globalControlGroup, throughputQueryMono); + } + + /*** + * Enable the server throughput bucket control group. + * + * + * 
+     * ThroughputControlGroupConfig groupConfig =
+     *     new ThroughputControlGroupConfigBuilder()
+     *         .groupName("localControlGroup")
+     *         .throughputBucket(2)
+     *         .build();
+     *
+     * container.enableServerThroughputControlGroup(groupConfig);
+     *
+ * + * + * @param groupConfig the throughput control group config, see {@link ThroughputControlGroupConfig}. + */ + @Beta(value = Beta.SinceVersion.V4_74_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public void enableServerThroughputControlGroup(ThroughputControlGroupConfig groupConfig) { + if (groupConfig.getPriorityLevel() == null && groupConfig.getThroughputBucket() == null) { + throw new IllegalArgumentException("Config 'priorityLevel' and 'throughputBucket' can not be null for both."); + } + + ServerThroughputControlGroup serverThroughputControlGroup = + ThroughputControlGroupFactory.createServerThroughputControlGroup(groupConfig, this); + + this.database.getClient().enableServerThroughputControlGroup(serverThroughputControlGroup); } void configureFaultInjectionProvider(IFaultInjectorProvider injectorProvider) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java index b0cbc7c37b85..04a6060c1927 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java @@ -3,7 +3,7 @@ package com.azure.cosmos; -import com.azure.cosmos.implementation.throughputControl.config.GlobalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.sdk.config.GlobalThroughputControlGroup; import com.azure.cosmos.models.CosmosBatch; import com.azure.cosmos.models.CosmosBatchOperationResult; import com.azure.cosmos.models.CosmosBatchRequestOptions; @@ -28,6 +28,7 @@ import com.azure.cosmos.models.SqlQuerySpec; import com.azure.cosmos.models.ThroughputProperties; import com.azure.cosmos.models.ThroughputResponse; +import com.azure.cosmos.util.Beta; import com.azure.cosmos.util.CosmosPagedFlux; import com.azure.cosmos.util.CosmosPagedIterable; import org.slf4j.Logger; @@ -948,7 +949,6 @@ public CosmosScripts getScripts() { } // TODO: should make partitionkey public in CosmosAsyncItem and fix the below call - private CosmosPagedIterable getCosmosPagedIterable(CosmosPagedFlux cosmosPagedFlux) { return new CosmosPagedIterable<>(cosmosPagedFlux); } @@ -1030,6 +1030,28 @@ public void enableGlobalThroughputControlGroup(ThroughputControlGroupConfig grou this.asyncContainer.enableGlobalThroughputControlGroup(groupConfig, globalControlConfig); } + /*** + * Enable the server throughput control group. + * + * + * 
+     * ThroughputControlGroupConfig groupConfig =
+     *     new ThroughputControlGroupConfigBuilder()
+     *         .groupName("localControlGroup")
+     *         .throughputBucket(2)
+     *         .build();
+     *
+     * container.enableServerThroughputControlGroup(groupConfig);
+     *
+ * + * + * @param groupConfig the throughput control group config, see {@link ThroughputControlGroupConfig}. + */ + @Beta(value = Beta.SinceVersion.V4_74_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public void enableServerThroughputControlGroup(ThroughputControlGroupConfig groupConfig) { + this.asyncContainer.enableServerThroughputControlGroup(groupConfig); + } + /** * Initializes the container by warming up the caches and connections for the current read region. * diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfig.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfig.java index 81b8417666b5..b5f3cedac8fe 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfig.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ThroughputControlGroupConfig.java @@ -4,6 +4,7 @@ package com.azure.cosmos; import com.azure.cosmos.models.PriorityLevel; +import com.azure.cosmos.util.Beta; /** * Throughput control group configuration. @@ -13,6 +14,7 @@ public final class ThroughputControlGroupConfig { private final Integer targetThroughput; private final Double targetThroughputThreshold; private final PriorityLevel priorityLevel; + private final Integer throughputBucket; private final boolean isDefault; private final boolean continueOnInitError; @@ -21,12 +23,14 @@ public final class ThroughputControlGroupConfig { Integer targetThroughput, Double targetThroughputThreshold, PriorityLevel priorityLevel, + Integer throughputBucket, boolean isDefault, boolean continueOnInitError) { this.groupName = groupName; this.targetThroughput = targetThroughput; this.targetThroughputThreshold = targetThroughputThreshold; this.priorityLevel = priorityLevel; + this.throughputBucket = throughputBucket; this.isDefault = isDefault; this.continueOnInitError = continueOnInitError; } @@ -76,6 +80,20 @@ public Double getTargetThroughputThreshold() { */ public PriorityLevel getPriorityLevel() { return this.priorityLevel; } + + /*** + * Get the throughput bucket. + *

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

+ * For more information about throughput bucket please visit + * Throughput buckets in Azure Cosmos DB + * + * @param throughputBucket the throughput bucket id. + * @return The {@link ThroughputControlGroupConfigBuilder}. + */ + @Beta(value = Beta.SinceVersion.V4_74_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public ThroughputControlGroupConfigBuilder throughputBucket(int throughputBucket) { + checkArgument(throughputBucket >= 0, "Throughput bucket should be no smaller than 0"); + this.throughputBucket = throughputBucket; + return this; + } + /** * Set whether allow request to continue on original request flow if throughput control controller failed on initialization. * If set to true, requests will be able to fall back to original request flow if throughput control controller failed on initialization. @@ -178,9 +195,15 @@ public ThroughputControlGroupConfig build() { if (StringUtils.isEmpty(this.groupName)) { throw new IllegalArgumentException("Group name cannot be null nor empty"); } - if (this.targetThroughput == null && this.targetThroughputThreshold == null && this.priorityLevel == null) { - throw new IllegalArgumentException("All targetThroughput, targetThroughputThreshold and priorityLevel cannot be null or empty."); + + if (this.targetThroughput == null + && this.targetThroughputThreshold == null + && this.priorityLevel == null + && this.throughputBucket == null) { + throw new IllegalArgumentException( + "All targetThroughput, targetThroughputThreshold, priorityLevel and throughput bucket cannot be null or empty."); } + if (this.targetThroughput == null && this.targetThroughputThreshold == null) { this.targetThroughput = Integer.MAX_VALUE; } @@ -190,6 +213,7 @@ public ThroughputControlGroupConfig build() { this.targetThroughput, this.targetThroughputThreshold, this.priorityLevel, + this.throughputBucket, this.isDefault, this.continueOnInitError); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java index 957224b96304..f19ccb503027 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java @@ -21,7 +21,8 @@ import com.azure.cosmos.implementation.directconnectivity.AddressSelector; import com.azure.cosmos.implementation.faultinjection.IFaultInjectorProvider; import com.azure.cosmos.implementation.query.PartitionedQueryExecutionInfo; -import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroup; import com.azure.cosmos.models.CosmosAuthorizationTokenResolver; import com.azure.cosmos.models.CosmosBatchResponse; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; @@ -1631,11 +1632,18 @@ Flux> readAllDocuments( CosmosItemSerializer getEffectiveItemSerializer(CosmosItemSerializer requestOptionsItemSerializer); /** - * Enable throughput control group. + * Enable sdk throughput control group. * * @param group the throughput control group. */ - void enableThroughputControlGroup(ThroughputControlGroupInternal group, Mono throughputQueryMono); + void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono throughputQueryMono); + + /*** + * Enable server throughput control group. + * + * @param group the server throughput control group. + */ + void enableServerThroughputControlGroup(ServerThroughputControlGroup group); /** * Submits open connection tasks and warms up caches for replicas for containers specified by diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientSideRequestStatistics.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientSideRequestStatistics.java index 3503a46e11c1..fbfaf776edc8 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientSideRequestStatistics.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientSideRequestStatistics.java @@ -270,6 +270,8 @@ public void recordGatewayResponse( gatewayStatistics.faultInjectionRuleId = storeResponseDiagnostics.getFaultInjectionRuleId(); gatewayStatistics.faultInjectionEvaluationResults = storeResponseDiagnostics.getFaultInjectionEvaluationResults(); gatewayStatistics.endpoint = storeResponseDiagnostics.getEndpoint(); + gatewayStatistics.requestThroughputControlGroupName = storeResponseDiagnostics.getRequestThroughputControlGroupName(); + gatewayStatistics.requestThroughputControlGroupConfig = storeResponseDiagnostics.getRequestThroughputControlGroupConfig(); this.activityId = storeResponseDiagnostics.getActivityId() != null ? storeResponseDiagnostics.getActivityId() : rxDocumentServiceRequest.getActivityId().toString(); @@ -910,6 +912,8 @@ public static class GatewayStatistics { private PerPartitionCircuitBreakerInfoHolder perPartitionCircuitBreakerInfoHolder; private PerPartitionFailoverInfoHolder perPartitionFailoverInfoHolder; private String endpoint; + private String requestThroughputControlGroupName; + private String requestThroughputControlGroupConfig; public String getSessionToken() { return sessionToken; @@ -979,6 +983,14 @@ public String getEndpoint() { return this.endpoint; } + public String getRequestThroughputControlGroupName() { + return this.requestThroughputControlGroupName; + } + + public String getRequestThroughputControlGroupConfig() { + return this.requestThroughputControlGroupConfig; + } + public static class GatewayStatisticsSerializer extends StdSerializer { private static final long serialVersionUID = 1L; @@ -1016,6 +1028,8 @@ public void serialize(GatewayStatistics gatewayStatistics, this.writeNonNullObjectField(jsonGenerator, "perPartitionCircuitBreakerInfoHolder", gatewayStatistics.getPerPartitionCircuitBreakerInfoHolder()); this.writeNonNullObjectField(jsonGenerator, "perPartitionFailoverInfoHolder", gatewayStatistics.getPerPartitionFailoverInfoHolder()); + this.writeNonNullStringField(jsonGenerator, "requestTCG", gatewayStatistics.getRequestThroughputControlGroupName()); + this.writeNonNullStringField(jsonGenerator, "requestTCGConfig", gatewayStatistics.getRequestThroughputControlGroupConfig()); jsonGenerator.writeEndObject(); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DocumentServiceRequestContext.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DocumentServiceRequestContext.java index 426c7493091a..18da5250c458 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DocumentServiceRequestContext.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DocumentServiceRequestContext.java @@ -18,6 +18,7 @@ import com.azure.cosmos.implementation.directconnectivity.Uri; import com.azure.cosmos.implementation.routing.PartitionKeyInternal; import com.azure.cosmos.implementation.routing.RegionalRoutingContext; +import com.azure.cosmos.implementation.throughputControl.ThroughputControlRequestContext; import java.util.ArrayList; import java.util.List; @@ -54,7 +55,7 @@ public class DocumentServiceRequestContext implements Cloneable { public volatile PartitionKeyInternal effectivePartitionKey; public volatile CosmosDiagnostics cosmosDiagnostics; public volatile String resourcePhysicalAddress; - public volatile String throughputControlCycleId; + public ThroughputControlRequestContext throughputControlRequestContext; public volatile boolean replicaAddressValidationEnabled = Configs.isReplicaAddressValidationEnabled(); private final Set failedEndpoints = ConcurrentHashMap.newKeySet(); private CosmosEndToEndOperationLatencyPolicyConfig endToEndOperationLatencyPolicyConfig; @@ -132,6 +133,10 @@ public void addToFailedEndpoints(Exception exception, Uri address) { } } + public void setThroughputControlRequestContext(ThroughputControlRequestContext throughputControlRequestContext) { + this.throughputControlRequestContext = throughputControlRequestContext; + } + @Override public DocumentServiceRequestContext clone() { DocumentServiceRequestContext context = new DocumentServiceRequestContext(); @@ -158,7 +163,7 @@ public DocumentServiceRequestContext clone() { context.performedBackgroundAddressRefresh = this.performedBackgroundAddressRefresh; context.cosmosDiagnostics = this.cosmosDiagnostics; context.resourcePhysicalAddress = this.resourcePhysicalAddress; - context.throughputControlCycleId = this.throughputControlCycleId; + context.throughputControlRequestContext = this.throughputControlRequestContext; context.replicaAddressValidationEnabled = this.replicaAddressValidationEnabled; context.endToEndOperationLatencyPolicyConfig = this.endToEndOperationLatencyPolicyConfig; context.unavailableRegionsForPartition = this.unavailableRegionsForPartition; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/HttpConstants.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/HttpConstants.java index 27694742a257..fea6efeccfc6 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/HttpConstants.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/HttpConstants.java @@ -291,6 +291,9 @@ public static class HttpHeaders { public static final String GLOBAL_DATABASE_ACCOUNT_NAME = "GlobalDatabaseAccountName"; public static final String THINCLIENT_START_EPK = "x-ms-thinclient-range-min"; public static final String THINCLIENT_END_EPK = "x-ms-thinclient-range-max"; + + // Throughput bucket header + public static final String THROUGHPUT_BUCKET = "x-ms-cosmos-throughput-bucket"; } public static class A_IMHeaderValues { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 85051d9a23f3..f7727b95f9d0 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -69,7 +69,8 @@ import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple; import com.azure.cosmos.implementation.spark.OperationListener; import com.azure.cosmos.implementation.throughputControl.ThroughputControlStore; -import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroup; import com.azure.cosmos.models.CosmosAuthorizationTokenResolver; import com.azure.cosmos.models.CosmosBatchResponse; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; @@ -6362,9 +6363,22 @@ public void close() { } } @Override - public synchronized void enableThroughputControlGroup(ThroughputControlGroupInternal group, Mono throughputQueryMono) { + public void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono throughputQueryMono) { checkNotNull(group, "Throughput control group can not be null"); + this.enableThroughputControlStore(); + this.throughputControlStore.enableSDKThroughputControlGroup(group, throughputQueryMono); + } + + @Override + public void enableServerThroughputControlGroup(ServerThroughputControlGroup group) { + checkNotNull(group, "Argument 'group' can not be null"); + + this.enableThroughputControlStore(); + this.throughputControlStore.enableServerThroughputControlGroup(group); + } + + private synchronized void enableThroughputControlStore() { if (this.throughputControlEnabled.compareAndSet(false, true)) { this.throughputControlStore = new ThroughputControlStore( @@ -6378,8 +6392,6 @@ public synchronized void enableThroughputControlGroup(ThroughputControlGroupInte this.gatewayProxy.enableThroughputControl(throughputControlStore); } } - - this.throughputControlStore.enableThroughputControlGroup(group, throughputQueryMono); } @Override 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 b1619d048681..a7369d852625 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 @@ -1173,6 +1173,12 @@ public void setPriorityLevel(PriorityLevel priorityLevel) { } } + public void setThroughputBucket(Integer throughputBucket) { + if (throughputBucket != null) { + this.headers.put(HttpConstants.HttpHeaders.THROUGHPUT_BUCKET, throughputBucket.toString()); + } + } + public Duration getResponseTimeout() { return responseTimeout; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResponseDiagnostics.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResponseDiagnostics.java index eebcda67db69..936913f98cb9 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResponseDiagnostics.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResponseDiagnostics.java @@ -52,6 +52,8 @@ public class StoreResponseDiagnostics { private final String faultInjectionRuleId; private final List faultInjectionEvaluationResults; private final String endpoint; + private final String requestThroughputControlGroupName; + private final String requestThroughputControlGroupConfig; public static StoreResponseDiagnostics createStoreResponseDiagnostics( StoreResponse storeResponse, @@ -91,6 +93,9 @@ private StoreResponseDiagnostics(StoreResponse storeResponse, RxDocumentServiceR this.faultInjectionRuleId = storeResponse.getFaultInjectionRuleId(); this.faultInjectionEvaluationResults = storeResponse.getFaultInjectionRuleEvaluationResults(); this.endpoint = storeResponse.getEndpoint(); + this.requestThroughputControlGroupName = rxDocumentServiceRequest.throughputControlGroupName; + this.requestThroughputControlGroupConfig = + rxDocumentServiceRequest.requestContext.throughputControlRequestContext != null ? rxDocumentServiceRequest.requestContext.throughputControlRequestContext.getConfigString() : null; } private StoreResponseDiagnostics(CosmosException e, RxDocumentServiceRequest rxDocumentServiceRequest) { @@ -122,6 +127,9 @@ private StoreResponseDiagnostics(CosmosException e, RxDocumentServiceRequest rxD } else { this.endpoint = ""; } + this.requestThroughputControlGroupName = rxDocumentServiceRequest.throughputControlGroupName; + this.requestThroughputControlGroupConfig = + rxDocumentServiceRequest.requestContext.throughputControlRequestContext != null ? rxDocumentServiceRequest.requestContext.throughputControlRequestContext.getConfigString() : null; } public int getStatusCode() { @@ -205,4 +213,12 @@ public List getFaultInjectionEvaluationResults() { public String getEndpoint() { return this.endpoint; } + + public String getRequestThroughputControlGroupName() { + return this.requestThroughputControlGroupName; + } + + public String getRequestThroughputControlGroupConfig() { + return this.requestThroughputControlGroupConfig; + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResultDiagnostics.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResultDiagnostics.java index 6770ea63d09a..c6c7bd76ef81 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResultDiagnostics.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResultDiagnostics.java @@ -240,6 +240,14 @@ public void serialize(StoreResultDiagnostics storeResultDiagnostics, jsonGenerator.writeObjectField("channelStatistics", storeResponseDiagnostics.getRntbdChannelStatistics()); jsonGenerator.writeObjectField("serviceEndpointStatistics", storeResponseDiagnostics.getRntbdEndpointStatistics()); + this.writeNonNullStringField( + jsonGenerator, + "requestTCG", + storeResponseDiagnostics.getRequestThroughputControlGroupName()); + this.writeNonNullStringField( + jsonGenerator, + "requestTCGConfig", + storeResponseDiagnostics.getRequestThroughputControlGroupConfig()); jsonGenerator.writeEndObject(); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java index e332263685cf..0129ac84f10a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java @@ -598,7 +598,8 @@ public enum RntbdRequestHeader implements RntbdHeader { SDKSupportedCapabilities((short) 0x00A2, RntbdTokenType.ULong, false), ChangeFeedWireFormatVersion((short) 0x00B2, RntbdTokenType.String, false), PriorityLevel((short) 0x00BF, RntbdTokenType.Byte, false), - GlobalDatabaseAccountName((short) 0x00CE, RntbdTokenType.String, false); + GlobalDatabaseAccountName((short) 0x00CE, RntbdTokenType.String, false), + ThroughputBucket((short)0x00DB, RntbdTokenType.Byte, false); //TODO: does the throughput bucket supported in thin client public static final List thinClientHeadersInOrderList = Arrays.asList( EffectivePartitionKey, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java index 8f258eed4487..0c3f9615312f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java @@ -132,6 +132,7 @@ final class RntbdRequestHeaders extends RntbdTokenStream { this.addChangeFeedWireFormatVersion(headers); this.addPriorityLevel(headers); this.addGlobalDatabaseAccountName(headers); + this.addThroughputBucket(headers); // Normal headers (Strings, Ints, Longs, etc.) @@ -291,6 +292,8 @@ private RntbdToken getCorrelatedActivityId() { private RntbdToken getPriorityLevel() { return this.get(RntbdRequestHeader.PriorityLevel); } + private RntbdToken getThroughputBucket() { return this.get(RntbdRequestHeader.ThroughputBucket); } + private RntbdToken getGlobalDatabaseAccountName() { return this.get(RntbdRequestHeader.GlobalDatabaseAccountName); } @@ -792,6 +795,16 @@ private void addPriorityLevel(final Map headers) } } + private void addThroughputBucket(final Map headers) + { + final String value = headers.get(HttpHeaders.THROUGHPUT_BUCKET); + + if (StringUtils.isNotEmpty(value)) { + final int throughputBucket = Integer.valueOf(value); + this.getThroughputBucket().setValue((byte)throughputBucket); + } + } + private void addGlobalDatabaseAccountName(final Map headers) { final String value = headers.get(HttpHeaders.GLOBAL_DATABASE_ACCOUNT_NAME); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/EmptyThroughputContainerController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/EmptyThroughputContainerController.java similarity index 91% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/EmptyThroughputContainerController.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/EmptyThroughputContainerController.java index 987d9f652bfe..72aeb8e479a0 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/EmptyThroughputContainerController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/EmptyThroughputContainerController.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.container; +package com.azure.cosmos.implementation.throughputControl; import com.azure.cosmos.implementation.RxDocumentServiceRequest; import reactor.core.publisher.Mono; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/IThroughputContainerController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/IThroughputContainerController.java similarity index 57% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/IThroughputContainerController.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/IThroughputContainerController.java index 35a0629f20e9..7933ac3e040c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/IThroughputContainerController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/IThroughputContainerController.java @@ -1,13 +1,10 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.container; - -import com.azure.cosmos.implementation.throughputControl.controller.IThroughputController; +package com.azure.cosmos.implementation.throughputControl; /** * Represents a throughput container controller. */ public interface IThroughputContainerController extends IThroughputController { - } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/IThroughputControlGroup.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/IThroughputControlGroup.java new file mode 100644 index 000000000000..46fbef4cb58b --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/IThroughputControlGroup.java @@ -0,0 +1,8 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl; + +public interface IThroughputControlGroup { + String getDiagnosticsString(); +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/IThroughputController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/IThroughputController.java similarity index 94% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/IThroughputController.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/IThroughputController.java index 203479f8f9e4..c82398dcaa7f 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/IThroughputController.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller; +package com.azure.cosmos.implementation.throughputControl; import com.azure.cosmos.implementation.RxDocumentServiceRequest; import reactor.core.publisher.Mono; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/ThroughputControlGroupFactory.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupFactory.java similarity index 70% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/ThroughputControlGroupFactory.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupFactory.java index c2a95c91593c..0a52d45b831b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/ThroughputControlGroupFactory.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlGroupFactory.java @@ -1,12 +1,15 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.config; +package com.azure.cosmos.implementation.throughputControl; import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.ThroughputControlGroupConfig; import com.azure.cosmos.GlobalThroughputControlConfig; +import com.azure.cosmos.implementation.throughputControl.sdk.config.GlobalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.sdk.config.LocalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroup; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; @@ -48,4 +51,19 @@ public static GlobalThroughputControlGroup createThroughputGlobalControlGroup( globalControlConfig.getControlItemExpireInterval()); } + + public static ServerThroughputControlGroup createServerThroughputControlGroup( + ThroughputControlGroupConfig groupConfig, + CosmosAsyncContainer targetContainer) { + + checkNotNull(groupConfig, "Throughput control group config can not be null"); + checkNotNull(targetContainer, "Throughput target container can not be null"); + + return new ServerThroughputControlGroup( + groupConfig.getGroupName(), + groupConfig.isDefault(), + groupConfig.getPriorityLevel(), + groupConfig.getThroughputBucket(), + targetContainer); + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlRequestContext.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlRequestContext.java new file mode 100644 index 000000000000..0f025e280857 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlRequestContext.java @@ -0,0 +1,33 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl; + +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; + +import java.util.concurrent.atomic.AtomicReference; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; + +public class ThroughputControlRequestContext { + private final String configString; + private final AtomicReference throughputControlCycleId; + + public ThroughputControlRequestContext(String configString) { + checkArgument(StringUtils.isNotEmpty(configString), "Argument 'configString' cannot be null or empty."); + this.configString = configString; + this.throughputControlCycleId = new AtomicReference<>(); + } + + public String getConfigString() { + return this.configString; + } + + public String getThroughputControlCycleId() { + return this.throughputControlCycleId.get(); + } + + public void setThroughputControlCycleId(String throughputControlCycleId) { + this.throughputControlCycleId.set(throughputControlCycleId); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlStore.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlStore.java index e565a7774bc0..feb207b9abb5 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlStore.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlStore.java @@ -5,86 +5,26 @@ import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.ConnectionMode; -import com.azure.cosmos.CosmosException; import com.azure.cosmos.implementation.ResourceType; import com.azure.cosmos.implementation.RxDocumentServiceRequest; import com.azure.cosmos.implementation.Utils; -import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; -import com.azure.cosmos.implementation.apachecommons.lang.tuple.Pair; -import com.azure.cosmos.implementation.caches.AsyncCache; import com.azure.cosmos.implementation.caches.RxClientCollectionCache; import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; -import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal; -import com.azure.cosmos.implementation.throughputControl.controller.IThroughputController; -import com.azure.cosmos.implementation.throughputControl.controller.container.EmptyThroughputContainerController; -import com.azure.cosmos.implementation.throughputControl.controller.container.IThroughputContainerController; -import com.azure.cosmos.implementation.throughputControl.controller.container.ThroughputContainerController; -import com.azure.cosmos.implementation.throughputControl.exceptions.ThroughputControlInitializationException; +import com.azure.cosmos.implementation.throughputControl.sdk.SDKThroughputControlStore; +import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.server.ServerThroughputControlStore; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroup; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import reactor.core.Exceptions; import reactor.core.publisher.Mono; -import java.util.concurrent.ConcurrentHashMap; - -import static com.azure.cosmos.implementation.Exceptions.isNameCacheStale; -import static com.azure.cosmos.implementation.Exceptions.isPartitionKeyMismatchException; -import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; -/** - * This is the entrance class for the whole throughput control work flow pipeline. - * The pipeline will consist of controllers which is implementation of {@link IThroughputController} and {@link ThroughputRequestThrottler}. - * - * Following is a high-level diagram of the pipeline: - * - * +-------------------+ - * | Client | - * +-------------------+ - * | - * | - * | - * +---------------------------------------------------------+ - * | ThroughputControlStore | - * +---------------------------------------------------------+ - * / \ - * / \ - * / \ - * +---------------------------------------+ +---------------------------------------+ - * | Container A controller | | Container B controller | - * +---------------------------------------+ ... +---------------------------------------+ - * / \ - * / \ - * / \ - * +--------------------+ +---------------------+ - * | Group 1 controller | | Group 2 controller | - * +--------------------+ ... +---------------------+ - * | \ - * | \-------------\ - * | \ - * +---------------------------+ +-----------------------------+ - * |Global Request controller | OR | PkRanges Request controller | - * +--------------------------+ +------------------------------+ - * | / \ - * | / \ - * | / \ - * +------------------+ +------------------+ +------------------+ - * |Request throttler | |Request throttler | |Request throttler | - * +------------------+ +------------------+ ... +------------------+ - * - * - */ + public class ThroughputControlStore { private static final Logger logger = LoggerFactory.getLogger(ThroughputControlStore.class); - - private final RxClientCollectionCache collectionCache; - private final ConnectionMode connectionMode; - private final AsyncCache containerControllerCache; - private final ConcurrentHashMap containerMap; - private final RxPartitionKeyRangeCache partitionKeyRangeCache; - - private final LinkedCancellationTokenSource cancellationTokenSource; - private final ConcurrentHashMap cancellationTokenMap; + private final SDKThroughputControlStore sdkThroughputControlStore; + private final ServerThroughputControlStore serverThroughputControlStore; public ThroughputControlStore( RxClientCollectionCache collectionCache, @@ -94,41 +34,37 @@ public ThroughputControlStore( checkNotNull(collectionCache,"RxClientCollectionCache can not be null"); checkNotNull(partitionKeyRangeCache, "PartitionKeyRangeCache can not be null"); - this.collectionCache = collectionCache; - this.connectionMode = connectionMode; - this.containerControllerCache = new AsyncCache<>(); - this.containerMap = new ConcurrentHashMap<>(); - this.partitionKeyRangeCache = partitionKeyRangeCache; - - this.cancellationTokenSource = new LinkedCancellationTokenSource(); - this.cancellationTokenMap = new ConcurrentHashMap<>(); + this.sdkThroughputControlStore = new SDKThroughputControlStore(collectionCache, connectionMode, partitionKeyRangeCache); + this.serverThroughputControlStore = new ServerThroughputControlStore(); } - public void enableThroughputControlGroup(ThroughputControlGroupInternal group, Mono throughputQueryMono) { + public synchronized void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono throughputQueryMono) { checkNotNull(group, "Throughput control group cannot be null"); - String containerNameLink = Utils.trimBeginningAndEndingSlashes(BridgeInternal.extractContainerSelfLink(group.getTargetContainer())); - this.containerMap.compute(containerNameLink, (key, throughputControlContainerProperties) -> { - if (throughputControlContainerProperties == null) { - throughputControlContainerProperties = new ContainerThroughputControlGroupProperties(containerNameLink); + if (group.isDefault()) { + //verify a default group is not being defined in the server throughput control store + String containerNameLink = Utils.trimBeginningAndEndingSlashes(BridgeInternal.extractContainerSelfLink(group.getTargetContainer())); + if (this.serverThroughputControlStore.hasDefaultGroup(containerNameLink)) { + throw new IllegalArgumentException("A default group already exists"); } + } - int groupSizeBefore = throughputControlContainerProperties.getThroughputControlGroupSet().size(); - Pair stateAfterEnabling = - throughputControlContainerProperties.enableThroughputControlGroup(group, throughputQueryMono); + this.sdkThroughputControlStore.enableThroughputControlGroup(group, throughputQueryMono); + } - int groupSizeAfter = stateAfterEnabling.getLeft(); - boolean wasGroupConfigUpdated = stateAfterEnabling.getRight(); + public synchronized void enableServerThroughputControlGroup(ServerThroughputControlGroup group) { + checkNotNull(group, "Throughput control group cannot be null"); - if ((groupSizeAfter > groupSizeBefore && groupSizeAfter == 1) || wasGroupConfigUpdated) { - // This is the first enabled group for the target container or an existing group was modified - // Clean the current cache in case we have built EmptyThroughputContainerController or an existing - // group was modified - this.containerControllerCache.remove(containerNameLink); + if (group.isDefault()) { + //verify a default group is not being defined in the sdk throughput control store + String containerNameLink = + Utils.trimBeginningAndEndingSlashes(BridgeInternal.extractContainerSelfLink(group.getTargetContainer())); + if (this.sdkThroughputControlStore.hasDefaultGroup(containerNameLink)) { + throw new IllegalArgumentException("A default group already exists"); } + } - return throughputControlContainerProperties; - }); + this.serverThroughputControlStore.enableThroughputControlGroup(group); } public Mono processRequest(RxDocumentServiceRequest request, Mono originalRequestMono) { @@ -142,162 +78,28 @@ public Mono processRequest(RxDocumentServiceRequest request, Mono orig } String collectionNameLink = Utils.getCollectionName(request.getResourceAddress()); - return this.resolveContainerController(collectionNameLink) - .flatMap(containerController -> { - if (containerController.canHandleRequest(request)) { - return containerController.processRequest(request, originalRequestMono) - .doOnError(throwable -> this.handleException(collectionNameLink, request, throwable)); - } - - // Unable to find container controller to handle the request, - // It is caused by control store out of sync or the request has staled info. - // We will handle the first scenario by creating a new container controller, - // while fall back to original request Mono for the second scenario. - return this.updateControllerAndRetry(collectionNameLink, request, originalRequestMono); - }) - .onErrorResume(throwable -> { - - Exception unwrappedException = Utils.as(Exceptions.unwrap(throwable), Exception.class); - if (unwrappedException instanceof ThroughputControlInitializationException) { - if (this.shouldContinueRequestOnInitError(request, collectionNameLink, unwrappedException)) { - return originalRequestMono; - } - - return Mono.error(unwrappedException.getCause()); - } - - return Mono.error(throwable); - }); - } - - private boolean shouldContinueRequestOnInitError(RxDocumentServiceRequest request, String collectionNameLink, Throwable throwable) { - if (throwable instanceof ThroughputControlInitializationException) { - ContainerThroughputControlGroupProperties throughputControlContainerProperties = this.containerMap.get(collectionNameLink); - - checkNotNull( - throughputControlContainerProperties, - "Throughput control container properties should not be null"); - checkArgument( - throughputControlContainerProperties.getThroughputControlGroupSet().size() > 0, - "There should be more than one throughput control group"); - - return throughputControlContainerProperties.allowRequestToContinueOnInitError(request); + if (this.serverThroughputControlStore.hasGroup(collectionNameLink, request.getThroughputControlGroupName())) { + return this.serverThroughputControlStore.processRequest(request, originalRequestMono); } - return false; - } - - private Mono updateControllerAndRetry( - String containerNameLink, - RxDocumentServiceRequest request, - Mono originalRequestMono) { - - return this.shouldRefreshContainerController(containerNameLink, request) - .flatMap(shouldRefresh -> { - if (shouldRefresh) { - this.cancellationTokenMap.compute(containerNameLink, (key, cancellationToken) -> { - if (cancellationToken != null) { - cancellationToken.cancel(); - } - - return null; - }); - - this.containerControllerCache.refresh(containerNameLink, () -> this.createAndInitContainerController(containerNameLink)); - return this.resolveContainerController(containerNameLink) - .flatMap(updatedContainerController -> { - if (updatedContainerController.canHandleRequest(request)) { - return updatedContainerController.processRequest(request, originalRequestMono) - .doOnError(throwable -> this.handleException(containerNameLink, request, throwable)); - } else { - // still can not handle the request - logger.warn( - "Can not find container controller to process request {} with collectionRid {} ", - request.getActivityId(), - request.requestContext.resolvedCollectionRid); - - return originalRequestMono; - } - }); - } - - return originalRequestMono; - }); - } - - private Mono resolveContainerController(String containerNameLink) { - checkArgument(StringUtils.isNotEmpty(containerNameLink), "Container name link can not be null or empty"); - - return this.containerControllerCache.getAsync( - containerNameLink, - null, - () -> this.createAndInitContainerController(containerNameLink)) - .onErrorResume(throwable -> Mono.error(new ThroughputControlInitializationException(throwable))); - } - - private Mono createAndInitContainerController(String containerNameLink) { - checkArgument(StringUtils.isNotEmpty(containerNameLink), "Container link should not be null or empty"); - - if (this.containerMap.containsKey(containerNameLink)) { - return Mono.just(this.containerMap.get(containerNameLink)) - .flatMap(throughputControlContainerProperties -> { - LinkedCancellationToken parentToken = - this.cancellationTokenMap.compute( - containerNameLink, - (key, cancellationToken) -> this.cancellationTokenSource.getToken()); - - ThroughputContainerController containerController = - new ThroughputContainerController( - this.collectionCache, - this.connectionMode, - throughputControlContainerProperties.getThroughputControlGroupSet(), - this.partitionKeyRangeCache, - parentToken, - throughputControlContainerProperties.getThroughputQueryMono()); - - return containerController.init(); - }); - } else { - return Mono.just(new EmptyThroughputContainerController()) - .flatMap(EmptyThroughputContainerController::init); + if (this.sdkThroughputControlStore.hasGroup(collectionNameLink, request.getThroughputControlGroupName())) { + return this.sdkThroughputControlStore.processRequest(request, originalRequestMono); } - } - - private Mono shouldRefreshContainerController(String containerLink, RxDocumentServiceRequest request) { - // TODO: populate diagnostics - return this.collectionCache.resolveByNameAsync(null, containerLink, null) - .flatMap(documentCollection -> - Mono.just(StringUtils.equals(documentCollection.getResourceId(), request.requestContext.resolvedCollectionRid))); - } - - private void handleException(String containerNameLink, RxDocumentServiceRequest request, Throwable throwable) { - checkArgument(StringUtils.isNotEmpty(containerNameLink), "Container name link can not be null nor empty"); - checkNotNull(request, "Request can not be null"); - checkNotNull(throwable, "Exception can not be null"); - CosmosException cosmosException = Utils.as(Exceptions.unwrap(throwable), CosmosException.class); - - if (cosmosException != null && - (isNameCacheStale(cosmosException) || isPartitionKeyMismatchException(cosmosException))) { - - this.cancellationTokenMap.compute(containerNameLink,(key, cancellationToken) -> { - if (cancellationToken != null) { - cancellationToken.cancel(); - } - return null; - }); - - String containerLink = Utils.getCollectionName(request.getResourceAddress()); + // can not find exact throughput control group mapping, using default group if any + if (this.serverThroughputControlStore.hasDefaultGroup(collectionNameLink)) { + return this.serverThroughputControlStore.processRequest(request, originalRequestMono); + } - this.collectionCache.refresh(null, containerLink, null); - this.containerControllerCache.refresh( - containerLink, - () -> createAndInitContainerController(containerLink) - ); + if (this.sdkThroughputControlStore.hasDefaultGroup(collectionNameLink)) { + return this.sdkThroughputControlStore.processRequest(request, originalRequestMono); } + + // neither store can process the request, fallback to just use the original mono + return originalRequestMono; } public void close() { - this.cancellationTokenSource.close(); + this.sdkThroughputControlStore.close(); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ContainerThroughputControlGroupProperties.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerSDKThroughputControlGroupProperties.java similarity index 73% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ContainerThroughputControlGroupProperties.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerSDKThroughputControlGroupProperties.java index d67f0e400266..7975c2a8b9de 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ContainerThroughputControlGroupProperties.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ContainerSDKThroughputControlGroupProperties.java @@ -1,16 +1,17 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl; +package com.azure.cosmos.implementation.throughputControl.sdk; import com.azure.cosmos.implementation.RxDocumentServiceRequest; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.implementation.apachecommons.lang.tuple.Pair; -import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Mono; +import java.util.Map; import java.util.Objects; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; @@ -20,21 +21,21 @@ import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; -public class ContainerThroughputControlGroupProperties { - private static Logger logger = LoggerFactory.getLogger(ContainerThroughputControlGroupProperties.class); +public class ContainerSDKThroughputControlGroupProperties { + private static Logger logger = LoggerFactory.getLogger(ContainerSDKThroughputControlGroupProperties.class); - private final AtomicReference defaultGroup; - private final Set throughputControlGroupSet; + private final AtomicReference defaultGroup; + private final Map throughputControlGroups; private final Set supressInitErrorGroupSet; private final AtomicReference> throughputQueryMonoReference; private final String containerNameLink; - public ContainerThroughputControlGroupProperties(String containerNameLink) { + public ContainerSDKThroughputControlGroupProperties(String containerNameLink) { checkArgument(StringUtils.isNotEmpty(containerNameLink), "Argument 'containerNameLink' should not be empty"); - this.containerNameLink = containerNameLink; this.defaultGroup = new AtomicReference<>(); - this.throughputControlGroupSet = ConcurrentHashMap.newKeySet(); + this.containerNameLink = containerNameLink; + this.throughputControlGroups = new ConcurrentHashMap<>(); this.supressInitErrorGroupSet = ConcurrentHashMap.newKeySet(); this.throughputQueryMonoReference = new AtomicReference<>(); } @@ -42,11 +43,11 @@ public ContainerThroughputControlGroupProperties(String containerNameLink) { /*** * Enable a throughput control group. * - * @param group a {@link ThroughputControlGroupInternal}. + * @param group a {@link SDKThroughputControlGroupInternal}. * * @return the total size of distinct throughput control groups enabled on the container. */ - public Pair enableThroughputControlGroup(ThroughputControlGroupInternal group, Mono throughputQueryMono) { + public Pair enableThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono throughputQueryMono) { checkNotNull(group, "Throughput control group should not be null"); if (group.isDefault()) { @@ -65,8 +66,8 @@ public Pair enableThroughputControlGroup(ThroughputControlGrou // Only throw when two different groups are using the same id (databaseId + containerId + groupName) this - .throughputControlGroupSet - .stream() + .throughputControlGroups + .values() .forEach(existingGroup -> { if (Objects.equals(existingGroup.getIdPrefix(), group.getIdPrefix()) && !existingGroup.equals(group)) { @@ -79,7 +80,7 @@ public Pair enableThroughputControlGroup(ThroughputControlGrou "cannot be changed."); } - this.throughputControlGroupSet.remove(existingGroup); + this.throughputControlGroups.remove(existingGroup); logger.info( "Throughput control group with id-prefix {} already exists with different config. " + "Will update config.", @@ -88,16 +89,16 @@ public Pair enableThroughputControlGroup(ThroughputControlGrou } }); - this.throughputControlGroupSet.add(group); + this.throughputControlGroups.put(group.getGroupName(), group); if (!this.throughputQueryMonoReference.compareAndSet(null, throughputQueryMono)) { logger.debug("ThroughputQueryMono has exists for container {}", containerNameLink); } - return Pair.of(this.throughputControlGroupSet.size(), updatedGroupConfig.get()); + return Pair.of(this.throughputControlGroups.size(), updatedGroupConfig.get()); } - public Set getThroughputControlGroupSet() { - return this.throughputControlGroupSet; + public Map getThroughputControlGroups() { + return this.throughputControlGroups; } public Mono getThroughputQueryMono() { @@ -117,4 +118,12 @@ public boolean allowRequestToContinueOnInitError(RxDocumentServiceRequest reques return this.supressInitErrorGroupSet.contains(requestGroupName); } + + public boolean hasDefaultGroup() { + return this.defaultGroup.get() != null; + } + + public boolean hasGroup(String groupName) { + return this.throughputControlGroups.containsKey(groupName); + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationToken.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/LinkedCancellationToken.java similarity index 96% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationToken.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/LinkedCancellationToken.java index c53ff1316fd0..90b09d70c4dc 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationToken.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/LinkedCancellationToken.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl; +package com.azure.cosmos.implementation.throughputControl.sdk; import java.util.ArrayList; import java.util.List; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationTokenSource.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/LinkedCancellationTokenSource.java similarity index 96% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationTokenSource.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/LinkedCancellationTokenSource.java index 0bbfd014fb7d..cf65b9ba2b5e 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/LinkedCancellationTokenSource.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/LinkedCancellationTokenSource.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl; +package com.azure.cosmos.implementation.throughputControl.sdk; import java.io.Closeable; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/SDKThroughputControlStore.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/SDKThroughputControlStore.java new file mode 100644 index 000000000000..d48cf3262d60 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/SDKThroughputControlStore.java @@ -0,0 +1,323 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl.sdk; + +import com.azure.cosmos.BridgeInternal; +import com.azure.cosmos.ConnectionMode; +import com.azure.cosmos.CosmosException; +import com.azure.cosmos.implementation.ResourceType; +import com.azure.cosmos.implementation.RxDocumentServiceRequest; +import com.azure.cosmos.implementation.Utils; +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.implementation.apachecommons.lang.tuple.Pair; +import com.azure.cosmos.implementation.caches.AsyncCache; +import com.azure.cosmos.implementation.caches.RxClientCollectionCache; +import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; +import com.azure.cosmos.implementation.throughputControl.EmptyThroughputContainerController; +import com.azure.cosmos.implementation.throughputControl.IThroughputContainerController; +import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.IThroughputController; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.container.SDKThroughputContainerController; +import com.azure.cosmos.implementation.throughputControl.sdk.exceptions.ThroughputControlInitializationException; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import reactor.core.Exceptions; +import reactor.core.publisher.Mono; + +import java.util.concurrent.ConcurrentHashMap; + +import static com.azure.cosmos.implementation.Exceptions.isNameCacheStale; +import static com.azure.cosmos.implementation.Exceptions.isPartitionKeyMismatchException; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +/** + * This is the entrance class for the whole throughput control work flow pipeline. + * The pipeline will consist of controllers which is implementation of {@link IThroughputController} and {@link ThroughputRequestThrottler}. + * + * Following is a high-level diagram of the pipeline: + * + * +-------------------+ + * | Client | + * +-------------------+ + * | + * | + * | + * +---------------------------------------------------------+ + * | ThroughputControlStore | + * +---------------------------------------------------------+ + * / \ + * / \ + * / \ + * +---------------------------------------+ +---------------------------------------+ + * | Container A controller | | Container B controller | + * +---------------------------------------+ ... +---------------------------------------+ + * / \ + * / \ + * / \ + * +--------------------+ +---------------------+ + * | Group 1 controller | | Group 2 controller | + * +--------------------+ ... +---------------------+ + * | \ + * | \-------------\ + * | \ + * +---------------------------+ +-----------------------------+ + * |Global Request controller | OR | PkRanges Request controller | + * +--------------------------+ +------------------------------+ + * | / \ + * | / \ + * | / \ + * +------------------+ +------------------+ +------------------+ + * |Request throttler | |Request throttler | |Request throttler | + * +------------------+ +------------------+ ... +------------------+ + * + * + */ +public class SDKThroughputControlStore { + private static final Logger logger = LoggerFactory.getLogger(SDKThroughputControlStore.class); + + private final RxClientCollectionCache collectionCache; + private final ConnectionMode connectionMode; + private final AsyncCache containerControllerCache; + private final ConcurrentHashMap containerMap; + private final RxPartitionKeyRangeCache partitionKeyRangeCache; + + private final LinkedCancellationTokenSource cancellationTokenSource; + private final ConcurrentHashMap cancellationTokenMap; + + public SDKThroughputControlStore( + RxClientCollectionCache collectionCache, + ConnectionMode connectionMode, + RxPartitionKeyRangeCache partitionKeyRangeCache) { + + checkNotNull(collectionCache,"RxClientCollectionCache can not be null"); + checkNotNull(partitionKeyRangeCache, "PartitionKeyRangeCache can not be null"); + + this.collectionCache = collectionCache; + this.connectionMode = connectionMode; + this.containerControllerCache = new AsyncCache<>(); + this.containerMap = new ConcurrentHashMap<>(); + this.partitionKeyRangeCache = partitionKeyRangeCache; + + this.cancellationTokenSource = new LinkedCancellationTokenSource(); + this.cancellationTokenMap = new ConcurrentHashMap<>(); + } + + public void enableThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono throughputQueryMono) { + checkNotNull(group, "Throughput control group cannot be null"); + + String containerNameLink = Utils.trimBeginningAndEndingSlashes(BridgeInternal.extractContainerSelfLink(group.getTargetContainer())); + this.containerMap.compute(containerNameLink, (key, throughputControlContainerProperties) -> { + if (throughputControlContainerProperties == null) { + throughputControlContainerProperties = new ContainerSDKThroughputControlGroupProperties(containerNameLink); + } + + int groupSizeBefore = throughputControlContainerProperties.getThroughputControlGroups().size(); + Pair stateAfterEnabling = + throughputControlContainerProperties.enableThroughputControlGroup(group, throughputQueryMono); + + int groupSizeAfter = stateAfterEnabling.getLeft(); + boolean wasGroupConfigUpdated = stateAfterEnabling.getRight(); + + if ((groupSizeAfter > groupSizeBefore && groupSizeAfter == 1) || wasGroupConfigUpdated) { + // This is the first enabled group for the target container or an existing group was modified + // Clean the current cache in case we have built EmptyThroughputContainerController or an existing + // group was modified + this.containerControllerCache.remove(containerNameLink); + } + + return throughputControlContainerProperties; + }); + } + + public Mono processRequest(RxDocumentServiceRequest request, Mono originalRequestMono) { + checkNotNull(request, "Request can not be null"); + checkNotNull(originalRequestMono, "originalRequestMono can not be null"); + + // Currently, we will only target two resource types. + // If in the future we find other useful scenarios for throughput control, add more resource type here. + if (request.getResourceType() != ResourceType.Document && request.getResourceType() != ResourceType.StoredProcedure) { + return originalRequestMono; + } + + String collectionNameLink = Utils.getCollectionName(request.getResourceAddress()); + return this.resolveContainerController(collectionNameLink) + .flatMap(containerController -> { + if (containerController.canHandleRequest(request)) { + return containerController.processRequest(request, originalRequestMono) + .doOnError(throwable -> this.handleException(collectionNameLink, request, throwable)); + } + + // Unable to find container controller to handle the request, + // It is caused by control store out of sync or the request has staled info. + // We will handle the first scenario by creating a new container controller, + // while fall back to original request Mono for the second scenario. + return this.updateControllerAndRetry(collectionNameLink, request, originalRequestMono); + }) + .onErrorResume(throwable -> { + + Exception unwrappedException = Utils.as(Exceptions.unwrap(throwable), Exception.class); + if (unwrappedException instanceof ThroughputControlInitializationException) { + if (this.shouldContinueRequestOnInitError(request, collectionNameLink, unwrappedException)) { + return originalRequestMono; + } + + return Mono.error(unwrappedException.getCause()); + } + + return Mono.error(throwable); + }); + } + + private boolean shouldContinueRequestOnInitError(RxDocumentServiceRequest request, String collectionNameLink, Throwable throwable) { + if (throwable instanceof ThroughputControlInitializationException) { + ContainerSDKThroughputControlGroupProperties throughputControlContainerProperties = this.containerMap.get(collectionNameLink); + + checkNotNull( + throughputControlContainerProperties, + "Throughput control container properties should not be null"); + checkArgument( + throughputControlContainerProperties.getThroughputControlGroups().size() > 0, + "There should be more than one throughput control group"); + + return throughputControlContainerProperties.allowRequestToContinueOnInitError(request); + } + + return false; + } + + private Mono updateControllerAndRetry( + String containerNameLink, + RxDocumentServiceRequest request, + Mono originalRequestMono) { + + return this.shouldRefreshContainerController(containerNameLink, request) + .flatMap(shouldRefresh -> { + if (shouldRefresh) { + this.cancellationTokenMap.compute(containerNameLink, (key, cancellationToken) -> { + if (cancellationToken != null) { + cancellationToken.cancel(); + } + + return null; + }); + + this.containerControllerCache.refresh(containerNameLink, () -> this.createAndInitContainerController(containerNameLink)); + return this.resolveContainerController(containerNameLink) + .flatMap(updatedContainerController -> { + if (updatedContainerController.canHandleRequest(request)) { + return updatedContainerController.processRequest(request, originalRequestMono) + .doOnError(throwable -> this.handleException(containerNameLink, request, throwable)); + } else { + // still can not handle the request + logger.warn( + "Can not find container controller to process request {} with collectionRid {} ", + request.getActivityId(), + request.requestContext.resolvedCollectionRid); + + return originalRequestMono; + } + }); + } + + return originalRequestMono; + }); + } + + private Mono resolveContainerController(String containerNameLink) { + checkArgument(StringUtils.isNotEmpty(containerNameLink), "Container name link can not be null or empty"); + + return this.containerControllerCache.getAsync( + containerNameLink, + null, + () -> this.createAndInitContainerController(containerNameLink)) + .onErrorResume(throwable -> Mono.error(new ThroughputControlInitializationException(throwable))); + } + + private Mono createAndInitContainerController(String containerNameLink) { + checkArgument(StringUtils.isNotEmpty(containerNameLink), "Container link should not be null or empty"); + + if (this.containerMap.containsKey(containerNameLink)) { + return Mono.just(this.containerMap.get(containerNameLink)) + .flatMap(throughputControlContainerProperties -> { + LinkedCancellationToken parentToken = + this.cancellationTokenMap.compute( + containerNameLink, + (key, cancellationToken) -> this.cancellationTokenSource.getToken()); + + SDKThroughputContainerController containerController = + new SDKThroughputContainerController( + this.collectionCache, + this.connectionMode, + throughputControlContainerProperties.getThroughputControlGroups(), + this.partitionKeyRangeCache, + parentToken, + throughputControlContainerProperties.getThroughputQueryMono()); + + return containerController.init(); + }); + } else { + return Mono.just(new EmptyThroughputContainerController()) + .flatMap(EmptyThroughputContainerController::init); + } + } + + private Mono shouldRefreshContainerController(String containerLink, RxDocumentServiceRequest request) { + // TODO: populate diagnostics + return this.collectionCache.resolveByNameAsync(null, containerLink, null) + .flatMap(documentCollection -> + Mono.just(StringUtils.equals(documentCollection.getResourceId(), request.requestContext.resolvedCollectionRid))); + } + + private void handleException(String containerNameLink, RxDocumentServiceRequest request, Throwable throwable) { + checkArgument(StringUtils.isNotEmpty(containerNameLink), "Container name link can not be null nor empty"); + checkNotNull(request, "Request can not be null"); + checkNotNull(throwable, "Exception can not be null"); + + CosmosException cosmosException = Utils.as(Exceptions.unwrap(throwable), CosmosException.class); + + if (cosmosException != null && + (isNameCacheStale(cosmosException) || isPartitionKeyMismatchException(cosmosException))) { + + this.cancellationTokenMap.compute(containerNameLink,(key, cancellationToken) -> { + if (cancellationToken != null) { + cancellationToken.cancel(); + } + return null; + }); + + String containerLink = Utils.getCollectionName(request.getResourceAddress()); + + this.collectionCache.refresh(null, containerLink, null); + this.containerControllerCache.refresh( + containerLink, + () -> createAndInitContainerController(containerLink) + ); + } + } + + public boolean hasDefaultGroup(String containerNameLink) { + if (this.containerMap.containsKey(containerNameLink)) { + return this.containerMap.get(containerNameLink).hasDefaultGroup(); + } + + return false; + } + + public boolean hasGroup(String containerNameLink, String throughputControlGroupName) { + if (StringUtils.isEmpty(throughputControlGroupName)) { + return false; + } + + if (this.containerMap.containsKey(containerNameLink)) { + return this.containerMap.get(containerNameLink).hasGroup(throughputControlGroupName); + } + + return false; + } + + public void close() { + this.cancellationTokenSource.close(); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTrackingUnit.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputControlTrackingUnit.java similarity index 98% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTrackingUnit.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputControlTrackingUnit.java index 020f56c8f172..1d24df68a90f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputControlTrackingUnit.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputControlTrackingUnit.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl; +package com.azure.cosmos.implementation.throughputControl.sdk; import com.azure.cosmos.implementation.OperationType; import org.slf4j.Logger; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputRequestThrottler.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputRequestThrottler.java similarity index 96% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputRequestThrottler.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputRequestThrottler.java index e54632cea1e7..1b75c2ffd138 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/ThroughputRequestThrottler.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/ThroughputRequestThrottler.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl; +package com.azure.cosmos.implementation.throughputControl.sdk; import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.CosmosException; @@ -20,7 +20,6 @@ import reactor.core.Exceptions; import reactor.core.publisher.Mono; -import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicReference; import java.util.concurrent.locks.ReentrantReadWriteLock; @@ -93,8 +92,8 @@ public Mono processRequest(RxDocumentServiceRequest request, Mono orig })); if (this.availableThroughput.get() > 0) { - if (StringUtils.isEmpty(request.requestContext.throughputControlCycleId)) { - request.requestContext.throughputControlCycleId = this.cycleId; + if (StringUtils.isEmpty(request.requestContext.throughputControlRequestContext.getThroughputControlCycleId())) { + request.requestContext.throughputControlRequestContext.setThroughputControlCycleId(this.cycleId); } trackingUnit.increasePassedRequest(); @@ -179,7 +178,7 @@ private void trackRequestCharge (RxDocumentServiceRequest request, T respons } // If the response comes back in a different cycle, discard it. - if (StringUtils.equals(this.cycleId, request.requestContext.throughputControlCycleId)) { + if (StringUtils.equals(this.cycleId, request.requestContext.throughputControlRequestContext.getThroughputControlCycleId())) { this.availableThroughput.getAndAccumulate(requestCharge, (available, consumed) -> available - consumed); } else { if (trackingUnit != null) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/GlobalThroughputControlGroup.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/GlobalThroughputControlGroup.java similarity index 89% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/GlobalThroughputControlGroup.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/GlobalThroughputControlGroup.java index 5ee5b55fd3e5..bd0574e7d354 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/GlobalThroughputControlGroup.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/GlobalThroughputControlGroup.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.config; +package com.azure.cosmos.implementation.throughputControl.sdk.config; import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.models.PriorityLevel; @@ -10,7 +10,7 @@ import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; -public class GlobalThroughputControlGroup extends ThroughputControlGroupInternal { +public class GlobalThroughputControlGroup extends SDKThroughputControlGroupInternal { private static final Duration DEFAULT_CONTROL_ITEM_RENEW_INTERVAL = Duration.ofSeconds(5); private final CosmosAsyncContainer globalControlContainer; @@ -29,7 +29,14 @@ public GlobalThroughputControlGroup( Duration controlItemRenewInterval, Duration controlItemExpireInterval) { - super (groupName, targetContainer, targetThroughput, targetThroughputThreshold, priorityLevel, isDefault, continueOnInitError); + super ( + groupName, + targetContainer, + targetThroughput, + targetThroughputThreshold, + priorityLevel, + isDefault, + continueOnInitError); checkNotNull(globalControlContainer, "Global control container can not be null"); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/LocalThroughputControlGroup.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/LocalThroughputControlGroup.java similarity index 56% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/LocalThroughputControlGroup.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/LocalThroughputControlGroup.java index 15476deb9fb4..6d1a2caa672f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/LocalThroughputControlGroup.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/LocalThroughputControlGroup.java @@ -1,12 +1,12 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.config; +package com.azure.cosmos.implementation.throughputControl.sdk.config; import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.models.PriorityLevel; -public class LocalThroughputControlGroup extends ThroughputControlGroupInternal { +public class LocalThroughputControlGroup extends SDKThroughputControlGroupInternal { public LocalThroughputControlGroup( String groupName, @@ -16,6 +16,13 @@ public LocalThroughputControlGroup( PriorityLevel priorityLevel, boolean isDefault, boolean continueOnInitError) { - super (groupName, targetContainer, targetThroughput, targetThroughputThreshold, priorityLevel, isDefault, continueOnInitError); + super ( + groupName, + targetContainer, + targetThroughput, + targetThroughputThreshold, + priorityLevel, + isDefault, + continueOnInitError); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/ThroughputControlGroupInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/SDKThroughputControlGroupInternal.java similarity index 84% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/ThroughputControlGroupInternal.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/SDKThroughputControlGroupInternal.java index 3ffa06c2933a..981fd440dd34 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/config/ThroughputControlGroupInternal.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/config/SDKThroughputControlGroupInternal.java @@ -1,10 +1,11 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.config; +package com.azure.cosmos.implementation.throughputControl.sdk.config; import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.implementation.throughputControl.IThroughputControlGroup; import com.azure.cosmos.models.PriorityLevel; import java.util.Objects; @@ -12,7 +13,7 @@ import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; -public abstract class ThroughputControlGroupInternal { +public abstract class SDKThroughputControlGroupInternal implements IThroughputControlGroup { private final String groupName; private final String idPrefix; private final String id; @@ -21,9 +22,9 @@ public abstract class ThroughputControlGroupInternal { private final CosmosAsyncContainer targetContainer; private final Integer targetThroughput; private final Double targetThroughputThreshold; - private PriorityLevel priorityLevel; + private final PriorityLevel priorityLevel; - public ThroughputControlGroupInternal( + public SDKThroughputControlGroupInternal( String groupName, CosmosAsyncContainer targetContainer, Integer targetThroughput, @@ -59,7 +60,11 @@ public ThroughputControlGroupInternal( getThroughputIdSuffix(targetThroughput, targetThroughputThreshold, priorityLevel)); } - public static String getThroughputIdSuffix(Integer throughput, Double throughputThreshold, PriorityLevel priorityLevel) { + public static String getThroughputIdSuffix( + Integer throughput, + Double throughputThreshold, + PriorityLevel priorityLevel) { + StringBuilder sb = new StringBuilder(); if (throughput != null) { sb.append("t-"); @@ -186,7 +191,7 @@ public boolean equals(Object other) { return false; } - ThroughputControlGroupInternal that = (ThroughputControlGroupInternal) other; + SDKThroughputControlGroupInternal that = (SDKThroughputControlGroupInternal) other; return Objects.equals(this.id, that.id) && this.isDefault == that.isDefault @@ -205,13 +210,33 @@ public boolean hasSameIdentity(Object other) { return false; } - ThroughputControlGroupInternal that = (ThroughputControlGroupInternal) other; + SDKThroughputControlGroupInternal that = (SDKThroughputControlGroupInternal) other; return Objects.equals(this.idPrefix, that.idPrefix) && this.isDefault == that.isDefault && this.continueOnInitError == that.continueOnInitError; } + @Override + public String getDiagnosticsString() { + StringBuilder sb = new StringBuilder(); + sb.append("("); + sb.append("name=" + this.groupName); + sb.append(", default=" + this.isDefault); + if (this.priorityLevel != null) { + sb.append(", priorityLevel=" + this.priorityLevel); + } + if (this.targetThroughputThreshold != null) { + sb.append(", targetRUThreshold=" + this.targetThroughputThreshold); + } + if (this.targetThroughput != null) { + sb.append(", targetRU=" + this.targetThroughput); + } + + sb.append(")"); + return sb.toString(); + } + @Override public int hashCode() { int hash = 0; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/ThroughputContainerController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/container/SDKThroughputContainerController.java similarity index 85% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/ThroughputContainerController.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/container/SDKThroughputContainerController.java index 9f5e938f22db..eb688dd297ec 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/ThroughputContainerController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/container/SDKThroughputContainerController.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.container; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.container; import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.ConnectionMode; @@ -22,12 +22,13 @@ import com.azure.cosmos.implementation.caches.AsyncCache; import com.azure.cosmos.implementation.caches.RxCollectionCache; import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; -import com.azure.cosmos.implementation.throughputControl.LinkedCancellationToken; -import com.azure.cosmos.implementation.throughputControl.LinkedCancellationTokenSource; -import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal; -import com.azure.cosmos.implementation.throughputControl.controller.group.ThroughputGroupControllerBase; -import com.azure.cosmos.implementation.throughputControl.controller.group.ThroughputGroupControllerFactory; -import com.azure.cosmos.implementation.throughputControl.exceptions.ThroughputControlInitializationException; +import com.azure.cosmos.implementation.throughputControl.IThroughputContainerController; +import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationToken; +import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationTokenSource; +import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.group.SDKThroughputGroupControllerBase; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.group.ThroughputGroupControllerFactory; +import com.azure.cosmos.implementation.throughputControl.sdk.exceptions.ThroughputControlInitializationException; import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.ModelBridgeInternal; import com.azure.cosmos.models.ThroughputProperties; @@ -40,7 +41,7 @@ import reactor.util.retry.RetrySpec; import java.time.Duration; -import java.util.Set; +import java.util.Map; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicReference; @@ -53,8 +54,8 @@ * 2. Create and initialize throughput group controllers * 3. Start throughput refresh task if necessary */ -public class ThroughputContainerController implements IThroughputContainerController { - private static final Logger logger = LoggerFactory.getLogger(ThroughputContainerController.class); +public class SDKThroughputContainerController implements IThroughputContainerController { + private static final Logger logger = LoggerFactory.getLogger(SDKThroughputContainerController.class); private static final Duration DEFAULT_THROUGHPUT_REFRESH_INTERVAL = Duration.ofMinutes(15); private static final int NO_OFFER_EXCEPTION_STATUS_CODE = HttpConstants.StatusCodes.BADREQUEST; @@ -63,8 +64,8 @@ public class ThroughputContainerController implements IThroughputContainerContro private final AsyncDocumentClient client; private final RxCollectionCache collectionCache; private final ConnectionMode connectionMode; - private final AsyncCache groupControllerCache; - private final Set groups; + private final AsyncCache groupControllerCache; + private final Map groups; private final AtomicReference maxContainerThroughput; private final RxPartitionKeyRangeCache partitionKeyRangeCache; private final CosmosAsyncContainer targetContainer; @@ -73,15 +74,15 @@ public class ThroughputContainerController implements IThroughputContainerContro private final ConcurrentHashMap cancellationTokenMap; private final Mono throughputQueryMono; - private ThroughputGroupControllerBase defaultGroupController; + private SDKThroughputGroupControllerBase defaultGroupController; private String targetContainerRid; private String targetDatabaseRid; private ThroughputProvisioningScope throughputProvisioningScope; - public ThroughputContainerController( + public SDKThroughputContainerController( RxCollectionCache collectionCache, ConnectionMode connectionMode, - Set groups, + Map groups, RxPartitionKeyRangeCache partitionKeyRangeCache, LinkedCancellationToken parentToken, Mono throughputQueryMono) { @@ -98,7 +99,7 @@ public ThroughputContainerController( this.maxContainerThroughput = new AtomicReference<>(); this.partitionKeyRangeCache = partitionKeyRangeCache; - this.targetContainer = groups.iterator().next().getTargetContainer(); + this.targetContainer = groups.values().iterator().next().getTargetContainer(); this.client = CosmosBridgeInternal.getContextClient(this.targetContainer); this.throughputProvisioningScope = this.getThroughputResolveLevel(groups); @@ -108,8 +109,8 @@ public ThroughputContainerController( this.throughputQueryMono = throughputQueryMono == null ? this.resolveContainerMaxThroughputCore() : throughputQueryMono; } - private ThroughputProvisioningScope getThroughputResolveLevel(Set groupConfigs) { - if (groupConfigs.stream().anyMatch(groupConfig -> groupConfig.getTargetThroughputThreshold() != null)) { + private ThroughputProvisioningScope getThroughputResolveLevel(Map groupConfigs) { + if (groupConfigs.values().stream().anyMatch(groupConfig -> groupConfig.getTargetThroughputThreshold() != null)) { // Throughput can be provisioned on container level or database level, will start from container return ThroughputProvisioningScope.CONTAINER; } else { @@ -177,7 +178,7 @@ private Mono resolveContainerThroughput() { } } - private Mono resolveContainerMaxThroughput() { + private Mono resolveContainerMaxThroughput() { return this.throughputQueryMono .flatMap(maxThroughput -> { this.maxContainerThroughput.set(maxThroughput); @@ -304,22 +305,20 @@ 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.just(new Utils.ValueHolder<>(this.defaultGroupController)); } - for (ThroughputControlGroupInternal group : this.groups) { - if (StringUtils.equals(groupName, group.getGroupName())) { - return this.resolveThroughputGroupController(group) - .map(Utils.ValueHolder::new); - } + SDKThroughputControlGroupInternal group = this.groups.get(groupName); + if (group == null) { + // If the request is associated with a group not enabled, will fall back to the default one. + return Mono.just(new Utils.ValueHolder<>(this.defaultGroupController)); } - // If the request is associated with a group not enabled, will fall back to the default one. - return Mono.just(new Utils.ValueHolder<>(this.defaultGroupController)); + return this.resolveThroughputGroupController(group).map(Utils.ValueHolder::new); } public String getTargetContainerRid() { @@ -333,13 +332,13 @@ public boolean canHandleRequest(RxDocumentServiceRequest request) { return StringUtils.equals(this.targetContainerRid, request.requestContext.resolvedCollectionRid); } - private Mono createAndInitializeGroupControllers() { - return Flux.fromIterable(this.groups) + private Mono createAndInitializeGroupControllers() { + return Flux.fromIterable(this.groups.values()) .flatMap(this::resolveThroughputGroupController) .then(Mono.just(this)); } - private Mono resolveThroughputGroupController(ThroughputControlGroupInternal group) { + private Mono resolveThroughputGroupController(SDKThroughputControlGroupInternal group) { return this.groupControllerCache.getAsync( group.getGroupName(), null, @@ -347,13 +346,13 @@ private Mono resolveThroughputGroupController(Thr .onErrorResume(throwable -> Mono.error(new ThroughputControlInitializationException(throwable))); } - private Mono createAndInitializeGroupController(ThroughputControlGroupInternal group) { + private Mono createAndInitializeGroupController(SDKThroughputControlGroupInternal group) { LinkedCancellationToken parentToken = this.cancellationTokenMap.compute( group.getGroupName(), (key, cancellationToken) -> this.cancellationTokenSource.getToken()); - ThroughputGroupControllerBase groupController = ThroughputGroupControllerFactory.createController( + SDKThroughputGroupControllerBase groupController = ThroughputGroupControllerFactory.createController( this.connectionMode, group, this.maxContainerThroughput.get(), @@ -363,7 +362,7 @@ private Mono createAndInitializeGroupController(T return groupController .init() - .cast(ThroughputGroupControllerBase.class) + .cast(SDKThroughputGroupControllerBase.class) .doOnSuccess(controller -> { if (controller.isDefault()) { this.defaultGroupController = controller; @@ -387,7 +386,7 @@ private Flux refreshContainerMaxThroughputTask(LinkedCancellationToken can return this.resolveContainerMaxThroughput(); } }) - .flatMapIterable(controller -> this.groups) + .flatMapIterable(controller -> this.groups.values()) .flatMap(this::resolveThroughputGroupController) .doOnNext(groupController -> groupController.onContainerMaxThroughputRefresh(this.maxContainerThroughput.get())) .onErrorResume(throwable -> { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/ThroughputProvisioningScope.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/container/ThroughputProvisioningScope.java similarity index 68% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/ThroughputProvisioningScope.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/container/ThroughputProvisioningScope.java index a161cb8ed830..c61a06dd9b71 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/container/ThroughputProvisioningScope.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/container/ThroughputProvisioningScope.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.container; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.container; public enum ThroughputProvisioningScope { NONE, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/ThroughputGroupControllerBase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/SDKThroughputGroupControllerBase.java similarity index 87% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/ThroughputGroupControllerBase.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/SDKThroughputGroupControllerBase.java index 54bae311e044..c9cc558050e1 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/ThroughputGroupControllerBase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/SDKThroughputGroupControllerBase.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.group; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.group; import com.azure.cosmos.ConnectionMode; import com.azure.cosmos.CosmosException; @@ -11,14 +11,15 @@ import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.implementation.caches.AsyncCache; import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; -import com.azure.cosmos.implementation.throughputControl.LinkedCancellationToken; -import com.azure.cosmos.implementation.throughputControl.LinkedCancellationTokenSource; -import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal; -import com.azure.cosmos.implementation.throughputControl.controller.IThroughputController; -import com.azure.cosmos.implementation.throughputControl.controller.request.GlobalThroughputRequestController; -import com.azure.cosmos.implementation.throughputControl.controller.request.IThroughputRequestController; -import com.azure.cosmos.implementation.throughputControl.controller.request.PkRangesThroughputRequestController; -import com.azure.cosmos.implementation.throughputControl.exceptions.ThroughputControlInitializationException; +import com.azure.cosmos.implementation.throughputControl.ThroughputControlRequestContext; +import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationToken; +import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationTokenSource; +import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.IThroughputController; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.request.GlobalThroughputRequestController; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.request.IThroughputRequestController; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.request.PkRangesThroughputRequestController; +import com.azure.cosmos.implementation.throughputControl.sdk.exceptions.ThroughputControlInitializationException; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.Exceptions; @@ -39,12 +40,12 @@ * 1. Create and initialize request controller based on connection mode * 2. Schedule reset throughput usage every 1s. */ -public abstract class ThroughputGroupControllerBase implements IThroughputController { - private final static Logger logger = LoggerFactory.getLogger(ThroughputGroupControllerBase.class); +public abstract class SDKThroughputGroupControllerBase implements IThroughputController { + private final static Logger logger = LoggerFactory.getLogger(SDKThroughputGroupControllerBase.class); private final Duration DEFAULT_THROUGHPUT_USAGE_RESET_DURATION = Duration.ofSeconds(1); private final ConnectionMode connectionMode; - private final ThroughputControlGroupInternal group; + private final SDKThroughputControlGroupInternal group; private final AtomicInteger maxContainerThroughput; private final RxPartitionKeyRangeCache partitionKeyRangeCache; private final AsyncCache requestControllerAsyncCache; @@ -53,9 +54,9 @@ public abstract class ThroughputGroupControllerBase implements IThroughputContro protected final AtomicReference groupThroughput; protected final LinkedCancellationTokenSource cancellationTokenSource; - public ThroughputGroupControllerBase( + public SDKThroughputGroupControllerBase( ConnectionMode connectionMode, - ThroughputControlGroupInternal group, + SDKThroughputControlGroupInternal group, Integer maxContainerThroughput, RxPartitionKeyRangeCache partitionKeyRangeCache, String targetContainerRid, @@ -155,6 +156,11 @@ public void onContainerMaxThroughputRefresh(int maxContainerThroughput) { @Override public Mono processRequest(RxDocumentServiceRequest request, Mono originalRequestMono) { + if (request.requestContext != null) { + request.requestContext.throughputControlRequestContext = + new ThroughputControlRequestContext(this.group.getDiagnosticsString()); + } + return this.resolveRequestController() .flatMap(requestController -> { if (requestController.canHandleRequest(request)) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/ThroughputGroupControllerFactory.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/ThroughputGroupControllerFactory.java similarity index 61% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/ThroughputGroupControllerFactory.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/ThroughputGroupControllerFactory.java index 01d1c9dbb040..7a26cf4ef508 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/ThroughputGroupControllerFactory.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/ThroughputGroupControllerFactory.java @@ -1,23 +1,22 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.group; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.group; import com.azure.cosmos.ConnectionMode; -import com.azure.cosmos.implementation.GlobalEndpointManager; import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; -import com.azure.cosmos.implementation.throughputControl.LinkedCancellationToken; -import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal; -import com.azure.cosmos.implementation.throughputControl.config.GlobalThroughputControlGroup; -import com.azure.cosmos.implementation.throughputControl.config.LocalThroughputControlGroup; -import com.azure.cosmos.implementation.throughputControl.controller.group.global.GlobalThroughputControlGroupController; -import com.azure.cosmos.implementation.throughputControl.controller.group.local.LocalThroughputControlGroupController; +import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationToken; +import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal; +import com.azure.cosmos.implementation.throughputControl.sdk.config.GlobalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.sdk.config.LocalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.group.global.GlobalThroughputControlGroupController; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.group.local.LocalThroughputControlGroupController; public class ThroughputGroupControllerFactory { - public static ThroughputGroupControllerBase createController( + public static SDKThroughputGroupControllerBase createController( ConnectionMode connectionMode, - ThroughputControlGroupInternal group, + SDKThroughputControlGroupInternal group, Integer maxContainerThroughput, RxPartitionKeyRangeCache partitionKeyRangeCache, String targetCollectionRid, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlClientItem.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlClientItem.java similarity index 94% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlClientItem.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlClientItem.java index a20578e3f4e9..5e220c89261e 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlClientItem.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlClientItem.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.group.global; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.group.global; import com.fasterxml.jackson.annotation.JsonProperty; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlConfigItem.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlConfigItem.java similarity index 97% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlConfigItem.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlConfigItem.java index b1578ea9d435..314a49bd9c41 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlConfigItem.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlConfigItem.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.group.global; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.group.global; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.fasterxml.jackson.annotation.JsonIgnore; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlGroupController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlGroupController.java similarity index 92% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlGroupController.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlGroupController.java index 1916c38de7f3..0beebbcd8105 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlGroupController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlGroupController.java @@ -1,15 +1,15 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.group.global; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.group.global; import com.azure.cosmos.ConnectionMode; import com.azure.cosmos.implementation.CosmosSchedulers; import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; import com.azure.cosmos.implementation.guava25.collect.EvictingQueue; -import com.azure.cosmos.implementation.throughputControl.LinkedCancellationToken; -import com.azure.cosmos.implementation.throughputControl.config.GlobalThroughputControlGroup; -import com.azure.cosmos.implementation.throughputControl.controller.group.ThroughputGroupControllerBase; +import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationToken; +import com.azure.cosmos.implementation.throughputControl.sdk.config.GlobalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.group.SDKThroughputGroupControllerBase; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; @@ -19,7 +19,7 @@ import java.time.Instant; import java.util.concurrent.atomic.AtomicReference; -public class GlobalThroughputControlGroupController extends ThroughputGroupControllerBase { +public class GlobalThroughputControlGroupController extends SDKThroughputGroupControllerBase { private static final Logger logger = LoggerFactory.getLogger(GlobalThroughputControlGroupController.class); private static final double INITIAL_CLIENT_THROUGHPUT_RU_SHARE = 1.0; private static final double INITIAL_THROUGHPUT_USAGE = 1.0; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlItem.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlItem.java similarity index 93% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlItem.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlItem.java index bfaeb5e2ba04..af6d87c239bf 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/GlobalThroughputControlItem.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/GlobalThroughputControlItem.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.group.global; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.group.global; import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.annotation.JsonProperty; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/ThroughputControlContainerManager.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/ThroughputControlContainerManager.java similarity index 98% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/ThroughputControlContainerManager.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/ThroughputControlContainerManager.java index fd382a10ead9..e5c87df99d43 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/ThroughputControlContainerManager.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/ThroughputControlContainerManager.java @@ -1,14 +1,14 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.group.global; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.group.global; import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.CosmosException; import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.UUIDs; import com.azure.cosmos.implementation.Utils; -import com.azure.cosmos.implementation.throughputControl.config.GlobalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.sdk.config.GlobalThroughputControlGroup; import com.azure.cosmos.models.CosmosItemRequestOptions; import com.azure.cosmos.models.PartitionKey; import com.azure.cosmos.models.SqlParameter; @@ -23,7 +23,6 @@ import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.List; -import java.util.UUID; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/ThroughputUsageSnapshot.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/ThroughputUsageSnapshot.java similarity index 92% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/ThroughputUsageSnapshot.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/ThroughputUsageSnapshot.java index e26c038cd4e9..c2b23da2e5c7 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/global/ThroughputUsageSnapshot.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/global/ThroughputUsageSnapshot.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.group.global; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.group.global; import java.time.Duration; import java.time.Instant; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/local/LocalThroughputControlGroupController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/local/LocalThroughputControlGroupController.java similarity index 74% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/local/LocalThroughputControlGroupController.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/local/LocalThroughputControlGroupController.java index 2c5410f5df6e..5e9498fe5924 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/group/local/LocalThroughputControlGroupController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/group/local/LocalThroughputControlGroupController.java @@ -1,17 +1,17 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.group.local; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.group.local; import com.azure.cosmos.ConnectionMode; import com.azure.cosmos.implementation.CosmosSchedulers; import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; -import com.azure.cosmos.implementation.throughputControl.LinkedCancellationToken; -import com.azure.cosmos.implementation.throughputControl.config.LocalThroughputControlGroup; -import com.azure.cosmos.implementation.throughputControl.controller.group.ThroughputGroupControllerBase; +import com.azure.cosmos.implementation.throughputControl.sdk.LinkedCancellationToken; +import com.azure.cosmos.implementation.throughputControl.sdk.config.LocalThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.sdk.controller.group.SDKThroughputGroupControllerBase; import reactor.core.publisher.Mono; -public class LocalThroughputControlGroupController extends ThroughputGroupControllerBase { +public class LocalThroughputControlGroupController extends SDKThroughputGroupControllerBase { public LocalThroughputControlGroupController( ConnectionMode connectionMode, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/request/GlobalThroughputRequestController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/request/GlobalThroughputRequestController.java similarity index 89% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/request/GlobalThroughputRequestController.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/request/GlobalThroughputRequestController.java index 6fdec7d63b4f..e2424a7ee0fb 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/request/GlobalThroughputRequestController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/request/GlobalThroughputRequestController.java @@ -1,11 +1,11 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.request; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.request; import com.azure.cosmos.implementation.RxDocumentServiceRequest; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; -import com.azure.cosmos.implementation.throughputControl.ThroughputRequestThrottler; +import com.azure.cosmos.implementation.throughputControl.sdk.ThroughputRequestThrottler; import reactor.core.publisher.Mono; import java.util.concurrent.atomic.AtomicReference; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/request/IThroughputRequestController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/request/IThroughputRequestController.java similarity index 81% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/request/IThroughputRequestController.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/request/IThroughputRequestController.java index 2e1fcb4437c3..e2c0e541ceed 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/request/IThroughputRequestController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/request/IThroughputRequestController.java @@ -1,9 +1,9 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.request; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.request; -import com.azure.cosmos.implementation.throughputControl.controller.IThroughputController; +import com.azure.cosmos.implementation.throughputControl.IThroughputController; /** * Represents a throughput request controller. diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/request/PkRangesThroughputRequestController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/request/PkRangesThroughputRequestController.java similarity index 96% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/request/PkRangesThroughputRequestController.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/request/PkRangesThroughputRequestController.java index 9d57c793e249..fd7f56918e3f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/controller/request/PkRangesThroughputRequestController.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/controller/request/PkRangesThroughputRequestController.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.controller.request; +package com.azure.cosmos.implementation.throughputControl.sdk.controller.request; import com.azure.cosmos.implementation.PartitionKeyRange; import com.azure.cosmos.implementation.RxDocumentServiceRequest; @@ -9,7 +9,7 @@ import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper; import com.azure.cosmos.implementation.routing.Range; -import com.azure.cosmos.implementation.throughputControl.ThroughputRequestThrottler; +import com.azure.cosmos.implementation.throughputControl.sdk.ThroughputRequestThrottler; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Mono; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/exceptions/ThroughputControlInitializationException.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/exceptions/ThroughputControlInitializationException.java similarity index 93% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/exceptions/ThroughputControlInitializationException.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/exceptions/ThroughputControlInitializationException.java index ece6c75d1c79..a1345fe0c64b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/exceptions/ThroughputControlInitializationException.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/sdk/exceptions/ThroughputControlInitializationException.java @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -package com.azure.cosmos.implementation.throughputControl.exceptions; +package com.azure.cosmos.implementation.throughputControl.sdk.exceptions; import com.azure.cosmos.implementation.throughputControl.ThroughputControlStore; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ContainerServerThroughputControlGroupProperties.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ContainerServerThroughputControlGroupProperties.java new file mode 100644 index 000000000000..472302f6e405 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ContainerServerThroughputControlGroupProperties.java @@ -0,0 +1,68 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl.server; + +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroup; + +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicReference; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +public class ContainerServerThroughputControlGroupProperties { + + private final String containerNameLink; + private final AtomicReference defaultGroup; + private final Map throughputControlGroups; + + public ContainerServerThroughputControlGroupProperties(String containerNameLink) { + checkArgument(StringUtils.isNotEmpty(containerNameLink), "Argument 'containerNameLink' should not be empty"); + + this.containerNameLink = containerNameLink; + this.defaultGroup = new AtomicReference<>(); + this.throughputControlGroups = new ConcurrentHashMap<>(); + } + + /*** + * Enable a server throughput control group. + * + * @param group a {@link ServerThroughputControlGroup}. + * + * @return the total size of distinct server throughput control groups enabled on the container. + */ + public int enableThroughputControlGroup(ServerThroughputControlGroup group) { + checkNotNull(group, "Argument 'group' should not be null'"); + + if (group.isDefault()) { + if (!this.defaultGroup.compareAndSet(null, group)) { + if (!this.defaultGroup.get().equals(group)) { + throw new IllegalArgumentException("A default group already exists"); + } + } + } + + ServerThroughputControlGroup serverThroughputControlGroup = + this.throughputControlGroups.computeIfAbsent(group.getGroupName(), (key) -> group); + if (!serverThroughputControlGroup.equals(group)) { + throw new IllegalArgumentException("A group with same name already exists, name: " + group.getGroupName()); + } + + return this.throughputControlGroups.size(); + } + + public Map getThroughputControlGroups() { + return this.throughputControlGroups; + } + + public boolean hasDefaultGroup() { + return this.defaultGroup.get() != null; + } + + public boolean hasGroup(String groupName) { + return this.throughputControlGroups.containsKey(groupName); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ServerThroughputControlStore.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ServerThroughputControlStore.java new file mode 100644 index 000000000000..1146e2e72bf9 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/ServerThroughputControlStore.java @@ -0,0 +1,109 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl.server; + +import com.azure.cosmos.BridgeInternal; +import com.azure.cosmos.implementation.ResourceType; +import com.azure.cosmos.implementation.RxDocumentServiceRequest; +import com.azure.cosmos.implementation.Utils; +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.implementation.caches.AsyncCache; +import com.azure.cosmos.implementation.throughputControl.EmptyThroughputContainerController; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroup; +import com.azure.cosmos.implementation.throughputControl.server.controller.ServerThroughputContainerController; +import reactor.core.publisher.Mono; + +import java.util.concurrent.ConcurrentHashMap; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +public class ServerThroughputControlStore { + + private final AsyncCache containerControllerCache; + private final ConcurrentHashMap containerMap; + + public ServerThroughputControlStore() { + this.containerControllerCache = new AsyncCache<>(); + this.containerMap = new ConcurrentHashMap<>(); + } + + public void enableThroughputControlGroup(ServerThroughputControlGroup group) { + checkNotNull(group, "Throughput control group cannot be null"); + + String containerNameLink = Utils.trimBeginningAndEndingSlashes(BridgeInternal.extractContainerSelfLink(group.getTargetContainer())); + this.containerMap.compute(containerNameLink, (key, throughputControlContainerProperties) -> { + if (throughputControlContainerProperties == null) { + throughputControlContainerProperties = new ContainerServerThroughputControlGroupProperties(containerNameLink); + } + + int groupSizeAfterEnabling = throughputControlContainerProperties.enableThroughputControlGroup(group); + if (groupSizeAfterEnabling == 1) { + // This is the first enabled group for the target container or an existing group was modified + // Clean the current cache in case we have built EmptyThroughputContainerController or an existing + // group was modified + this.containerControllerCache.remove(containerNameLink); + } + + return throughputControlContainerProperties; + }); + } + + public Mono processRequest(RxDocumentServiceRequest request, Mono originalRequestMono) { + checkNotNull(request, "Request can not be null"); + checkNotNull(originalRequestMono, "originalRequestMono can not be null"); + + // Currently, we will only target two resource types. + // If in the future we find other useful scenarios for throughput control, add more resource type here. + if (request.getResourceType() != ResourceType.Document && request.getResourceType() != ResourceType.StoredProcedure) { + return originalRequestMono; + } + + String collectionNameLink = Utils.getCollectionName(request.getResourceAddress()); + return this.resolveContainerController(collectionNameLink) + .flatMap(containerController -> containerController.processRequest(request, originalRequestMono)); + } + + public boolean hasDefaultGroup(String containerNameLink) { + if (this.containerMap.containsKey(containerNameLink)) { + return this.containerMap.get(containerNameLink).hasDefaultGroup(); + } + + return false; + } + + public boolean hasGroup(String containerNameLink, String throughputControlGroupName) { + if (StringUtils.isEmpty(throughputControlGroupName)) { + return false; + } + + if (this.containerMap.containsKey(containerNameLink)) { + return this.containerMap.get(containerNameLink).hasGroup(throughputControlGroupName); + } + + return false; + } + + private Mono resolveContainerController(String containerNameLink) { + return this.containerControllerCache.getAsync( + containerNameLink, + null, + () -> this.createAndInitContainerController(containerNameLink)); + } + + private Mono createAndInitContainerController(String containerNameLink) { + checkArgument(StringUtils.isNotEmpty(containerNameLink), "Container link should not be null or empty"); + + if (this.containerMap.containsKey(containerNameLink)) { + ContainerServerThroughputControlGroupProperties containerProperties = + this.containerMap.get(containerNameLink); + return Mono + .just(new ServerThroughputContainerController(containerProperties.getThroughputControlGroups())) + .flatMap(ServerThroughputContainerController::init); + } else { + return Mono.just(new EmptyThroughputContainerController()) + .flatMap(EmptyThroughputContainerController::init); + } + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ServerThroughputControlGroup.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ServerThroughputControlGroup.java new file mode 100644 index 000000000000..ee317cb9a096 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/config/ServerThroughputControlGroup.java @@ -0,0 +1,102 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl.server.config; + +import com.azure.cosmos.CosmosAsyncContainer; +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.implementation.throughputControl.IThroughputControlGroup; +import com.azure.cosmos.models.PriorityLevel; + +import java.util.Objects; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +public class ServerThroughputControlGroup implements IThroughputControlGroup { + private final String groupName; + private final boolean isDefault; + private final CosmosAsyncContainer targetContainer; + private final PriorityLevel priorityLevel; + private final Integer throughputBucket; + + public ServerThroughputControlGroup( + String groupName, + boolean isDefault, + PriorityLevel priorityLevel, + Integer throughputBucket, + CosmosAsyncContainer targetContainer) { + + checkArgument(StringUtils.isNotEmpty(groupName), "Argument 'groupName' cannot be null or empty."); + checkNotNull(targetContainer, "Argument 'targetContainer' can not be null"); + checkArgument( + priorityLevel != null || throughputBucket != null, + "At least one of 'priorityLevel' or 'throughputBucket' must be provided."); + checkArgument(throughputBucket == null || throughputBucket >= 0, "Target throughput bucket should be no less than 0"); + + this.groupName = groupName; + this.isDefault = isDefault; + this.targetContainer = targetContainer; + this.priorityLevel = priorityLevel; + this.throughputBucket = throughputBucket; + } + + public String getGroupName() { + return this.groupName; + } + + public boolean isDefault() { + return this.isDefault; + } + + public CosmosAsyncContainer getTargetContainer() { + return this.targetContainer; + } + + public PriorityLevel getPriorityLevel() { + return this.priorityLevel; + } + + public Integer getThroughputBucket() { + return this.throughputBucket; + } + + @Override + public boolean equals(Object o) { + if (o == null || getClass() != o.getClass()) return false; + ServerThroughputControlGroup that = (ServerThroughputControlGroup) o; + return isDefault == that.isDefault + && Objects.equals(groupName, that.groupName) + && Objects.equals(targetContainer, that.targetContainer) + && Objects.equals(priorityLevel, that.priorityLevel) + && Objects.equals(throughputBucket, that.throughputBucket); + } + + @Override + public int hashCode() { + return Objects.hash(groupName, isDefault, targetContainer, priorityLevel, throughputBucket); + } + + /*** + * Config show up in the diagnostics. + * @return the config string. + */ + @Override + public String getDiagnosticsString() { + StringBuilder sb = new StringBuilder(); + sb.append("("); + sb.append("name=" + this.groupName); + sb.append(", default=" + this.isDefault); + + if (this.priorityLevel != null) { + sb.append(", priorityLevel=" + this.priorityLevel); + } + + if (this.throughputBucket != null) { + sb.append(", throughputBucket=" + this.throughputBucket); + } + + sb.append(")"); + return sb.toString(); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputContainerController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputContainerController.java new file mode 100644 index 000000000000..4aa3e51a97d9 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputContainerController.java @@ -0,0 +1,100 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl.server.controller; + +import com.azure.cosmos.implementation.RxDocumentServiceRequest; +import com.azure.cosmos.implementation.Utils; +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.implementation.caches.AsyncCache; +import com.azure.cosmos.implementation.throughputControl.IThroughputContainerController; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroup; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +import java.util.Map; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +public class ServerThroughputContainerController implements IThroughputContainerController { + private static final Logger logger = LoggerFactory.getLogger(ServerThroughputContainerController.class); + + private final AsyncCache groupControllerCache; + private final Map groups; + + private ServerThroughputGroupController defaultGroupController; + public ServerThroughputContainerController(Map groups) { + + checkArgument(groups != null && !groups.isEmpty(), "Throughput groups can not be null or empty"); + + this.groupControllerCache = new AsyncCache<>(); + this.groups = groups; + } + + @Override + @SuppressWarnings("unchecked") + public Mono init() { + return this.createAndInitializeGroupControllers().thenReturn((T) this); + } + + @Override + public Mono processRequest(RxDocumentServiceRequest request, Mono originalRequestMono) { + checkNotNull(request, "Request can not be null"); + checkNotNull(originalRequestMono, "Original request mono can not be null"); + + return this.getOrCreateThroughputGroupController(request.getThroughputControlGroupName()) + .flatMap(groupController -> { + if (groupController.v != null) { + return groupController.v.processRequest(request, originalRequestMono); + } + + return originalRequestMono; + }); + } + + private Mono> getOrCreateThroughputGroupController(String groupName) { + + // If there is no control group defined, using the default group controller + if (StringUtils.isEmpty(groupName)) { + return Mono.just(new Utils.ValueHolder<>(this.defaultGroupController)); + } + + ServerThroughputControlGroup group = this.groups.get(groupName); + if (group == null) { + // If the request is associated with a group not enabled, will fall back to the default one. + return Mono.just(new Utils.ValueHolder<>(this.defaultGroupController)); + } + + return this.resolveThroughputGroupController(group).map(Utils.ValueHolder::new); + } + + @Override + public boolean canHandleRequest(RxDocumentServiceRequest request) { + return true; + } + + private Mono createAndInitializeGroupControllers() { + return Flux.fromIterable(this.groups.values()) + .flatMap(this::resolveThroughputGroupController) + .then(Mono.just(this)); + } + + private Mono resolveThroughputGroupController(ServerThroughputControlGroup group) { + return this.groupControllerCache.getAsync( + group.getGroupName(), + null, + () -> this.createAndInitializeGroupController(group)); + } + + private Mono createAndInitializeGroupController(ServerThroughputControlGroup group) { + // create throughput bucket group controller + ServerThroughputGroupController throughputGroupController = new ServerThroughputGroupController(group); + if (throughputGroupController.isDefault()) { + this.defaultGroupController = throughputGroupController; + } + return Mono.just(throughputGroupController); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupController.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupController.java new file mode 100644 index 000000000000..b419b5778939 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/throughputControl/server/controller/ServerThroughputGroupController.java @@ -0,0 +1,55 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.throughputControl.server.controller; + +import com.azure.cosmos.implementation.RxDocumentServiceRequest; +import com.azure.cosmos.implementation.throughputControl.IThroughputController; +import com.azure.cosmos.implementation.throughputControl.ThroughputControlRequestContext; +import com.azure.cosmos.implementation.throughputControl.server.config.ServerThroughputControlGroup; +import reactor.core.publisher.Mono; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +public class ServerThroughputGroupController implements IThroughputController { + private final ServerThroughputControlGroup serverThroughputControlGroup; + + public ServerThroughputGroupController(ServerThroughputControlGroup serverThroughputControlGroup) { + checkNotNull(serverThroughputControlGroup, "Argument 'serverThroughputControlGroup' cannot be null."); + this.serverThroughputControlGroup = serverThroughputControlGroup; + } + + @Override + public boolean canHandleRequest(RxDocumentServiceRequest request) { + return true; + } + + @Override + @SuppressWarnings("unchecked") + public Mono init() { + return (Mono) Mono.just(this); + } + + @Override + public Mono processRequest(RxDocumentServiceRequest request, Mono originalRequestMono) { + if (this.serverThroughputControlGroup.getPriorityLevel() != null) { + request.setPriorityLevel(this.serverThroughputControlGroup.getPriorityLevel()); + } + + if (this.serverThroughputControlGroup.getThroughputBucket() != null) { + request.setThroughputBucket(this.serverThroughputControlGroup.getThroughputBucket()); + } + + if (request.requestContext != null) { + request.requestContext.setThroughputControlRequestContext( + new ThroughputControlRequestContext(this.serverThroughputControlGroup.getDiagnosticsString()) + ); + } + + return originalRequestMono; + } + + public boolean isDefault() { + return this.serverThroughputControlGroup.isDefault(); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java index 857cb893ed4b..5ccd99db2fd0 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java @@ -105,6 +105,8 @@ public enum SinceVersion { /** v4.69.0 */ V4_69_0, /** v4.71.0 */ - V4_71_0 + V4_71_0, + /** v4.74.0 */ + V4_74_0 } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/module-info.java b/sdk/cosmos/azure-cosmos/src/main/java/module-info.java index 345b4a00e929..82b7f11e5056 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/module-info.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/module-info.java @@ -69,7 +69,7 @@ opens com.azure.cosmos.implementation.clienttelemetry to com.fasterxml.jackson.databind; opens com.azure.cosmos.util to com.fasterxml.jackson.databind; opens com.azure.cosmos.implementation.throughputControl to com.fasterxml.jackson.databind; - opens com.azure.cosmos.implementation.throughputControl.controller.group.global to com.fasterxml.jackson.databind; + opens com.azure.cosmos.implementation.throughputControl.sdk.controller.group.global to com.fasterxml.jackson.databind; opens com.azure.cosmos.implementation.perPartitionCircuitBreaker to com.fasterxml.jackson.databind; opens com.azure.cosmos.implementation.perPartitionAutomaticFailover to com.fasterxml.jackson.databind; @@ -79,6 +79,8 @@ exports com.azure.cosmos.implementation.routing to com.azure.cosmos.test; opens com.azure.cosmos to com.azure.cosmos.test, com.azure.spring.data.cosmos, com.fasterxml.jackson.databind, com.fasterxml.jackson.module.afterburner, java.logging; opens com.azure.cosmos.models to com.azure.cosmos.test, com.azure.spring.data.cosmos, com.fasterxml.jackson.databind, com.fasterxml.jackson.module.afterburner, java.logging; + opens com.azure.cosmos.implementation.throughputControl.sdk to com.fasterxml.jackson.databind; + opens com.azure.cosmos.implementation.throughputControl.sdk.config to com.fasterxml.jackson.databind; uses com.azure.cosmos.implementation.guava25.base.PatternCompiler; uses com.azure.core.util.tracing.Tracer; diff --git a/sdk/cosmos/azure-cosmos/src/samples/java/com/azure/cosmos/ThroughputControlCodeSnippet.java b/sdk/cosmos/azure-cosmos/src/samples/java/com/azure/cosmos/ThroughputControlCodeSnippet.java index 51d482c8efe9..22c173f928dd 100644 --- a/sdk/cosmos/azure-cosmos/src/samples/java/com/azure/cosmos/ThroughputControlCodeSnippet.java +++ b/sdk/cosmos/azure-cosmos/src/samples/java/com/azure/cosmos/ThroughputControlCodeSnippet.java @@ -53,4 +53,16 @@ public void codeSnippetForEnableGlobalThroughputControl() { container.enableGlobalThroughputControlGroup(groupConfig, globalControlConfig); // END: com.azure.cosmos.throughputControl.globalControl } + + public void codeSnippetForEnableServerThroughputControl() { + // BEGIN: com.azure.cosmos.throughputControl.serverControl + ThroughputControlGroupConfig groupConfig = + new ThroughputControlGroupConfigBuilder() + .groupName("localControlGroup") + .throughputBucket(2) + .build(); + + container.enableServerThroughputControlGroup(groupConfig); + // END: com.azure.cosmos.throughputControl.serverControl + } }