Query prefetch fix

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)

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 <T> Flux<FeedResponse<T>> 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

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<ObjectNode> cosmosPagedIterable = cosmosContainer.queryItems("select * from c",
+                cosmosQueryRequestOptions, ObjectNode.class);
+
+        Iterable<FeedResponse<ObjectNode>> feedResponses = cosmosPagedIterable.iterableByPage(10);
+        //  Just creating iterator drains all the results!
+        Iterator<FeedResponse<ObjectNode>> iterator = feedResponses.iterator();
+        if (iterator.hasNext()) {
+            FeedResponse<ObjectNode> 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<ObjectNode> cosmosPagedFlux = cosmosAsyncContainer.queryItems("select * from c",
+                cosmosQueryRequestOptions, ObjectNode.class);
+
+        CosmosPagedIterable<ObjectNode> cosmosPagedIterable = new CosmosPagedIterable<>(cosmosPagedFlux, 10, 1);
+        Iterator<FeedResponse<ObjectNode>> iterator = cosmosPagedIterable.iterableByPage().iterator();
+        if (iterator.hasNext()) {
+            FeedResponse<ObjectNode> 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<FeedResponse<Long>> 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<Flux<FeedResponse<Long>>> fluxList1 = Arrays.asList(validatePrefetchControl(numPages2, 10, prefetchEager2));
+        Flux<FeedResponse<Long>> 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<FeedResponse<Long>> 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<Flux<FeedResponse<Long>>> fluxList2 = Arrays.asList(validatePrefetchControl(numPages4, 10, prefetchLazy2));
+        Flux<FeedResponse<Long>> 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<FeedResponse<Long>> flux, AtomicInteger pagesPrefetched, int batchSize) {
+        Iterator<FeedResponse<Long>> iterator = flux.toIterable(batchSize).iterator();
+        if (iterator.hasNext()) {
+            iterator.next();
+        }
+        return pagesPrefetched;
+    }
+
+    private Flux<FeedResponse<Long>> validatePrefetchControl(int numPages, int pageSize, AtomicInteger pagesFetched) {
+        return Flux.generate(Tuple::new, (Tuple state, SynchronousSink<FeedResponse<Long>> 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<Long> feedResponse;
+
+        Tuple() {
+            pageIdx = new AtomicInteger(0);
+            feedResponse = ModelBridgeInternal.createFeedResponse(new ArrayList<>(), new HashMap<>());
+        }
+    }
+
     @Test(groups = { "unit" })
     public void PagePrefetchCountReasonablyLow() {