Skip to slice after progress token has been updated, after dist bit count change

2 files changed, 23 insertions, 3 deletions

diff --git a/documentapi/src/main/java/com/yahoo/documentapi/VisitorIterator.java b/documentapi/src/main/java/com/yahoo/documentapi/VisitorIterator.java
index e15512ca71b..b76dee5efa9 100755
--- a/documentapi/src/main/java/com/yahoo/documentapi/VisitorIterator.java
+++ b/documentapi/src/main/java/com/yahoo/documentapi/VisitorIterator.java
@@ -216,7 +216,6 @@ public class VisitorIterator {
                 assert(p.getActiveBucketCount() == 0);
                 p.clearAllBuckets();
                 p.setBucketCursor(0);
-                skipToSlice();
                 return;
             }
 
@@ -409,6 +408,7 @@ public class VisitorIterator {
 
                 correctTruncatedBucketCursor();
                 progressToken.setInconsistentState(false);
+                skipToSlice();
             }
         }
 
diff --git a/documentapi/src/test/java/com/yahoo/documentapi/VisitorIteratorTestCase.java b/documentapi/src/test/java/com/yahoo/documentapi/VisitorIteratorTestCase.java
index fb5f5bd2cfb..5457b0be471 100755
--- a/documentapi/src/test/java/com/yahoo/documentapi/VisitorIteratorTestCase.java
+++ b/documentapi/src/test/java/com/yahoo/documentapi/VisitorIteratorTestCase.java
@@ -118,8 +118,28 @@ public class VisitorIteratorTestCase {
 
         // Iterator with a single distribution bit ignores slicing.
         assertTrue(iter.hasNext());
-        assertEquals(ProgressToken.toBucketId(0, 1), iter.getNext().getSuperbucket());
-        assertEquals(ProgressToken.toBucketId(1, 1), iter.getNext().getSuperbucket());
+
+        VisitorIterator.BucketProgress first = iter.getNext();
+        assertEquals(ProgressToken.toBucketId(0, 1), first.getSuperbucket());
+
+        VisitorIterator.BucketProgress second = iter.getNext();
+        assertEquals(ProgressToken.toBucketId(1, 1), second.getSuperbucket());
+
+        assertFalse(iter.hasNext());
+
+        // Handling wrong distributino reply, but inconsistent state due to the other, active bucket.
+        iter.update(first.getSuperbucket(), first.getProgress());
+        iter.setDistributionBitCount(2);
+        assertEquals(2, iter.getDistributionBitCount());
+        assertEquals(1, progress.getDistributionBitCount());
+
+        // All buckets returned, now the actual distribution bit count change occurs.
+        iter.update(second.getSuperbucket(), second.getProgress());
+        assertEquals(2, iter.getDistributionBitCount());
+        assertEquals(2, progress.getDistributionBitCount());
+
+        assertTrue(iter.hasNext());
+        assertEquals(ProgressToken.toBucketId(2, 2), iter.getNext().getSuperbucket());
         assertFalse(iter.hasNext());
     }