Fix DbKvs.getRangeOfTimestamps() only returning the first page

The expectation is to return all results that fall into the range
regardless of the page size. This is how all other KVSs work
at least, and also consistent with getRange().

I'm planning to delete all of this code soon, and just need it to
work in order to execute the intermediate steps of my big sweep
refactor/rewrite.

atlasdb-dbkvs/src/main/java/com/palantir/atlasdb/keyvalue/dbkvs/impl/DbKvs.java

@@ -702,7 +702,7 @@ private TimestampsByCellResultWithToken getTimestampsByCell(
             Token token) {
         try (ClosableIterator<AgnosticLightResultRow> iterator = table
                 .getAllRows(rows, columns, timestamp, false, DbReadTable.Order.fromBoolean(isReverse))) {
-            return TimestampsByCellResultWithToken.create(iterator, token, batchSize);
+            return TimestampsByCellResultWithToken.create(iterator, token, batchSize, isReverse);
         }
     }
 

atlasdb-dbkvs/src/main/java/com/palantir/atlasdb/keyvalue/dbkvs/impl/TimestampsByCellResultWithToken.java

@@ -23,6 +23,7 @@
 import com.google.common.collect.SetMultimap;
 import com.google.common.primitives.UnsignedBytes;
 import com.palantir.atlasdb.keyvalue.api.Cell;
+import com.palantir.atlasdb.keyvalue.api.RangeRequests;
 import com.palantir.common.base.ClosableIterator;
 import com.palantir.nexus.db.sql.AgnosticLightResultRow;
 
@@ -35,17 +36,20 @@ final class TimestampsByCellResultWithToken {
     private SetMultimap<Cell, Long> rowBuffer;
     private boolean moreResults = false;
     private Token token = Token.INITIAL;
+    private final boolean reverse;
 
-    private TimestampsByCellResultWithToken(ClosableIterator<AgnosticLightResultRow> iterator) {
+    private TimestampsByCellResultWithToken(ClosableIterator<AgnosticLightResultRow> iterator, boolean reverse) {
         entries = HashMultimap.create();
         rowBuffer = HashMultimap.create();
         this.iterator = Iterators.peekingIterator(iterator);
+        this.reverse = reverse;
     }
 
     static TimestampsByCellResultWithToken create(ClosableIterator<AgnosticLightResultRow> iterator,
             Token oldToken,
-            long batchSize) {
-        return new TimestampsByCellResultWithToken(iterator)
+            long batchSize,
+            boolean reverse) {
+        return new TimestampsByCellResultWithToken(iterator, reverse)
                 .moveForward(oldToken)
                 .getBatchOfTimestamps(batchSize)
                 .checkNextEntryAndCreateToken();
@@ -123,6 +127,13 @@ private TimestampsByCellResultWithToken checkNextEntryAndCreateToken() {
             }
         } else {
             flushRowBuffer();
+            if (currentRow != null) {
+                byte[] nextRow = RangeRequests.getNextStartRowUnlessTerminal(reverse, currentRow);
+                if (nextRow != null) {
+                    moreResults = true;
+                    token = ImmutableToken.builder().row(nextRow).shouldSkip(false).build();
+                }
+            }
         }
         return this;
     }

atlasdb-tests-shared/src/main/java/com/palantir/atlasdb/keyvalue/impl/AbstractKeyValueServiceTest.java

@@ -1078,6 +1078,23 @@ private void testGetRangeWithTimestamps(boolean reverse) {
         assertTrue(cell0.getValue().contains(TEST_TIMESTAMP + 1));
     }
 
+    @Test
+    public void testGetRangeOfTimestampsReturnsAllRows() {
+        keyValueService.put(TEST_TABLE,
+            ImmutableMap.of(
+                Cell.create(row0, column0), value0_t0,
+                Cell.create(row1, column0), value0_t0,
+                Cell.create(row2, column0), value0_t0),
+            TEST_TIMESTAMP);
+        RangeRequest range = RangeRequest.all().withBatchHint(1);
+        List<RowResult<Set<Long>>> results = ImmutableList.copyOf(
+                keyValueService.getRangeOfTimestamps(TEST_TABLE, range, TEST_TIMESTAMP + 1));
+        assertEquals(3, results.size());
+        assertTrue(Arrays.equals(row0, results.get(0).getRowName()));
+        assertTrue(Arrays.equals(row1, results.get(1).getRowName()));
+        assertTrue(Arrays.equals(row2, results.get(2).getRowName()));
+    }
+
     @Test
     public void testKeyAlreadyExists() {
         // Test that it does not throw some random exceptions

docs/source/release_notes/release-notes.rst

@@ -71,6 +71,10 @@ develop
          - Relax the signature of KeyValueService.addGarbageCollectionSentinelValues() to take an Iterable instead of a Set.
            (`Pull Request <https://github.com/palantir/atlasdb/pull/1843>`__)
 
+    *    - |fixed|
+         - Fixed DbKvs.getRangeOfTimestamps() only returning the first page of results rather than paging through the whole range.
+           (`Pull Request <https://github.com/palantir/atlasdb/pull/1872>`__)
+
 .. <<<<------------------------------------------------------------------------------------------------------------->>>>