Restore local checkpoint tracker on promotion

When a shard is promoted to replica, it's possible that it was
previously a replica that started following a new primary. When it
started following this new primary, the state of its local checkpoint
tracker was reset. Upon promotion, it's possible that the state of the
local checkpoint tracker has not yet restored from a successful
primary-replica re-sync. To account for this, we must restore the state
of the local checkpoint tracker when a replica shard is promoted to
primary. To do this, we stream the operations in the translog, marking
the operations that are in the translog as completed. We do this before
we fill the gaps on the newly promoted primary, ensuring that we have a
primary shard with a complete history up to the largest maximum sequence
number it has ever seen.

Author: jasontedor

core/src/main/java/org/elasticsearch/index/engine/Engine.java

@@ -1445,6 +1445,15 @@ public interface Warmer {
      */
     public abstract void deactivateThrottling();
 
+    /**
+     * Marks operations in the translog from the specified local checkpoint as completed. This is used to restore the state of the local
+     * checkpoint tracker on primary promotion.
+     *
+     * @param localCheckpoint the local checkpoint
+     * @throws IOException if an I/O exception occurred reading the translog
+     */
+    public abstract void restoreLocalCheckpointTracker(long localCheckpoint) throws IOException;
+
     /**
      * Fills up the local checkpoints history with no-ops until the local checkpoint
      * and the max seen sequence ID are identical.

core/src/main/java/org/elasticsearch/index/engine/InternalEngine.java

@@ -232,6 +232,22 @@ public InternalEngine(EngineConfig engineConfig) throws EngineException {
         logger.trace("created new InternalEngine");
     }
 
+    @Override
+    public void restoreLocalCheckpointTracker(final long localCheckpoint) throws IOException {
+        try (ReleasableLock ignored = writeLock.acquire()) {
+            ensureOpen();
+            try (Translog.View view = getTranslog().newView()) {
+                final Translog.Snapshot snapshot = view.snapshot(localCheckpoint + 1);
+                Translog.Operation operation;
+                while ((operation = snapshot.next()) != null) {
+                    if (operation.seqNo() > localCheckpoint) {
+                        seqNoService().markSeqNoAsCompleted(operation.seqNo());
+                    }
+                }
+            }
+        }
+    }
+
     @Override
     public int fillSeqNoGaps(long primaryTerm) throws IOException {
         try (ReleasableLock ignored = writeLock.acquire()) {

core/src/main/java/org/elasticsearch/index/shard/IndexShard.java

@@ -467,6 +467,7 @@ private void updatePrimaryTerm(
                 () -> {
                     latch.await();
                     try {
+                        getEngine().restoreLocalCheckpointTracker(getLocalCheckpoint());
                         getEngine().fillSeqNoGaps(newPrimaryTerm);
                         primaryReplicaSyncer.accept(IndexShard.this, new ActionListener<ResyncTask>() {
                             @Override
@@ -1660,6 +1661,15 @@ public long getLocalCheckpoint() {
         return getEngine().seqNoService().getLocalCheckpoint();
     }
 
+    /**
+     * Returns the maximum sequence number for the shard.
+     *
+     * @return the maximum sequence number
+     */
+    public long getMaxSeqNo() {
+        return getEngine().seqNoService().getMaxSeqNo();
+    }
+
     /**
      * Returns the global checkpoint for the shard.
      *

core/src/test/java/org/elasticsearch/index/engine/InternalEngineTests.java

@@ -3987,6 +3987,59 @@ private Tuple<Long, Long> getSequenceID(Engine engine, Engine.Get get) throws En
         }
     }
 
+    public void testRestoreLocalCheckpointTracker() throws IOException {
+        engine.close();
+        InternalEngine actualEngine = null;
+        try {
+            final Set<Long> completedSeqNos = new HashSet<>();
+            final SequenceNumbersService seqNoService =
+                    new SequenceNumbersService(
+                            shardId,
+                            defaultSettings,
+                            SequenceNumbersService.NO_OPS_PERFORMED,
+                            SequenceNumbersService.NO_OPS_PERFORMED,
+                            SequenceNumbersService.UNASSIGNED_SEQ_NO) {
+                @Override
+                public void markSeqNoAsCompleted(long seqNo) {
+                    super.markSeqNoAsCompleted(seqNo);
+                    completedSeqNos.add(seqNo);
+                }
+            };
+            actualEngine = new InternalEngine(copy(engine.config(), EngineConfig.OpenMode.OPEN_INDEX_AND_TRANSLOG)) {
+                @Override
+                public SequenceNumbersService seqNoService() {
+                    return seqNoService;
+                }
+            };
+            final int operations = randomIntBetween(0, 1024);
+            final Set<Long> expectedCompletedSeqNos = new HashSet<>();
+            for (int i = 0; i < operations; i++) {
+                if (rarely() && i < operations - 1) {
+                    continue;
+                }
+                final String id = Integer.toString(i);
+                final ParsedDocument doc = testParsedDocument(id, null, testDocumentWithTextField(), SOURCE, null);
+                final Term uid = newUid(doc);
+                actualEngine.index(
+                        new Engine.Index(uid, doc, i, 1, 1, VersionType.EXTERNAL, REPLICA, System.nanoTime(), System.nanoTime(), false));
+                expectedCompletedSeqNos.add((long) i);
+            }
+            final long currentLocalCheckpoint = actualEngine.seqNoService().getLocalCheckpoint();
+            final long resetLocalCheckpoint =
+                    randomIntBetween(Math.toIntExact(SequenceNumbersService.NO_OPS_PERFORMED), Math.toIntExact(currentLocalCheckpoint));
+            actualEngine.seqNoService().resetLocalCheckpoint(resetLocalCheckpoint);
+            completedSeqNos.clear();
+            actualEngine.restoreLocalCheckpointTracker(actualEngine.seqNoService().getLocalCheckpoint());
+            final Set<Long> intersection = new HashSet<>(expectedCompletedSeqNos);
+            intersection.retainAll(LongStream.range(resetLocalCheckpoint + 1, operations).boxed().collect(Collectors.toSet()));
+            assertThat(completedSeqNos, equalTo(intersection));
+            assertThat(actualEngine.seqNoService().getLocalCheckpoint(), equalTo(currentLocalCheckpoint));
+            assertThat(actualEngine.seqNoService().generateSeqNo(), equalTo((long) operations));
+        } finally {
+            IOUtils.close(actualEngine);
+        }
+    }
+
     public void testFillUpSequenceIdGapsOnRecovery() throws IOException {
         final int docs = randomIntBetween(1, 32);
         int numDocsOnReplica = 0;

core/src/test/java/org/elasticsearch/index/shard/IndexShardTests.java

@@ -80,7 +80,6 @@
 import org.elasticsearch.index.mapper.ParsedDocument;
 import org.elasticsearch.index.mapper.SeqNoFieldMapper;
 import org.elasticsearch.index.mapper.SourceToParse;
-import org.elasticsearch.index.seqno.SequenceNumbers;
 import org.elasticsearch.index.seqno.SequenceNumbersService;
 import org.elasticsearch.index.snapshots.IndexShardSnapshotStatus;
 import org.elasticsearch.index.store.Store;
@@ -699,6 +698,59 @@ private void finish() {
         closeShards(indexShard);
     }
 
+    public void testRestoreLocalCheckpointTrackerStateOnPromotion() throws IOException, InterruptedException {
+        final IndexShard indexShard = newStartedShard(false);
+        final int operations = 1024 - scaledRandomIntBetween(0, 1024);
+        indexOnReplicaWithGaps(indexShard, operations, Math.toIntExact(SequenceNumbersService.NO_OPS_PERFORMED));
+
+        final long maxSeqNo = indexShard.getMaxSeqNo();
+        final long globalCheckpointOnReplica = SequenceNumbersService.UNASSIGNED_SEQ_NO;
+        randomIntBetween(
+                Math.toIntExact(SequenceNumbersService.UNASSIGNED_SEQ_NO),
+                Math.toIntExact(indexShard.getLocalCheckpoint()));
+        indexShard.updateGlobalCheckpointOnReplica(globalCheckpointOnReplica);
+
+        final int globalCheckpoint =
+                randomIntBetween(
+                        Math.toIntExact(SequenceNumbersService.UNASSIGNED_SEQ_NO),
+                        Math.toIntExact(indexShard.getLocalCheckpoint()));
+
+        final CountDownLatch latch = new CountDownLatch(1);
+        indexShard.acquireReplicaOperationPermit(
+                indexShard.getPrimaryTerm() + 1,
+                globalCheckpoint,
+                new ActionListener<Releasable>() {
+                    @Override
+                    public void onResponse(Releasable releasable) {
+                        releasable.close();
+                        latch.countDown();
+                    }
+
+                    @Override
+                    public void onFailure(Exception e) {
+
+                    }
+                },
+                ThreadPool.Names.SAME);
+
+        latch.await();
+
+        final ShardRouting newRouting = indexShard.routingEntry().moveActiveReplicaToPrimary();
+        final CountDownLatch resyncLatch = new CountDownLatch(1);
+        indexShard.updateShardState(
+                newRouting,
+                indexShard.getPrimaryTerm() + 1,
+                (s, r) -> resyncLatch.countDown(),
+                1L,
+                Collections.singleton(newRouting.allocationId().getId()),
+                Collections.emptySet());
+        resyncLatch.await();
+        assertThat(indexShard.getLocalCheckpoint(), equalTo(maxSeqNo));
+        assertThat(indexShard.getMaxSeqNo(), equalTo(maxSeqNo));
+
+        closeShards(indexShard);
+    }
+
     public void testThrowBackLocalCheckpointOnReplica() throws IOException, InterruptedException {
         final IndexShard indexShard = newStartedShard(false);