Repository: foundationdb/fdb-record-layer
Issue: The FlatMapPipelinedCursor can return null continuations when not done

There are situations where the `FlatMapPipelinedCursor` can return `null` continuations even if the code isn't done yet. In particular, this can happen in situations where the inner cursor hits a limit before a single result is returned. It also probably means the cursor is less aggressive about updating its continuations than it should be, which might cause performance problems.

Changed files:
fdb-record-layer-core/test/unit/com/apple/foundationdb/record/RecordCursorTest.java
fdb-record-layer-core/src/com/apple/foundationdb/record/cursors/FlatMapPipelinedCursor.java

diff --git a/fdb-record-layer-core/src/com/apple/foundationdb/record/cursors/FlatMapPipelinedCursor.java b/fdb-record-layer-core/src/com/apple/foundationdb/record/cursors/FlatMapPipelinedCursor.java
index e7346488..dc25a85b 100644
--- a/fdb-record-layer-core/src/com/apple/foundationdb/record/cursors/FlatMapPipelinedCursor.java
+++ b/fdb-record-layer-core/src/com/apple/foundationdb/record/cursors/FlatMapPipelinedCursor.java
@@ -100,7 +100,7 @@ public class FlatMapPipelinedCursor<T, V> implements RecordCursor<V> {
         if (nextFuture == null) {
             mayGetContinuation = false;
             nextFuture = AsyncUtil.whileTrue(this::tryToFillPipeline, getExecutor()).thenApply(vignore -> {
-                boolean result = !(pipeline.isEmpty() || (innerReason != null && innerReason.isOutOfBand()));
+                boolean result = !(pipeline.isEmpty() || (innerReason != null && !innerReason.isSourceExhausted()));
                 mayGetContinuation = !result;
                 return result;
             });
@@ -168,7 +168,7 @@ public class FlatMapPipelinedCursor<T, V> implements RecordCursor<V> {
 
     @Override
     public NoNextReason getNoNextReason() {
-        if (innerReason != null && innerReason.isOutOfBand()) {
+        if (innerReason != null && !innerReason.isSourceExhausted()) {
             return innerReason;
         }
         return outerCursor.getNoNextReason();
@@ -239,14 +239,22 @@ public class FlatMapPipelinedCursor<T, V> implements RecordCursor<V> {
                 // Stop because next cursor has next.
                 return AsyncUtil.READY_FALSE;
             }
+            // Update the last entry to this entry. This will implicitly advance the continuation. This is safe
+            // because we have already hit the limit of this inner cursor. If the cursor is not exhausted, the
+            // whole cursor will stop and the advanced continuation will let it resume from where the cursor
+            // hit the limit. If the cursor is exhausted, then every item from this cursor has been returned to
+            // the user. That means that if the cursor is stopped right now, it will be resumed from the first
+            // item of the next element of the outer cursor, which is fine.
+            lastEntry = nextEntry;
             pipeline.remove();
             innerReason = nextEntry.innerCursor.getNoNextReason();
-            if (innerReason.isOutOfBand()) {
+            if (!innerReason.isSourceExhausted()) {
                 // Stop because inner cursor hit some limit.
                 // Not valid to take from later in the pipeline, even if available already.
                 return AsyncUtil.READY_FALSE;
+            } else {
+                return AsyncUtil.READY_TRUE;
             }
-            return AsyncUtil.READY_TRUE;
         }
         if (waitOuterNext == null) {
             // Recheck when cursor ready (pipeline is full).
diff --git a/fdb-record-layer-core/test/unit/com/apple/foundationdb/record/RecordCursorTest.java b/fdb-record-layer-core/test/unit/com/apple/foundationdb/record/RecordCursorTest.java
index 5e2b1872..cad549f3 100644
--- a/fdb-record-layer-core/test/unit/com/apple/foundationdb/record/RecordCursorTest.java
+++ b/fdb-record-layer-core/test/unit/com/apple/foundationdb/record/RecordCursorTest.java
@@ -27,11 +27,16 @@ import com.apple.foundationdb.record.cursors.LazyCursor;
 import com.apple.foundationdb.record.cursors.MapCursor;
 import com.apple.foundationdb.record.cursors.RowLimitedCursor;
 import com.apple.foundationdb.record.cursors.SkipCursor;
+import com.apple.foundationdb.record.provider.foundationdb.FDBStoreTimer;
 import com.google.common.collect.Iterators;
 import com.google.common.collect.Lists;
+import com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.commons.lang3.tuple.Pair;
 import org.hamcrest.Matchers;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
 
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -59,6 +64,9 @@ import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.greaterThan;
+import static org.hamcrest.Matchers.greaterThanOrEqualTo;
+import static org.hamcrest.Matchers.isOneOf;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
@@ -362,6 +370,62 @@ public class RecordCursorTest {
         assertEquals(adjusted, pieces);
     }
 
+    @EnumSource(TestHelpers.BooleanEnum.class)
+    @ParameterizedTest(name = "pipelineWithInnerLimits [outOfBand = {0}]")
+    public void pipelineWithInnerLimits(TestHelpers.BooleanEnum outOfBandEnum) throws InvalidProtocolBufferException {
+        final boolean outOfBand = outOfBandEnum.toBoolean();
+        final RecordCursor.NoNextReason[] possibleNoNextReasons = new RecordCursor.NoNextReason[]{
+                RecordCursor.NoNextReason.SOURCE_EXHAUSTED,
+                outOfBand ? RecordCursor.NoNextReason.TIME_LIMIT_REACHED : RecordCursor.NoNextReason.RETURN_LIMIT_REACHED
+        };
+        final List<Integer> ints = IntStream.range(0, 10).boxed().collect(Collectors.toList());
+
+        final FDBStoreTimer timer = new FDBStoreTimer();
+        final BiFunction<Integer, byte[], RecordCursor<Pair<Integer, Integer>>> innerFunc = (x, continuation) -> {
+            final RecordCursor<Integer> intCursor = RecordCursor.fromList(ints, continuation);
+            final RecordCursor<Integer> limitedCursor;
+            if (outOfBand) {
+                limitedCursor = new FakeOutOfBandCursor<>(intCursor, 3);
+            } else {
+                limitedCursor = intCursor.limitRowsTo(3);
+            }
+            return limitedCursor
+                    .filterInstrumented(y -> y < x, timer, FDBStoreTimer.Counts.QUERY_FILTER_GIVEN, FDBStoreTimer.Events.QUERY_FILTER, FDBStoreTimer.Counts.QUERY_FILTER_PASSED, FDBStoreTimer.Counts.QUERY_DISCARDED)
+                    .map(y -> Pair.of(x, y));
+        };
+        final Function<byte[], RecordCursor<Integer>> outerFunc = continuation -> RecordCursor.fromList(ints, continuation);
+
+        byte[] continuation = null;
+        int results = 0;
+        int leftSoFar = -1;
+        int rightSoFar = -1;
+        do {
+            RecordCursor<Pair<Integer, Integer>> cursor = RecordCursor.flatMapPipelined(outerFunc, innerFunc, continuation, 5);
+            while (cursor.hasNext()) {
+                Pair<Integer, Integer> value = cursor.next();
+                assertNotNull(value);
+                assertThat(value.getLeft(), greaterThan(value.getRight()));
+                assertThat(value.getLeft(), greaterThanOrEqualTo(leftSoFar));
+                if (value.getLeft() == leftSoFar) {
+                    assertThat(value.getRight(), greaterThan(rightSoFar));
+                    rightSoFar = value.getRight();
+                } else {
+                    leftSoFar = value.getLeft();
+                    rightSoFar = value.getRight();
+                }
+                results++;
+            }
+            assertThat(cursor.getNoNextReason(), isOneOf(possibleNoNextReasons));
+            continuation = cursor.getContinuation();
+        } while (continuation != null);
+
+        int expectedResults = ints.size() * (ints.size() - 1) / 2;
+        assertEquals(expectedResults, results);
+        assertEquals(ints.size() * ints.size(), timer.getCount(FDBStoreTimer.Counts.QUERY_FILTER_GIVEN));
+        assertEquals(expectedResults, timer.getCount(FDBStoreTimer.Counts.QUERY_FILTER_PASSED));
+        assertEquals(ints.size() * ints.size() - expectedResults, timer.getCount(FDBStoreTimer.Counts.QUERY_DISCARDED));
+    }
+
     @Test
     public void lazyCursorTest() {
         LazyCursor<Integer> cursor = new LazyCursor<>(