diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 10efb0c3d469..9e04b94b43f8 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -11,6 +11,7 @@ * Added improvement in `RntbdClientChannelHealthChecker` for detecting continuous transit timeout. - See [PR 31544](https://github.com/Azure/azure-sdk-for-java/pull/31544) * Fixed an issue in replica validation where addresses may have not sorted properly when replica validation is enabled. - See [PR 32022](https://github.com/Azure/azure-sdk-for-java/pull/32022) * Fixed unicode char handling in Uris in Cosmos Http Client. - See [PR 32058](https://github.com/Azure/azure-sdk-for-java/pull/32058) +* Fixed an eager prefetch issue to lazily prefetch pages on a query - See [PR 32122](https://github.com/Azure/azure-sdk-for-java/pull/32122) #### Other Changes * Shaded `MurmurHash3` of apache `commons-codec` to enable removing of the `guava` dependency - CVE-2020-8908 - See [PR 31761](https://github.com/Azure/azure-sdk-for-java/pull/31761) 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 dbf63c87a5b8..360f7c816890 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 @@ -72,6 +72,7 @@ import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.util.concurrent.Queues; import java.io.IOException; import java.io.UnsupportedEncodingException; @@ -951,7 +952,10 @@ private Flux> createQueryInternal( } return tFeedResponse; }); - }); + // concurrency is set to Queues.SMALL_BUFFER_SIZE to + // maximize the IDocumentQueryExecutionContext publisher instances to subscribe to concurrently + // prefetch is set to 1 to minimize the no. prefetched pages (result of merged executeAsync invocations) + }, Queues.SMALL_BUFFER_SIZE, 1); } @Override diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosPagedIterableTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosPagedIterableTest.java index 235b7a56e7a7..1c81c3f5a4a8 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosPagedIterableTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosPagedIterableTest.java @@ -24,10 +24,14 @@ import org.testng.annotations.Factory; import org.testng.annotations.Test; import reactor.core.publisher.Flux; +import reactor.core.publisher.SynchronousSink; import reactor.util.concurrent.Queues; +import java.util.ArrayList; +import java.util.Arrays; import java.util.HashMap; import java.util.Iterator; +import java.util.List; import java.util.UUID; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; @@ -113,6 +117,120 @@ public void readAllItemsBySubscribeWithCosmosPagedIterableHandler() throws Excep assertThat(handleCount.get() >= 1).isTrue(); } + @Test(groups = { "simple" }, timeOut = TIMEOUT, enabled = false) + public void queryItemsWithCosmosPagedIterable() throws Exception { + + CosmosQueryRequestOptions cosmosQueryRequestOptions = new CosmosQueryRequestOptions(); + cosmosQueryRequestOptions.setMaxBufferedItemCount(10); + CosmosPagedIterable cosmosPagedIterable = cosmosContainer.queryItems("select * from c", + cosmosQueryRequestOptions, ObjectNode.class); + + Iterable> feedResponses = cosmosPagedIterable.iterableByPage(10); + // Just creating iterator drains all the results! + Iterator> iterator = feedResponses.iterator(); + if (iterator.hasNext()) { + FeedResponse next = iterator.next(); + logger.info("Next is : {}", next.getResults().size()); + } + Thread.sleep(5 * 1000); + } + + @Test(groups = { "simple" }, timeOut = TIMEOUT, enabled = false) + public void queryItemsWithCosmosPagedFlux() throws Exception { + + CosmosQueryRequestOptions cosmosQueryRequestOptions = new CosmosQueryRequestOptions(); + cosmosQueryRequestOptions.setMaxBufferedItemCount(10); + CosmosAsyncContainer cosmosAsyncContainer = CosmosBridgeInternal.getCosmosAsyncContainer(cosmosContainer); + CosmosPagedFlux cosmosPagedFlux = cosmosAsyncContainer.queryItems("select * from c", + cosmosQueryRequestOptions, ObjectNode.class); + + CosmosPagedIterable cosmosPagedIterable = new CosmosPagedIterable<>(cosmosPagedFlux, 10, 1); + Iterator> iterator = cosmosPagedIterable.iterableByPage().iterator(); + if (iterator.hasNext()) { + FeedResponse next = iterator.next(); + logger.info("Next is : {}", next.getResults().size()); + } + Thread.sleep(5 * 1000); + } + + @Test(groups = {"unit"}) + public void validatePrefetchControl() { + AtomicInteger prefetchEager1 = new AtomicInteger(); + int bathSize1 = 1; + int numPages1 = 100; + Flux> eagerDrain1 = Flux.fromIterable(Arrays.asList(1)) + .flatMap(x -> validatePrefetchControl(numPages1, 10, prefetchEager1) + .flatMapSequential(Flux::just, 1, 1)); + // assert 32 to 37 pages are fetched eagerly even though batchSize is set to 1 + assertThat(validate(eagerDrain1, prefetchEager1, bathSize1).get()).isBetween(32, 37); + + int bathSize2 = 1; + int numPages2 = 100; + AtomicInteger prefetchEager2 = new AtomicInteger(); + List>> fluxList1 = Arrays.asList(validatePrefetchControl(numPages2, 10, prefetchEager2)); + Flux> fastDrain2 = Flux.fromIterable(Arrays.asList(1)) + .flatMap(x -> Flux.mergeSequential(fluxList1, 1, 1)); + // assert 32 to 37 pages are fetched eagerly even though batchSize is set to 1 + assertThat(validate(fastDrain2, prefetchEager2, bathSize2).get()).isBetween(32, 37); + + int batchSize3 = 19; + int numPages3 = 100; + AtomicInteger prefetchLazy1 = new AtomicInteger(); + Flux> lazyDrain1 = Flux.fromIterable(Arrays.asList(1)) + .flatMap(x -> validatePrefetchControl(numPages3, 10, prefetchLazy1) + .flatMapSequential(Flux::just, 1, 1), Queues.SMALL_BUFFER_SIZE, 1); + // assert that no. of pages fetched is close to the batch size + assertThat(validate(lazyDrain1, prefetchLazy1, batchSize3).get()) + .isLessThan(4 + batchSize3) + .isGreaterThanOrEqualTo(batchSize3); + + int batchSize4 = 37; + int numPages4 = 100; + AtomicInteger prefetchLazy2 = new AtomicInteger(); + List>> fluxList2 = Arrays.asList(validatePrefetchControl(numPages4, 10, prefetchLazy2)); + Flux> lazyDrain2 = Flux.just(Arrays.asList(1)) + .flatMap(x -> Flux + .mergeSequential(fluxList2, 1, 1), Queues.SMALL_BUFFER_SIZE, 1); + // assert that no. of pages fetched is close to the batch size + assertThat(validate(lazyDrain2, prefetchLazy2, batchSize4).get()) + .isLessThan(4 + batchSize4) + .isGreaterThanOrEqualTo(batchSize4); + } + + private AtomicInteger validate(Flux> flux, AtomicInteger pagesPrefetched, int batchSize) { + Iterator> iterator = flux.toIterable(batchSize).iterator(); + if (iterator.hasNext()) { + iterator.next(); + } + return pagesPrefetched; + } + + private Flux> validatePrefetchControl(int numPages, int pageSize, AtomicInteger pagesFetched) { + return Flux.generate(Tuple::new, (Tuple state, SynchronousSink> sink) -> { + if (state.pageIdx.get() < numPages) { + state.feedResponse = ModelBridgeInternal.createFeedResponse(LongStream.range(state.pageIdx.get(), state.pageIdx.get() + pageSize) + .boxed() + .collect(Collectors.toList()), + new HashMap<>()); + sink.next(state.feedResponse); + state.pageIdx.addAndGet(1); + } else { + sink.complete(); + } + return state; + }).doOnNext(response -> pagesFetched.addAndGet(1)); + } + + static class Tuple { + AtomicInteger pageIdx; + FeedResponse feedResponse; + + Tuple() { + pageIdx = new AtomicInteger(0); + feedResponse = ModelBridgeInternal.createFeedResponse(new ArrayList<>(), new HashMap<>()); + } + } + @Test(groups = { "unit" }) public void PagePrefetchCountReasonablyLow() {