Introduce concurrent translog recovery to accelerate segment replication primary promotion

server/src/main/java/org/opensearch/common/settings/ClusterSettings.java

@@ -346,6 +346,8 @@ public void apply(Settings value, Settings current, Settings previous) {
                 RecoverySettings.INDICES_RECOVERY_MAX_CONCURRENT_REMOTE_STORE_STREAMS_SETTING,
                 RecoverySettings.INDICES_INTERNAL_REMOTE_UPLOAD_TIMEOUT,
                 RecoverySettings.INDICES_RECOVERY_CHUNK_SIZE_SETTING,
+                RecoverySettings.INDICES_TRANSLOG_CONCURRENT_RECOVERY_ENABLE,
+                RecoverySettings.INDICES_TRANSLOG_CONCURRENT_RECOVERY_BATCH_SIZE,
                 ThrottlingAllocationDecider.CLUSTER_ROUTING_ALLOCATION_NODE_INITIAL_PRIMARIES_RECOVERIES_SETTING,
                 ThrottlingAllocationDecider.CLUSTER_ROUTING_ALLOCATION_NODE_INITIAL_REPLICAS_RECOVERIES_SETTING,
                 ThrottlingAllocationDecider.CLUSTER_ROUTING_ALLOCATION_NODE_CONCURRENT_INCOMING_RECOVERIES_SETTING,

server/src/main/java/org/opensearch/index/shard/IndexShard.java

@@ -103,6 +103,7 @@
 import org.opensearch.common.util.concurrent.AbstractRunnable;
 import org.opensearch.common.util.concurrent.AsyncIOProcessor;
 import org.opensearch.common.util.concurrent.BufferedAsyncIOProcessor;
+import org.opensearch.common.util.concurrent.OpenSearchThreadPoolExecutor;
 import org.opensearch.common.util.concurrent.RunOnce;
 import org.opensearch.common.util.concurrent.ThreadContext;
 import org.opensearch.common.util.io.IOUtils;
@@ -240,6 +241,7 @@
 import java.util.Set;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -5274,14 +5276,60 @@ public void close() throws IOException {
             newEngineReference.set(engineFactory.newReadWriteEngine(newEngineConfig(replicationTracker)));
             onNewEngine(newEngineReference.get());
         }
-        final TranslogRecoveryRunner translogRunner = (snapshot) -> runTranslogRecovery(
-            newEngineReference.get(),
-            snapshot,
-            Engine.Operation.Origin.LOCAL_RESET,
-            () -> {
-                // TODO: add a dedicate recovery stats for the reset translog
+        final TranslogRecoveryRunner translogRunner = (snapshot) -> {
+            Engine engine = newEngineReference.get();
+            assert null != engine;
+            int totalOperations = snapshot.totalOperations();
+            long batchSize = recoverySettings.getTranslogConcurrentRecoveryBatchSize();
+            boolean threadPoolNotBusy = ((OpenSearchThreadPoolExecutor) threadPool.executor(ThreadPool.Names.TRANSLOG_RECOVERY)).getQueue()
+                .isEmpty();
+            if (recoverySettings.isTranslogConcurrentRecoveryEnable()
+                && indexSettings.isSegRepEnabledOrRemoteNode()
+                && totalOperations > batchSize
+                && threadPoolNotBusy) {
+                long localCheckpoint = engine.getProcessedLocalCheckpoint();
+                List<Tuple<Future<Integer>, Translog.Snapshot>> translogSnapshotsFutureList = new ArrayList<>();
+                for (int i = 0; i < (totalOperations + batchSize - 1) / batchSize; i++) {
+                    long start = localCheckpoint + 1 + (long) i * batchSize;
+                    long end = (i == totalOperations / batchSize) ? Long.MAX_VALUE : start + batchSize - 1;
+                    Translog.Snapshot translogSnapshot = engine.translogManager().newChangesSnapshot(start, end, false);
+                    translogSnapshotsFutureList.add(
+                        new Tuple<>(
+                            threadPool.executor(ThreadPool.Names.TRANSLOG_RECOVERY)
+                                .submit(() -> runTranslogRecovery(engine, translogSnapshot, Engine.Operation.Origin.LOCAL_RESET, () -> {
+                                    // TODO: add a dedicate recovery stats for the reset translog
+                                })),
+                            translogSnapshot
+                        )
+                    );
+                }
+                Exception exception = null;
+                int totalRecovered = 0;
+                for (Tuple<Future<Integer>, Translog.Snapshot> translogSnapshotFuture : translogSnapshotsFutureList) {
+                    try {
+                        int recoveredOps = translogSnapshotFuture.v1().get();
+                        totalRecovered += recoveredOps;
+                    } catch (Exception e) {
+                        if (exception == null) {
+                            exception = e;
+                        } else {
+                            exception.addSuppressed(e);
+                        }
+                    } finally {
+                        Translog.Snapshot translogSnapshot = translogSnapshotFuture.v2();
+                        IOUtils.closeWhileHandlingException(translogSnapshot);
+                    }
+                }
+                if (exception != null) {
+                    throw new IOException("generate exception when concurrent translog recovery", exception);
+                }
+                return totalRecovered;
+            } else {
+                return runTranslogRecovery(newEngineReference.get(), snapshot, Engine.Operation.Origin.LOCAL_RESET, () -> {
+                    // TODO: add a dedicate recovery stats for the reset translog
+                });
             }
-        );
+        };
 
         // When the new engine is created, translogs are synced from remote store onto local. Since remote store is the source
         // of truth for translog, we play all translogs that exists locally. Otherwise, the recoverUpto happens upto global checkpoint.

server/src/main/java/org/opensearch/indices/recovery/RecoverySettings.java

@@ -230,6 +230,21 @@ public class RecoverySettings {
         Property.NodeScope
     );
 
+    public static final Setting<Boolean> INDICES_TRANSLOG_CONCURRENT_RECOVERY_ENABLE = Setting.boolSetting(
+        "indices.translog_concurrent_recovery.enable",
+        false,
+        Property.Dynamic,
+        Property.NodeScope
+    );
+
+    public static final Setting<Long> INDICES_TRANSLOG_CONCURRENT_RECOVERY_BATCH_SIZE = Setting.longSetting(
+        "indices.translog_concurrent_recovery.batch_size",
+        500000,
+        1,
+        Property.Dynamic,
+        Property.NodeScope
+    );
+
     private volatile ByteSizeValue mergedSegmentWarmerMinSegmentSizeThreshold;
     private volatile ByteSizeValue recoveryMaxBytesPerSec;
     private volatile ByteSizeValue replicationMaxBytesPerSec;
@@ -252,6 +267,9 @@ public class RecoverySettings {
     private volatile TimeValue internalRemoteUploadTimeout;
     private volatile TimeValue mergedSegmentReplicationTimeout;
 
+    private volatile boolean isTranslogConcurrentRecoveryEnable;
+    private volatile long translogConcurrentRecoveryBatchSize;
+
     public RecoverySettings(Settings settings, ClusterSettings clusterSettings) {
         this.retryDelayStateSync = INDICES_RECOVERY_RETRY_DELAY_STATE_SYNC_SETTING.get(settings);
         this.maxConcurrentFileChunks = INDICES_RECOVERY_MAX_CONCURRENT_FILE_CHUNKS_SETTING.get(settings);
@@ -286,6 +304,9 @@ public RecoverySettings(Settings settings, ClusterSettings clusterSettings) {
         this.internalRemoteUploadTimeout = INDICES_INTERNAL_REMOTE_UPLOAD_TIMEOUT.get(settings);
         this.chunkSize = INDICES_RECOVERY_CHUNK_SIZE_SETTING.get(settings);
 
+        this.isTranslogConcurrentRecoveryEnable = INDICES_TRANSLOG_CONCURRENT_RECOVERY_ENABLE.get(settings);
+        this.translogConcurrentRecoveryBatchSize = INDICES_TRANSLOG_CONCURRENT_RECOVERY_BATCH_SIZE.get(settings);
+
         clusterSettings.addSettingsUpdateConsumer(INDICES_RECOVERY_MAX_BYTES_PER_SEC_SETTING, this::setRecoveryMaxBytesPerSec);
         clusterSettings.addSettingsUpdateConsumer(INDICES_REPLICATION_MAX_BYTES_PER_SEC_SETTING, this::setReplicationMaxBytesPerSec);
         clusterSettings.addSettingsUpdateConsumer(
@@ -324,6 +345,11 @@ public RecoverySettings(Settings settings, ClusterSettings clusterSettings) {
             INDICES_RECOVERY_INTERNAL_ACTION_RETRY_TIMEOUT_SETTING,
             this::setInternalActionRetryTimeout
         );
+        clusterSettings.addSettingsUpdateConsumer(INDICES_TRANSLOG_CONCURRENT_RECOVERY_ENABLE, this::setTranslogConcurrentRecoveryEnable);
+        clusterSettings.addSettingsUpdateConsumer(
+            INDICES_TRANSLOG_CONCURRENT_RECOVERY_BATCH_SIZE,
+            this::setTranslogConcurrentRecoveryBatchSize
+        );
     }
 
     private void setMergedSegmentWarmerMinSegmentSizeThreshold(ByteSizeValue value) {
@@ -493,4 +519,19 @@ public void setIndicesMergedSegmentReplicationWarmerEnabled(boolean mergedSegmen
         this.mergedSegmentReplicationWarmerEnabled = mergedSegmentReplicationWarmerEnabled;
     }
 
+    public boolean isTranslogConcurrentRecoveryEnable() {
+        return isTranslogConcurrentRecoveryEnable;
+    }
+
+    private void setTranslogConcurrentRecoveryEnable(boolean translogConcurrentRecoveryEnable) {
+        isTranslogConcurrentRecoveryEnable = translogConcurrentRecoveryEnable;
+    }
+
+    public long getTranslogConcurrentRecoveryBatchSize() {
+        return translogConcurrentRecoveryBatchSize;
+    }
+
+    private void setTranslogConcurrentRecoveryBatchSize(long translogConcurrentRecoveryBatchSize) {
+        this.translogConcurrentRecoveryBatchSize = translogConcurrentRecoveryBatchSize;
+    }
 }

server/src/main/java/org/opensearch/threadpool/ThreadPool.java

@@ -99,6 +99,7 @@ public class ThreadPool implements ReportingService<ThreadPoolInfo>, Scheduler {
     public static class Names {
         public static final String SAME = "same";
         public static final String GENERIC = "generic";
+        public static final String TRANSLOG_RECOVERY = "translog_recovery";
         @Deprecated
         public static final String LISTENER = "listener";
         public static final String GET = "get";
@@ -179,6 +180,7 @@ public static ThreadPoolType fromType(String type) {
         HashMap<String, ThreadPoolType> map = new HashMap<>();
         map.put(Names.SAME, ThreadPoolType.DIRECT);
         map.put(Names.GENERIC, ThreadPoolType.SCALING);
+        map.put(Names.TRANSLOG_RECOVERY, ThreadPoolType.FIXED);
         map.put(Names.LISTENER, ThreadPoolType.FIXED);
         map.put(Names.GET, ThreadPoolType.FIXED);
         map.put(Names.ANALYZE, ThreadPoolType.FIXED);
@@ -258,6 +260,7 @@ public ThreadPool(
         final int genericThreadPoolMax = boundedBy(4 * allocatedProcessors, 128, 512);
         final int snapshotDeletionPoolMax = boundedBy(4 * allocatedProcessors, 64, 256);
         builders.put(Names.GENERIC, new ScalingExecutorBuilder(Names.GENERIC, 4, genericThreadPoolMax, TimeValue.timeValueSeconds(30)));
+        builders.put(Names.TRANSLOG_RECOVERY, new FixedExecutorBuilder(settings, Names.TRANSLOG_RECOVERY, allocatedProcessors, -1));
         builders.put(Names.WRITE, new FixedExecutorBuilder(settings, Names.WRITE, allocatedProcessors, 10000));
         builders.put(Names.GET, new FixedExecutorBuilder(settings, Names.GET, allocatedProcessors, 1000));
         builders.put(Names.ANALYZE, new FixedExecutorBuilder(settings, Names.ANALYZE, 1, 16));

server/src/test/java/org/opensearch/index/shard/SegmentReplicationWithNodeToNodeIndexShardTests.java

@@ -35,6 +35,7 @@
 import org.opensearch.indices.replication.SegmentReplicationSourceFactory;
 import org.opensearch.indices.replication.SegmentReplicationTarget;
 import org.opensearch.indices.replication.SegmentReplicationTargetService;
+import org.opensearch.indices.replication.checkpoint.MergedSegmentPublisher;
 import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.TransportService;
@@ -713,4 +714,88 @@ public void testSegmentReplication_Index_Update_Delete() throws Exception {
         }
     }
 
+    public void testPrimaryPromotionWithConcurrentTranslogRecovery() throws Exception {
+        final RecoverySettings recoverySettings = new RecoverySettings(
+            Settings.builder()
+                .put(RecoverySettings.INDICES_TRANSLOG_CONCURRENT_RECOVERY_ENABLE.getKey(), true)
+                .put(RecoverySettings.INDICES_TRANSLOG_CONCURRENT_RECOVERY_BATCH_SIZE.getKey(), 10)
+                .build(),
+            new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS)
+        );
+        try (
+            ReplicationGroup shards = createGroup(
+                2,
+                getIndexSettings(),
+                indexMapping,
+                new NRTReplicationEngineFactory(),
+                recoverySettings,
+                MergedSegmentPublisher.EMPTY
+            )
+        ) {
+            shards.startAll();
+            IndexShard oldPrimary = shards.getPrimary();
+            final IndexShard nextPrimary = shards.getReplicas().get(0);
+            final IndexShard replica = shards.getReplicas().get(1);
+
+            // 1. Create ops that are in the index and xlog of both shards but not yet part of a commit point.
+            final int numDocs = shards.indexDocs(randomInt(10));
+
+            // refresh and copy the segments over.
+            oldPrimary.refresh("Test");
+            replicateSegments(oldPrimary, shards.getReplicas());
+
+            // at this point both shards should have numDocs persisted and searchable.
+            assertDocCounts(oldPrimary, numDocs, numDocs);
+            for (IndexShard shard : shards.getReplicas()) {
+                assertDocCounts(shard, numDocs, numDocs);
+            }
+            assertEqualTranslogOperations(shards, oldPrimary);
+
+            // 2. Create ops that are in the replica's xlog, not in the index.
+            // index some more into both but don't replicate. replica will have only numDocs searchable, but should have totalDocs
+            // persisted.
+            final int additonalDocs = shards.indexDocs(randomIntBetween(200, 300));
+            final int totalDocs = numDocs + additonalDocs;
+            logger.info("total docs: {}, additonalDocs {}", totalDocs, additonalDocs);
+
+            assertDocCounts(oldPrimary, totalDocs, totalDocs);
+            assertEqualTranslogOperations(shards, oldPrimary);
+            for (IndexShard shard : shards.getReplicas()) {
+                assertDocCounts(shard, totalDocs, numDocs);
+            }
+            assertEquals(totalDocs, oldPrimary.translogStats().estimatedNumberOfOperations());
+            assertEquals(totalDocs, oldPrimary.translogStats().estimatedNumberOfOperations());
+            assertEquals(totalDocs, nextPrimary.translogStats().estimatedNumberOfOperations());
+            assertEquals(totalDocs, replica.translogStats().estimatedNumberOfOperations());
+            assertEquals(totalDocs, nextPrimary.translogStats().getUncommittedOperations());
+            assertEquals(totalDocs, replica.translogStats().getUncommittedOperations());
+
+            // promote the replica
+            shards.syncGlobalCheckpoint();
+            shards.promoteReplicaToPrimary(nextPrimary);
+
+            // close and start the oldPrimary as a replica.
+            oldPrimary.close("demoted", false, false);
+            oldPrimary.store().close();
+            oldPrimary = shards.addReplicaWithExistingPath(oldPrimary.shardPath(), oldPrimary.routingEntry().currentNodeId());
+            shards.recoverReplica(oldPrimary);
+
+            assertEquals(NRTReplicationEngine.class, oldPrimary.getEngine().getClass());
+            assertEquals(InternalEngine.class, nextPrimary.getEngine().getClass());
+            assertDocCounts(nextPrimary, totalDocs, totalDocs);
+            assertEquals(0, nextPrimary.translogStats().estimatedNumberOfOperations());
+
+            // refresh and push segments to our other replica.
+            nextPrimary.refresh("test");
+            replicateSegments(nextPrimary, asList(replica));
+
+            for (IndexShard shard : shards) {
+                assertConsistentHistoryBetweenTranslogAndLucene(shard);
+            }
+            final List<DocIdSeqNoAndSource> docsAfterRecovery = getDocIdAndSeqNos(shards.getPrimary());
+            for (IndexShard shard : shards.getReplicas()) {
+                assertThat(shard.routingEntry().toString(), getDocIdAndSeqNos(shard), equalTo(docsAfterRecovery));
+            }
+        }
+    }
 }

test/framework/src/main/java/org/opensearch/index/replication/OpenSearchIndexLevelReplicationTestCase.java

@@ -312,7 +312,7 @@ protected ReplicationGroup(
             this.indexMetadata = indexMetadata;
             updateAllocationIDsOnPrimary();
             for (int i = 0; i < indexMetadata.getNumberOfReplicas(); i++) {
-                addReplica(remotePath);
+                addReplica(remotePath, recoverySettings);
             }
         }
 
@@ -469,6 +469,23 @@ public IndexShard addReplica(Path remotePath) throws IOException {
             return replica;
         }
 
+        public IndexShard addReplica(Path remotePath, RecoverySettings recoverySettings) throws IOException {
+            final ShardRouting replicaRouting = createShardRouting("s" + replicaId.incrementAndGet(), false);
+            final IndexShard replica = newShard(
+                replicaRouting,
+                indexMetadata,
+                null,
+                getEngineFactory(replicaRouting),
+                () -> {},
+                retentionLeaseSyncer,
+                recoverySettings,
+                remotePath,
+                MergedSegmentPublisher.EMPTY
+            );
+            addReplica(replica);
+            return replica;
+        }
+
         public synchronized void addReplica(IndexShard replica) throws IOException {
             assert shardRoutings().stream().anyMatch(shardRouting -> shardRouting.isSameAllocation(replica.routingEntry())) == false
                 : "replica with aId [" + replica.routingEntry().allocationId() + "] already exists";