forked from Azure/azure-sdk-for-java
-
Notifications
You must be signed in to change notification settings - Fork 1
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Fix for bulk operations incorrect routing on split (Azure#21420)
* Fixing a scenario where bulk operations fail with 410 on container split. * Splits are taking longer somehow! So increasing the timeout for QueryValidationTests::splitQueryContinuationToken * Splits are taking longer somehow! So increasing the timeout for readFeedDocumentsAfterSplit() * Refactoring code Implementing PR comments Increasing timeout on tests
- Loading branch information
Showing
7 changed files
with
245 additions
and
47 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
177 changes: 177 additions & 0 deletions
177
sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosBulkGatewayTest.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,177 @@ | ||
| // Copyright (c) Microsoft Corporation. All rights reserved. | ||
| // Licensed under the MIT License. | ||
| package com.azure.cosmos; | ||
|
|
||
| import com.azure.cosmos.implementation.AsyncDocumentClient; | ||
| import com.azure.cosmos.implementation.PartitionKeyRange; | ||
| import com.azure.cosmos.models.CosmosContainerProperties; | ||
| import com.azure.cosmos.models.CosmosContainerResponse; | ||
| import com.azure.cosmos.models.CosmosQueryRequestOptions; | ||
| import com.azure.cosmos.models.FeedResponse; | ||
| import com.azure.cosmos.models.PartitionKey; | ||
| import com.azure.cosmos.models.ThroughputProperties; | ||
| import com.azure.cosmos.models.ThroughputResponse; | ||
| import io.netty.handler.codec.http.HttpResponseStatus; | ||
| import org.jetbrains.annotations.NotNull; | ||
| import org.slf4j.Logger; | ||
| import org.slf4j.LoggerFactory; | ||
| import org.testng.annotations.AfterClass; | ||
| import org.testng.annotations.BeforeClass; | ||
| import org.testng.annotations.Factory; | ||
| import org.testng.annotations.Test; | ||
| import reactor.core.publisher.Flux; | ||
| import reactor.core.publisher.Mono; | ||
|
|
||
| import java.util.ArrayList; | ||
| import java.util.List; | ||
| import java.util.Random; | ||
| import java.util.UUID; | ||
| import java.util.concurrent.atomic.AtomicInteger; | ||
|
|
||
| import static org.assertj.core.api.Assertions.assertThat; | ||
|
|
||
| public class CosmosBulkGatewayTest extends BatchTestBase { | ||
| private final static Logger logger = LoggerFactory.getLogger(CosmosBulkAsyncTest.class); | ||
|
|
||
| private CosmosAsyncClient bulkClient; | ||
| private CosmosAsyncDatabase createdDatabase; | ||
|
|
||
| @Factory(dataProvider = "simpleClientBuilderGatewaySession") | ||
| public CosmosBulkGatewayTest(CosmosClientBuilder clientBuilder) { | ||
| super(clientBuilder); | ||
| } | ||
|
|
||
| @BeforeClass(groups = {"simple"}, timeOut = SETUP_TIMEOUT) | ||
| public void before_CosmosBulkAsyncTest() { | ||
| assertThat(this.bulkClient).isNull(); | ||
| this.bulkClient = getClientBuilder().buildAsyncClient(); | ||
| createdDatabase = getSharedCosmosDatabase(this.bulkClient); | ||
| } | ||
|
|
||
| @AfterClass(groups = {"simple"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) | ||
| public void afterClass() { | ||
| safeCloseAsync(this.bulkClient); | ||
| } | ||
|
|
||
| @Test(groups = {"simple"}, timeOut = TIMEOUT * 20) | ||
| public void createItem_withBulk_split() throws InterruptedException { | ||
| String containerId = "bulksplittestcontainer_" + UUID.randomUUID(); | ||
| int totalRequest = getTotalRequest(); | ||
| CosmosContainerProperties containerProperties = new CosmosContainerProperties(containerId, "/mypk"); | ||
| CosmosContainerResponse containerResponse = createdDatabase.createContainer(containerProperties).block(); | ||
| CosmosAsyncContainer container = createdDatabase.getContainer(containerId); | ||
|
|
||
| Flux<CosmosItemOperation> cosmosItemOperationFlux1 = Flux.range(0, totalRequest).map(i -> { | ||
| String partitionKey = UUID.randomUUID().toString(); | ||
| TestDoc testDoc = this.populateTestDoc(partitionKey); | ||
|
|
||
| return BulkOperations.getCreateItemOperation(testDoc, new PartitionKey(partitionKey)); | ||
| }); | ||
|
|
||
| Flux<CosmosItemOperation> cosmosItemOperationFlux2 = Flux.range(0, totalRequest).map(i -> { | ||
| String partitionKey = UUID.randomUUID().toString(); | ||
| EventDoc eventDoc = new EventDoc(UUID.randomUUID().toString(), 2, 4, "type1", partitionKey); | ||
|
|
||
| return BulkOperations.getCreateItemOperation(eventDoc, new PartitionKey(partitionKey)); | ||
| }); | ||
|
|
||
| BulkProcessingOptions<CosmosBulkAsyncTest> bulkProcessingOptions = new BulkProcessingOptions<>(); | ||
| bulkProcessingOptions.setMaxMicroBatchSize(100); | ||
| bulkProcessingOptions.setMaxMicroBatchConcurrency(5); | ||
|
|
||
| Flux<CosmosBulkOperationResponse<CosmosBulkAsyncTest>> responseFlux = | ||
| container.processBulkOperations(cosmosItemOperationFlux1, bulkProcessingOptions); | ||
|
|
||
| AtomicInteger processedDoc = new AtomicInteger(0); | ||
| responseFlux | ||
| .flatMap((CosmosBulkOperationResponse<CosmosBulkAsyncTest> cosmosBulkOperationResponse) -> { | ||
|
|
||
| processedDoc.incrementAndGet(); | ||
|
|
||
| CosmosBulkItemResponse cosmosBulkItemResponse = cosmosBulkOperationResponse.getResponse(); | ||
| assertThat(cosmosBulkItemResponse.getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); | ||
| assertThat(cosmosBulkItemResponse.getRequestCharge()).isGreaterThan(0); | ||
| assertThat(cosmosBulkItemResponse.getCosmosDiagnostics().toString()).isNotNull(); | ||
| assertThat(cosmosBulkItemResponse.getSessionToken()).isNotNull(); | ||
| assertThat(cosmosBulkItemResponse.getActivityId()).isNotNull(); | ||
| assertThat(cosmosBulkItemResponse.getRequestCharge()).isNotNull(); | ||
|
|
||
| return Mono.just(cosmosBulkItemResponse); | ||
| }).blockLast(); | ||
|
|
||
| assertThat(processedDoc.get()).isEqualTo(totalRequest); | ||
|
|
||
| // introduce a split and continue bulk operations after split. The partition key range cache has to be | ||
| // refreshed and bulk processing should complete without errors | ||
| List<PartitionKeyRange> partitionKeyRanges = getPartitionKeyRanges(containerId, this.bulkClient); | ||
| // Scale up the throughput for a split | ||
| logger.info("Scaling up throughput for split"); | ||
| ThroughputProperties throughputProperties = ThroughputProperties.createManualThroughput(16000); | ||
| ThroughputResponse throughputResponse = container.replaceThroughput(throughputProperties).block(); | ||
| logger.info("Throughput replace request submitted for {} ", | ||
| throughputResponse.getProperties().getManualThroughput()); | ||
| throughputResponse = container.readThroughput().block(); | ||
|
|
||
|
|
||
| // Wait for the throughput update to complete so that we get the partition split | ||
| while (true) { | ||
| assert throughputResponse != null; | ||
| if (!throughputResponse.isReplacePending()) { | ||
| break; | ||
| } | ||
| logger.info("Waiting for split to complete"); | ||
| Thread.sleep(10 * 1000); | ||
| throughputResponse = container.readThroughput().block(); | ||
| } | ||
|
|
||
| // Read number of partitions. Should be greater than one | ||
| List<PartitionKeyRange> partitionKeyRangesAfterSplit = getPartitionKeyRanges(containerId, | ||
| this.bulkClient); | ||
| assertThat(partitionKeyRangesAfterSplit.size()).isGreaterThan(partitionKeyRanges.size()) | ||
| .as("Partition ranges should increase after split"); | ||
| logger.info("After split num partitions = {}", partitionKeyRangesAfterSplit.size()); | ||
|
|
||
| responseFlux = container.processBulkOperations(cosmosItemOperationFlux2, bulkProcessingOptions); | ||
|
|
||
| AtomicInteger processedDoc2 = new AtomicInteger(0); | ||
| responseFlux | ||
| .flatMap((CosmosBulkOperationResponse<CosmosBulkAsyncTest> cosmosBulkOperationResponse) -> { | ||
|
|
||
| processedDoc2.incrementAndGet(); | ||
|
|
||
| CosmosBulkItemResponse cosmosBulkItemResponse = cosmosBulkOperationResponse.getResponse(); | ||
| assertThat(cosmosBulkItemResponse.getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); | ||
| assertThat(cosmosBulkItemResponse.getRequestCharge()).isGreaterThan(0); | ||
| assertThat(cosmosBulkItemResponse.getCosmosDiagnostics().toString()).isNotNull(); | ||
| assertThat(cosmosBulkItemResponse.getSessionToken()).isNotNull(); | ||
| assertThat(cosmosBulkItemResponse.getActivityId()).isNotNull(); | ||
| assertThat(cosmosBulkItemResponse.getRequestCharge()).isNotNull(); | ||
|
|
||
| return Mono.just(cosmosBulkItemResponse); | ||
| }).blockLast(); | ||
|
|
||
| assertThat(processedDoc.get()).isEqualTo(totalRequest); | ||
| container.delete().block(); | ||
| } | ||
|
|
||
| @NotNull | ||
| private List<PartitionKeyRange> getPartitionKeyRanges( | ||
| String containerId, CosmosAsyncClient asyncClient) { | ||
| List<PartitionKeyRange> partitionKeyRanges = new ArrayList<>(); | ||
| AsyncDocumentClient asyncDocumentClient = BridgeInternal.getContextClient(asyncClient); | ||
| List<FeedResponse<PartitionKeyRange>> partitionFeedResponseList = asyncDocumentClient | ||
| .readPartitionKeyRanges("/dbs/" + createdDatabase.getId() | ||
| + "/colls/" + containerId, | ||
| new CosmosQueryRequestOptions()) | ||
| .collectList().block(); | ||
| partitionFeedResponseList.forEach(f -> partitionKeyRanges.addAll(f.getResults())); | ||
| return partitionKeyRanges; | ||
| } | ||
|
|
||
| private int getTotalRequest() { | ||
| int countRequest = new Random().nextInt(100) + 200; | ||
| logger.info("Total count of request for this test case: " + countRequest); | ||
|
|
||
| return countRequest; | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.