From a877f164e49fb90649cf0c9aed506ade0598cfd5 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Thu, 10 Nov 2022 20:43:47 -0500 Subject: [PATCH 01/15] Hardcoded prefetch value in createQueryInternal --- .../com/azure/cosmos/implementation/RxDocumentClientImpl.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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..645ca9e13d9b 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 @@ -951,7 +951,7 @@ private Flux> createQueryInternal( } return tFeedResponse; }); - }); + }, 1, 1); } @Override From d0657c195b0e3717667f87d6eb62e004ae922391 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Fri, 11 Nov 2022 13:41:22 -0500 Subject: [PATCH 02/15] Added prefetch unit test --- .../azure/cosmos/CosmosPagedIterableTest.java | 64 +++++++++++++++++++ 1 file changed, 64 insertions(+) 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..dfaff94ba561 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,66 @@ public void readAllItemsBySubscribeWithCosmosPagedIterableHandler() throws Excep assertThat(handleCount.get() >= 1).isTrue(); } + @Test(groups = {"unit"}) + public void validatePrefetchControl() { + AtomicInteger prefetchEager1 = new AtomicInteger(); + Flux> eagerDrain1 = validatePrefetchControl(10, 10, prefetchEager1) + .flatMapSequential(Flux::just, 1, 1) + .flatMap(Flux::just); + assertThat(validate(eagerDrain1, prefetchEager1).get()).isEqualTo(10); + + AtomicInteger prefetchEager2 = new AtomicInteger(); + List>> fluxList1 = Arrays.asList(validatePrefetchControl(10, 10, prefetchEager2)); + Flux> fastDrain2 = Flux + .mergeSequential(fluxList1, 1, 1) + .flatMap(Flux::just); + assertThat(validate(fastDrain2, prefetchEager2).get()).isEqualTo(10); + + AtomicInteger prefetchLazy1 = new AtomicInteger(); + Flux> lazyDrain1 = validatePrefetchControl(10, 10, prefetchLazy1) + .flatMapSequential(Flux::just, 1, 1) + .flatMap(Flux::just, 1, 1); + assertThat(validate(lazyDrain1, prefetchLazy1).get()).isLessThan(10); + + AtomicInteger prefetchLazy2 = new AtomicInteger(); + List>> fluxList2 = Arrays.asList(validatePrefetchControl(10, 10, prefetchLazy2)); + Flux> lazyDrain2 = Flux + .mergeSequential(fluxList2, 1, 1) + .flatMap(Flux::just, 1, 1); + assertThat(validate(lazyDrain2, prefetchLazy2).get()).isLessThan(10); + } + + private AtomicInteger validate(Flux> flux, AtomicInteger pagesPrefetched) { + Boolean hasNext = flux.toIterable(1).iterator().hasNext(); + 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); + pagesFetched.addAndGet(1); + } else { + sink.complete(); + } + return state; + }); + } + 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() { From b6c867df96e0066612b462eb6664f181164d6619 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Fri, 11 Nov 2022 15:32:32 -0500 Subject: [PATCH 03/15] Modified prefetch unit test --- .../azure/cosmos/CosmosPagedIterableTest.java | 41 ++++++++++++++----- 1 file changed, 30 insertions(+), 11 deletions(-) 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 dfaff94ba561..f88dba069fdd 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 @@ -120,34 +120,53 @@ public void readAllItemsBySubscribeWithCosmosPagedIterableHandler() throws Excep @Test(groups = {"unit"}) public void validatePrefetchControl() { AtomicInteger prefetchEager1 = new AtomicInteger(); - Flux> eagerDrain1 = validatePrefetchControl(10, 10, prefetchEager1) + int bathSize1 = 1; + int numPages1 = 100; + Flux> eagerDrain1 = validatePrefetchControl(numPages1, 10, prefetchEager1) .flatMapSequential(Flux::just, 1, 1) .flatMap(Flux::just); - assertThat(validate(eagerDrain1, prefetchEager1).get()).isEqualTo(10); + // assert all pages are fetched eagerly + assertThat(validate(eagerDrain1, prefetchEager1, bathSize1).get()).isEqualTo(numPages1); + int bathSize2 = 1; + int numPages2 = 100; AtomicInteger prefetchEager2 = new AtomicInteger(); - List>> fluxList1 = Arrays.asList(validatePrefetchControl(10, 10, prefetchEager2)); + List>> fluxList1 = Arrays.asList(validatePrefetchControl(numPages2, 10, prefetchEager2)); Flux> fastDrain2 = Flux .mergeSequential(fluxList1, 1, 1) .flatMap(Flux::just); - assertThat(validate(fastDrain2, prefetchEager2).get()).isEqualTo(10); + // assert all pages are fetched eagerly + assertThat(validate(fastDrain2, prefetchEager2, bathSize2).get()).isEqualTo(numPages2); + int batchSize3 = 19; + int numPages3 = 100; AtomicInteger prefetchLazy1 = new AtomicInteger(); - Flux> lazyDrain1 = validatePrefetchControl(10, 10, prefetchLazy1) + Flux> lazyDrain1 = validatePrefetchControl(numPages3, 10, prefetchLazy1) .flatMapSequential(Flux::just, 1, 1) .flatMap(Flux::just, 1, 1); - assertThat(validate(lazyDrain1, prefetchLazy1).get()).isLessThan(10); + // 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(10, 10, prefetchLazy2)); + List>> fluxList2 = Arrays.asList(validatePrefetchControl(numPages4, 10, prefetchLazy2)); Flux> lazyDrain2 = Flux .mergeSequential(fluxList2, 1, 1) .flatMap(Flux::just, 1, 1); - assertThat(validate(lazyDrain2, prefetchLazy2).get()).isLessThan(10); + // 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) { - Boolean hasNext = flux.toIterable(1).iterator().hasNext(); + private AtomicInteger validate(Flux> flux, AtomicInteger pagesPrefetched, int batchSize) { + Iterator> iterator = flux.toIterable(batchSize).iterator(); + if (iterator.hasNext()) { + iterator.next(); + } return pagesPrefetched; } @@ -159,8 +178,8 @@ private Flux> validatePrefetchControl(int numPages, int pageS .collect(Collectors.toList()), new HashMap<>()); sink.next(state.feedResponse); - state.pageIdx.addAndGet(1); pagesFetched.addAndGet(1); + state.pageIdx.addAndGet(1); } else { sink.complete(); } From a8ffdc1069b48409f733bb499b061454edc6bd33 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Sat, 12 Nov 2022 17:57:02 -0500 Subject: [PATCH 04/15] Format --- .../src/test/java/com/azure/cosmos/CosmosPagedIterableTest.java | 1 + 1 file changed, 1 insertion(+) 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 f88dba069fdd..508673cbdf3b 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 @@ -186,6 +186,7 @@ private Flux> validatePrefetchControl(int numPages, int pageS return state; }); } + static class Tuple { AtomicInteger pageIdx; FeedResponse feedResponse; From 33f26eabdace09f06c33dd54b8e1a39c341bd3fc Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Sat, 12 Nov 2022 18:16:39 -0500 Subject: [PATCH 05/15] Updated CHANGELOG.md --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index a1b503d0fb08..f73824216822 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -10,6 +10,7 @@ * Fixed a rare race condition for `query plan` cache exceeding the allowed size limit - See [PR 31859](https://github.com/Azure/azure-sdk-for-java/pull/31859) * 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 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) From 51e4a3ecda4a85aac3c9a26d00c610568eef1adb Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Sun, 13 Nov 2022 12:15:33 -0500 Subject: [PATCH 06/15] Test for regression --- .../com/azure/cosmos/implementation/RxDocumentClientImpl.java | 2 +- .../test/java/com/azure/cosmos/CosmosPagedIterableTest.java | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) 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 645ca9e13d9b..b61cace5b299 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 @@ -951,7 +951,7 @@ private Flux> createQueryInternal( } return tFeedResponse; }); - }, 1, 1); + }/*, 1, 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 508673cbdf3b..4f171dbb6ab0 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 @@ -178,13 +178,12 @@ private Flux> validatePrefetchControl(int numPages, int pageS .collect(Collectors.toList()), new HashMap<>()); sink.next(state.feedResponse); - pagesFetched.addAndGet(1); state.pageIdx.addAndGet(1); } else { sink.complete(); } return state; - }); + }).doOnNext(response -> pagesFetched.addAndGet(1)); } static class Tuple { From 21b56f26476e0988fb44b3cb1af08c2e87666679 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Sun, 13 Nov 2022 15:12:31 -0500 Subject: [PATCH 07/15] Test for regression --- .../com/azure/cosmos/implementation/RxDocumentClientImpl.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 b61cace5b299..645ca9e13d9b 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 @@ -951,7 +951,7 @@ private Flux> createQueryInternal( } return tFeedResponse; }); - }/*, 1, 1*/); + }, 1, 1); } @Override From 2057d60034318dc000d258fe243ce2f0798d52cd Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Mon, 14 Nov 2022 15:08:55 -0500 Subject: [PATCH 08/15] Set query flatMap concurrency to default --- .../implementation/RxDocumentClientImpl.java | 5 ++- .../azure/cosmos/CosmosPagedIterableTest.java | 44 +++++++++++++++++-- 2 files changed, 45 insertions(+), 4 deletions(-) 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 645ca9e13d9b..ae8c40c33532 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,9 @@ private Flux> createQueryInternal( } return tFeedResponse; }); - }, 1, 1); + // concurrency is set to the default to + // not affect latency by minimizing concurrency and prefetch + }, 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 4f171dbb6ab0..849eb1c3166c 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 @@ -117,7 +117,45 @@ public void readAllItemsBySubscribeWithCosmosPagedIterableHandler() throws Excep assertThat(handleCount.get() >= 1).isTrue(); } - @Test(groups = {"unit"}) + @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); + } + + // TODO: Investigate scenario + @Test(groups = {"unit"}, enabled = false) public void validatePrefetchControl() { AtomicInteger prefetchEager1 = new AtomicInteger(); int bathSize1 = 1; @@ -143,7 +181,7 @@ public void validatePrefetchControl() { AtomicInteger prefetchLazy1 = new AtomicInteger(); Flux> lazyDrain1 = validatePrefetchControl(numPages3, 10, prefetchLazy1) .flatMapSequential(Flux::just, 1, 1) - .flatMap(Flux::just, 1, 1); + .flatMap(Flux::just, 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) @@ -155,7 +193,7 @@ public void validatePrefetchControl() { List>> fluxList2 = Arrays.asList(validatePrefetchControl(numPages4, 10, prefetchLazy2)); Flux> lazyDrain2 = Flux .mergeSequential(fluxList2, 1, 1) - .flatMap(Flux::just, 1, 1); + .flatMap(Flux::just, 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) From da175358121c9863a435b3dbe779c97e9bc5bbb6 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Mon, 14 Nov 2022 15:12:24 -0500 Subject: [PATCH 09/15] Set query flatMap concurrency to default --- .../test/java/com/azure/cosmos/CosmosPagedIterableTest.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 849eb1c3166c..95e26c6beac9 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 @@ -154,7 +154,8 @@ public void queryItemsWithCosmosPagedFlux() throws Exception { Thread.sleep(5 * 1000); } - // TODO: Investigate scenario + // TODO: Check if eager prefetch can be avoided + // TODO: when concurrency is set to 256 @Test(groups = {"unit"}, enabled = false) public void validatePrefetchControl() { AtomicInteger prefetchEager1 = new AtomicInteger(); From 35057d01a308e5189c7c87c3bfe302478096b46f Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Mon, 14 Nov 2022 15:16:18 -0500 Subject: [PATCH 10/15] Set query flatMap concurrency to default --- .../test/java/com/azure/cosmos/CosmosPagedIterableTest.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 95e26c6beac9..f252e4369314 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 @@ -155,7 +155,8 @@ public void queryItemsWithCosmosPagedFlux() throws Exception { } // TODO: Check if eager prefetch can be avoided - // TODO: when concurrency is set to 256 + // TODO: when concurrency is set to 256 on a downstream flatMap + // TODO: the unit tests can be enabled then @Test(groups = {"unit"}, enabled = false) public void validatePrefetchControl() { AtomicInteger prefetchEager1 = new AtomicInteger(); From 5e1e4d8debe104f23dc717e2b8811594b2d53b9b Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Mon, 14 Nov 2022 15:17:53 -0500 Subject: [PATCH 11/15] Set query flatMap concurrency to default --- .../com/azure/cosmos/implementation/RxDocumentClientImpl.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 ae8c40c33532..faf5769cbf82 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 @@ -953,7 +953,7 @@ private Flux> createQueryInternal( return tFeedResponse; }); // concurrency is set to the default to - // not affect latency by minimizing concurrency and prefetch + // not affect latency by minimizing prefetch }, Queues.SMALL_BUFFER_SIZE, 1); } From f9f4e9bf20ee6a1c97245353de261b291ca52d82 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Mon, 14 Nov 2022 19:02:54 -0500 Subject: [PATCH 12/15] Fixed unit test --- .../implementation/RxDocumentClientImpl.java | 9 ++--- .../azure/cosmos/CosmosPagedIterableTest.java | 36 +++++++++---------- 2 files changed, 21 insertions(+), 24 deletions(-) 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 faf5769cbf82..2524017192e0 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 @@ -910,7 +910,7 @@ private Flux> createQuery( return ObservableHelper.fluxInlineIfPossibleAsObs( () -> createQueryInternal( - resourceLink, sqlQuery, options, klass, resourceTypeEnum, queryClient, correlationActivityId), + resourceLink, sqlQuery, options, klass, resourceTypeEnum, queryClient, correlationActivityId).log("913"), invalidPartitionExceptionRetryPolicy); } @@ -937,7 +937,7 @@ private Flux> createQueryInternal( } QueryInfo finalQueryInfo = queryInfo; - return iDocumentQueryExecutionContext.executeAsync() + return iDocumentQueryExecutionContext.executeAsync().log("940") .map(tFeedResponse -> { if (finalQueryInfo != null) { if (finalQueryInfo.hasSelectValue()) { @@ -952,8 +952,9 @@ private Flux> createQueryInternal( } return tFeedResponse; }); - // concurrency is set to the default to - // not affect latency by minimizing prefetch + // concurrency is set to Queues.SMALL_BUFFER_SIZE to + // maximize the IDocumentQueryExecutionContext instances to subscribe to concurrently + // prefetch is set to 1 to minimize the no. of requested pages (result of merged executeAsync invocations) }, Queues.SMALL_BUFFER_SIZE, 1); } 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 f252e4369314..ce5222919fa3 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 @@ -154,36 +154,32 @@ public void queryItemsWithCosmosPagedFlux() throws Exception { Thread.sleep(5 * 1000); } - // TODO: Check if eager prefetch can be avoided - // TODO: when concurrency is set to 256 on a downstream flatMap - // TODO: the unit tests can be enabled then - @Test(groups = {"unit"}, enabled = false) + @Test(groups = {"unit"}) public void validatePrefetchControl() { AtomicInteger prefetchEager1 = new AtomicInteger(); int bathSize1 = 1; int numPages1 = 100; - Flux> eagerDrain1 = validatePrefetchControl(numPages1, 10, prefetchEager1) - .flatMapSequential(Flux::just, 1, 1) - .flatMap(Flux::just); - // assert all pages are fetched eagerly - assertThat(validate(eagerDrain1, prefetchEager1, bathSize1).get()).isEqualTo(numPages1); + 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 - .mergeSequential(fluxList1, 1, 1) - .flatMap(Flux::just); - // assert all pages are fetched eagerly - assertThat(validate(fastDrain2, prefetchEager2, bathSize2).get()).isEqualTo(numPages2); + 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 = validatePrefetchControl(numPages3, 10, prefetchLazy1) - .flatMapSequential(Flux::just, 1, 1) - .flatMap(Flux::just, Queues.SMALL_BUFFER_SIZE, 1); + 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) @@ -193,9 +189,9 @@ public void validatePrefetchControl() { int numPages4 = 100; AtomicInteger prefetchLazy2 = new AtomicInteger(); List>> fluxList2 = Arrays.asList(validatePrefetchControl(numPages4, 10, prefetchLazy2)); - Flux> lazyDrain2 = Flux - .mergeSequential(fluxList2, 1, 1) - .flatMap(Flux::just, Queues.SMALL_BUFFER_SIZE, 1); + 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) From 2cabc249968a9df9b153e5106bc2bcf19bafa922 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Mon, 14 Nov 2022 19:06:11 -0500 Subject: [PATCH 13/15] Format --- .../com/azure/cosmos/implementation/RxDocumentClientImpl.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 2524017192e0..d5da1e92cba0 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 @@ -910,7 +910,7 @@ private Flux> createQuery( return ObservableHelper.fluxInlineIfPossibleAsObs( () -> createQueryInternal( - resourceLink, sqlQuery, options, klass, resourceTypeEnum, queryClient, correlationActivityId).log("913"), + resourceLink, sqlQuery, options, klass, resourceTypeEnum, queryClient, correlationActivityId), invalidPartitionExceptionRetryPolicy); } @@ -937,7 +937,7 @@ private Flux> createQueryInternal( } QueryInfo finalQueryInfo = queryInfo; - return iDocumentQueryExecutionContext.executeAsync().log("940") + return iDocumentQueryExecutionContext.executeAsync() .map(tFeedResponse -> { if (finalQueryInfo != null) { if (finalQueryInfo.hasSelectValue()) { From f21d995a9de138f2f74c5a5a1694dba0ab225d3b Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Tue, 15 Nov 2022 00:05:45 -0500 Subject: [PATCH 14/15] Updated comments --- .../com/azure/cosmos/implementation/RxDocumentClientImpl.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 d5da1e92cba0..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 @@ -953,8 +953,8 @@ private Flux> createQueryInternal( return tFeedResponse; }); // concurrency is set to Queues.SMALL_BUFFER_SIZE to - // maximize the IDocumentQueryExecutionContext instances to subscribe to concurrently - // prefetch is set to 1 to minimize the no. of requested pages (result of merged executeAsync invocations) + // 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); } From f1cfe5677df37687363143b572b3b44f068a516e Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Tue, 15 Nov 2022 19:55:37 -0500 Subject: [PATCH 15/15] Addressed review comments --- .../src/test/java/com/azure/cosmos/CosmosPagedIterableTest.java | 1 - 1 file changed, 1 deletion(-) 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 ce5222919fa3..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 @@ -135,7 +135,6 @@ public void queryItemsWithCosmosPagedIterable() throws Exception { Thread.sleep(5 * 1000); } - @Test(groups = { "simple" }, timeOut = TIMEOUT, enabled = false) public void queryItemsWithCosmosPagedFlux() throws Exception {