Allow slicing the bucket space for visitors

5 files changed, 197 insertions, 13 deletions

diff --git a/documentapi/abi-spec.json b/documentapi/abi-spec.json
index 9cc4f60ed7e..78a58f24a65 100644
--- a/documentapi/abi-spec.json
+++ b/documentapi/abi-spec.json
@@ -799,7 +799,7 @@
       "public"
     ],
     "methods": [
-      "public void <init>(int, com.yahoo.documentapi.ProgressToken)",
+      "public void <init>(int, com.yahoo.documentapi.ProgressToken, int, int)",
       "protected boolean isLosslessResetPossible()",
       "public boolean hasNext()",
       "public boolean shouldYield()",
@@ -851,6 +851,7 @@
       "public void setDistributionBitCount(int)",
       "public boolean visitsAllBuckets()",
       "public static com.yahoo.documentapi.VisitorIterator createFromDocumentSelection(java.lang.String, com.yahoo.document.BucketIdFactory, int, com.yahoo.documentapi.ProgressToken)",
+      "public static com.yahoo.documentapi.VisitorIterator createFromDocumentSelection(java.lang.String, com.yahoo.document.BucketIdFactory, int, com.yahoo.documentapi.ProgressToken, int, int)",
       "public static com.yahoo.documentapi.VisitorIterator createFromExplicitBucketSet(java.util.Set, int, com.yahoo.documentapi.ProgressToken)"
     ],
     "fields": []
@@ -931,6 +932,9 @@
       "public com.yahoo.documentapi.messagebus.loadtypes.LoadType getLoadType()",
       "public boolean skipBucketsOnFatalErrors()",
       "public void skipBucketsOnFatalErrors(boolean)",
+      "public void slice(int, int)",
+      "public int getSlices()",
+      "public int getSliceId()",
       "public void setDynamicallyIncreaseMaxBucketsPerVisitor(boolean)",
       "public void setDynamicMaxBucketsIncreaseFactor(float)",
       "public java.lang.String toString()"
diff --git a/documentapi/src/main/java/com/yahoo/documentapi/VisitorIterator.java b/documentapi/src/main/java/com/yahoo/documentapi/VisitorIterator.java
index e11bdf7f18c..0ccc64bb8f3 100755
--- a/documentapi/src/main/java/com/yahoo/documentapi/VisitorIterator.java
+++ b/documentapi/src/main/java/com/yahoo/documentapi/VisitorIterator.java
@@ -81,12 +81,24 @@ public class VisitorIterator {
     protected static class DistributionRangeBucketSource implements BucketSource {
         private boolean flushActive = false;
         private int distributionBitCount;
+        private final int slices;
+        private final int sliceId;
         // Wouldn't need this if this were a non-static class, but do it for
         // the sake of keeping things identical in Java and C++
         private ProgressToken progressToken;
 
         public DistributionRangeBucketSource(int distributionBitCount,
-                                             ProgressToken progress) {
+                                             ProgressToken progress,
+                                             int slices, int sliceId) {
+            if (slices < 1) {
+                throw new IllegalArgumentException("slices must be positive, but was " + slices);
+            }
+            if (sliceId < 0 || sliceId >= slices) {
+                throw new IllegalArgumentException("sliceId must be in [0, " + slices + "), but was " + sliceId);
+            }
+
+            this.slices = slices;
+            this.sliceId = sliceId;
             progressToken = progress;
 
             // New progress token (could also be empty, in which this is a
@@ -148,6 +160,7 @@ public class VisitorIterator {
             }
             // Should be all fixed up and good to go
             progressToken.setInconsistentState(false);
+            skipToSlice();
         }
 
         protected boolean isLosslessResetPossible() {
@@ -203,6 +216,7 @@ public class VisitorIterator {
                 assert(p.getActiveBucketCount() == 0);
                 p.clearAllBuckets();
                 p.setBucketCursor(0);
+                skipToSlice();
                 return;
             }
 
@@ -292,7 +306,14 @@ public class VisitorIterator {
         }
 
         public boolean hasNext() {
-            return progressToken.getBucketCursor() < (1L << distributionBitCount);
+            // There is a next bucket iff. there is a bucket no earlier than the cursor which
+            // is contained in the bucket space, and is also 0 modulo our sliceId; or if we're
+            // not yet properly initialised, with a real distribution bit count, we ignore this.
+            long nextBucket = progressToken.getBucketCursor();
+            if (distributionBitCount != 1) {
+                nextBucket += Math.floorMod(sliceId - nextBucket, slices);
+            }
+            return nextBucket < (1L << distributionBitCount);
         }
 
         public boolean shouldYield() {
@@ -311,13 +332,36 @@ public class VisitorIterator {
 
         public BucketProgress getNext() {
             assert(hasNext()) : "getNext() called with hasNext() == false";
-            long currentPosition = progressToken.getBucketCursor();
-            long key = ProgressToken.makeNthBucketKey(currentPosition, distributionBitCount);
-            ++currentPosition;
-            progressToken.setBucketCursor(currentPosition);
-            return new BucketProgress(
-                    new BucketId(ProgressToken.keyToBucketId(key)),
-                    new BucketId());
+
+            // Create the progress to return for creating visitors, and advance bucket cursor.
+            BucketProgress progress = new BucketProgress(toBucketId(progressToken.getBucketCursor(), distributionBitCount),
+                                                         new BucketId());
+            progressToken.setBucketCursor(progressToken.getBucketCursor() + 1);
+
+            // Skip ahead to our next next slice, to ensure we also exhaust the bucket space when
+            // hasNext() turns false, but there are still super buckets left after the current.
+            skipToSlice();
+
+            return progress;
+        }
+
+        // Advances the wrapped progress token's bucket cursor to our next slice, marking any skipped
+        // buckets as complete, but only if we've been initialised with a proper distribution bit count.
+        private void skipToSlice() {
+            if (distributionBitCount == 1)
+                return;
+
+            while (progressToken.getBucketCursor() < (1L << distributionBitCount) && (progressToken.getBucketCursor() - sliceId) % slices != 0) {
+                BucketId bucketId = toBucketId(progressToken.getBucketCursor(), distributionBitCount);
+                progressToken.addBucket(bucketId, ProgressToken.NULL_BUCKET, ProgressToken.BucketState.BUCKET_ACTIVE);
+                progressToken.updateProgress(toBucketId(progressToken.getBucketCursor(), distributionBitCount),
+                                             ProgressToken.FINISHED_BUCKET);
+                progressToken.setBucketCursor(progressToken.getBucketCursor() + 1);
+            }
+        }
+
+        private static BucketId toBucketId(long bucketCursor, int distributionBitCount) {
+            return new BucketId(ProgressToken.keyToBucketId(ProgressToken.makeNthBucketKey(bucketCursor, distributionBitCount)));
         }
 
         public int getDistributionBitCount() {
@@ -732,6 +776,13 @@ public class VisitorIterator {
         return bucketSource.visitsAllBuckets();
     }
 
+    public static VisitorIterator createFromDocumentSelection(
+            String documentSelection,
+            BucketIdFactory idFactory,
+            int distributionBitCount,
+            ProgressToken progress) throws ParseException {
+        return createFromDocumentSelection(documentSelection, idFactory, distributionBitCount, progress, 1, 0);
+    }
     /**
      * Create a new <code>VisitorIterator</code> instance based on the given document
      * selection string.
@@ -753,7 +804,9 @@ public class VisitorIterator {
             String documentSelection,
             BucketIdFactory idFactory,
             int distributionBitCount,
-            ProgressToken progress) throws ParseException {
+            ProgressToken progress,
+            int slices,
+            int sliceId) throws ParseException {
         BucketSelector bucketSel = new BucketSelector(idFactory);
         Set<BucketId> rawBuckets = bucketSel.getBucketList(documentSelection);
         BucketSource src;
@@ -763,7 +816,7 @@ public class VisitorIterator {
         // bit-based range source
         if (rawBuckets == null) {
             // Range source
-            src = new DistributionRangeBucketSource(distributionBitCount, progress);
+            src = new DistributionRangeBucketSource(distributionBitCount, progress, slices, sliceId);
         } else {
             // Explicit source
             src = new ExplicitBucketSource(rawBuckets, distributionBitCount, progress);
diff --git a/documentapi/src/main/java/com/yahoo/documentapi/VisitorParameters.java b/documentapi/src/main/java/com/yahoo/documentapi/VisitorParameters.java
index 8b0c8538855..44675d8d2ac 100644
--- a/documentapi/src/main/java/com/yahoo/documentapi/VisitorParameters.java
+++ b/documentapi/src/main/java/com/yahoo/documentapi/VisitorParameters.java
@@ -50,6 +50,8 @@ public class VisitorParameters extends Parameters {
     private int traceLevel = 0;
     private ThrottlePolicy throttlePolicy = null;
     private boolean skipBucketsOnFatalErrors = false;
+    private int slices = 1;
+    private int sliceId = 0;
 
     // Advanced parameter, only for internal use.
     Set<BucketId> bucketsToVisit = null;
@@ -101,6 +103,7 @@ public class VisitorParameters extends Parameters {
                 params.getDynamicMaxBucketsIncreaseFactor());
         setTraceLevel(params.getTraceLevel());
         skipBucketsOnFatalErrors(params.skipBucketsOnFatalErrors());
+        slice(params.getSlices(), getSliceId());
     }
 
     // Get functions
@@ -331,6 +334,15 @@ public class VisitorParameters extends Parameters {
 
     public void skipBucketsOnFatalErrors(boolean skipBucketsOnFatalErrors) { this.skipBucketsOnFatalErrors = skipBucketsOnFatalErrors; }
 
+    public void slice(int slices, int sliceId) {
+        this.slices = slices;
+        this.sliceId = sliceId;
+    }
+
+    public int getSlices() { return slices; }
+
+    public int getSliceId() { return sliceId; }
+
     /**
      * Set whether or not max buckets per visitor value should be dynamically
      * increased when using orderdoc and visitors do not return at least half
diff --git a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/MessageBusVisitorSession.java b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/MessageBusVisitorSession.java
index 5d07d433f18..f7242695490 100755
--- a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/MessageBusVisitorSession.java
+++ b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/MessageBusVisitorSession.java
@@ -572,7 +572,9 @@ public class MessageBusVisitorSession implements VisitorSession {
                     params.getDocumentSelection(),
                     bucketIdFactory,
                     1,
-                    progressToken);
+                    progressToken,
+                    params.getSlices(),
+                    params.getSliceId());
         } else {
             if (log.isLoggable(Level.FINE)) {
                 log.log(Level.FINE, "parameters specify explicit bucket set " +
diff --git a/documentapi/src/test/java/com/yahoo/documentapi/VisitorIteratorTestCase.java b/documentapi/src/test/java/com/yahoo/documentapi/VisitorIteratorTestCase.java
index 01cdad244a8..89a1125880d 100755
--- a/documentapi/src/test/java/com/yahoo/documentapi/VisitorIteratorTestCase.java
+++ b/documentapi/src/test/java/com/yahoo/documentapi/VisitorIteratorTestCase.java
@@ -77,6 +77,119 @@ public class VisitorIteratorTestCase {
     }
 
     @Test
+    public void testInvalidSlicing() throws ParseException {
+        int distBits = 4;
+        BucketIdFactory idFactory = new BucketIdFactory();
+        ProgressToken progress = new ProgressToken();
+
+        try {
+            VisitorIterator.createFromDocumentSelection(
+                    "id.group != \"yahoo.com\"", idFactory, distBits, progress, 0, 0);
+        }
+        catch (IllegalArgumentException e) {
+            assertEquals("slices must be positive, but was 0", e.getMessage());
+        }
+
+        try {
+            VisitorIterator.createFromDocumentSelection(
+                    "id.group != \"yahoo.com\"", idFactory, distBits, progress, 1, 1);
+        }
+        catch (IllegalArgumentException e) {
+            assertEquals("sliceId must be in [0, 1), but was 1", e.getMessage());
+        }
+
+        try {
+            VisitorIterator.createFromDocumentSelection(
+                    "id.group != \"yahoo.com\"", idFactory, distBits, progress, 1, -1);
+        }
+        catch (IllegalArgumentException e) {
+            assertEquals("sliceId must be in [0, 1), but was -1", e.getMessage());
+        }
+    }
+
+    @Test
+    public void testIgnoredSlicing() throws ParseException {
+        int distBits = 1;
+        BucketIdFactory idFactory = new BucketIdFactory();
+        ProgressToken progress = new ProgressToken();
+
+        VisitorIterator iter = VisitorIterator.createFromDocumentSelection(
+                "id.group != \"yahoo.com\"", idFactory, distBits, progress, 3, 2);
+
+        // Iterator with a single distribution bit ignores slicing.
+        assertTrue(iter.hasNext());
+        assertEquals(new BucketId(ProgressToken.keyToBucketId(ProgressToken.makeNthBucketKey(0, 1))),
+                     iter.getNext().getSuperbucket());
+
+        assertEquals(new BucketId(ProgressToken.keyToBucketId(ProgressToken.makeNthBucketKey(1, 1))),
+                     iter.getNext().getSuperbucket());
+
+        assertFalse(iter.hasNext());
+    }
+
+    @Test
+    public void testValidSlicing() throws ParseException {
+        int distBits = 4;
+        BucketIdFactory idFactory = new BucketIdFactory();
+        for (int slices = 1; slices <= 1 << distBits + 1; slices++) {
+            long bucketsTotal = 0;
+            for (int sliceId = 0; sliceId < slices; sliceId++) {
+                ProgressToken progress = new ProgressToken();
+
+                // docsel will be unknown --> entire bucket range will be covered
+                VisitorIterator iter = VisitorIterator.createFromDocumentSelection(
+                        "id.group != \"yahoo.com\"", idFactory, distBits, progress, slices, sliceId);
+
+                String context = "slices: " + slices + ", sliceId: " + sliceId;
+                assertEquals(context, progress.getDistributionBitCount(), distBits);
+                assertTrue(context, iter.getBucketSource() instanceof VisitorIterator.DistributionRangeBucketSource);
+
+                assertEquals(context, progress.getFinishedBucketCount(), Math.min(1 << distBits, sliceId));
+                assertEquals(context, progress.getTotalBucketCount(), 1 << distBits);
+
+                // First, get+update half of the buckets, marking them as done
+                long bucketCount = 0;
+
+                // Do buckets in th first half.
+                while (iter.hasNext() && progress.getFinishedBucketCount() < 1 << distBits - 1) {
+                    VisitorIterator.BucketProgress ids = iter.getNext();
+                    iter.update(ids.getSuperbucket(), ProgressToken.FINISHED_BUCKET);
+                    ++bucketCount;
+                    ++bucketsTotal;
+                }
+
+                if (slices + sliceId < 1 << distBits) { // Otherwise, we're already done ...
+                    assertEquals(context, ((1L << distBits - 1) + slices - sliceId - 1) / slices, bucketCount);
+                    // Should be no buckets in limbo at this point
+                    assertFalse(context, progress.hasActive());
+                    assertFalse(context, progress.hasPending());
+                    assertFalse(context, iter.isDone());
+                    assertTrue(context, iter.hasNext());
+                    assertEquals(context, progress.getFinishedBucketCount(), bucketCount * slices + sliceId);
+                    assertFalse(context, progress.isFinished());
+                }
+
+                while (iter.hasNext()) {
+                    VisitorIterator.BucketProgress ids = iter.getNext();
+                    iter.update(ids.getSuperbucket(), ProgressToken.FINISHED_BUCKET);
+                    ++bucketCount;
+                    ++bucketsTotal;
+                }
+
+                assertEquals(context, ((1L << distBits) + slices - sliceId - 1) / slices, bucketCount);
+                // Should be no buckets in limbo at this point
+                assertFalse(context, progress.hasActive());
+                assertFalse(context, progress.hasPending());
+                assertTrue(context, iter.isDone());
+                assertFalse(context, iter.hasNext());
+                assertEquals(context, progress.getFinishedBucketCount(), 1 << distBits);
+                assertTrue(context, progress.isFinished());
+            }
+            assertEquals("slices: " + slices, 1 << distBits, bucketsTotal);
+        }
+    }
+
+    @Test
     public void testProgressSerializationRange() throws ParseException {
         int distBits = 4;