Adding onNewCheckpoint to Start Replication on Replica Shard when Segment Replication is turned on

server/src/main/java/org/opensearch/index/engine/InternalEngine.java

@@ -2289,7 +2289,7 @@ protected SegmentInfos getLastCommittedSegmentInfos() {
     }
 
     @Override
-    public SegmentInfos getLatestSegmentInfos() {
+    protected SegmentInfos getLatestSegmentInfos() {
         OpenSearchDirectoryReader reader = null;
         try {
             reader = internalReaderManager.acquire();

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

@@ -160,7 +160,6 @@
 import org.opensearch.indices.recovery.RecoveryListener;
 import org.opensearch.indices.recovery.RecoveryState;
 import org.opensearch.indices.recovery.RecoveryTarget;
-import org.opensearch.indices.replication.checkpoint.PublishCheckpointRequest;
 import org.opensearch.indices.replication.checkpoint.SegmentReplicationCheckpointPublisher;
 import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint;
 import org.opensearch.repositories.RepositoriesService;
@@ -1374,18 +1373,24 @@ public GatedCloseable<IndexCommit> acquireSafeIndexCommit() throws EngineExcepti
     }
 
     /**
-     * Returns the lastest Replication Checkpoint that shard received
+     * Returns the lastest segmentInfos
      */
-    public ReplicationCheckpoint getLatestReplicationCheckpoint() {
-        return new ReplicationCheckpoint(shardId, 0, 0, 0, 0);
+    public SegmentInfos getLatestSegmentInfos() {
+        return getEngine().getSegmentInfosSnapshot().get();
     }
 
     /**
-     * Invoked when a new checkpoint is received from a primary shard.  Starts the copy process.
+     * Returns the lastest Replication Checkpoint that shard received
      */
-    public synchronized void onNewCheckpoint(final PublishCheckpointRequest request) {
-        assert shardRouting.primary() == false;
-        // TODO
+    public ReplicationCheckpoint getLatestReplicationCheckpoint() {
+        final SegmentInfos latestSegmentInfos = getLatestSegmentInfos();
+        return new ReplicationCheckpoint(
+            this.shardId,
+            getOperationPrimaryTerm(),
+            latestSegmentInfos.getGeneration(),
+            getProcessedLocalCheckpoint(),
+            latestSegmentInfos.getVersion()
+        );
     }
 
     /**

server/src/main/java/org/opensearch/indices/replication/SegmentReplicationTargetService.java

@@ -14,9 +14,11 @@
 import org.opensearch.OpenSearchException;
 import org.opensearch.action.ActionListener;
 import org.opensearch.common.Nullable;
+import org.opensearch.common.inject.Inject;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.index.shard.IndexEventListener;
 import org.opensearch.index.shard.IndexShard;
+import org.opensearch.index.shard.IndexShardState;
 import org.opensearch.index.shard.ShardId;
 import org.opensearch.indices.recovery.FileChunkRequest;
 import org.opensearch.indices.recovery.RecoverySettings;
@@ -38,7 +40,7 @@
  *
  * @opensearch.internal
  */
-public final class SegmentReplicationTargetService implements IndexEventListener {
+public class SegmentReplicationTargetService implements IndexEventListener {
 
     private static final Logger logger = LogManager.getLogger(SegmentReplicationTargetService.class);
 
@@ -58,6 +60,7 @@ public static class Actions {
         public static final String FILE_CHUNK = "internal:index/shard/replication/file_chunk";
     }
 
+    @Inject
     public SegmentReplicationTargetService(
         final ThreadPool threadPool,
         final RecoverySettings recoverySettings,
@@ -84,6 +87,33 @@ public void beforeIndexShardClosed(ShardId shardId, @Nullable IndexShard indexSh
         }
     }
 
+    /**
+     * Invoked when a new checkpoint is received from a primary shard.
+     * It checks if a new checkpoint should be processed or not and starts replication if needed.
+     * @param requestCheckpoint       received checkpoint that is checked for processing
+     * @param indexShard      replica shard on which checkpoint is received
+     */
+    public synchronized void onNewCheckpoint(final ReplicationCheckpoint requestCheckpoint, final IndexShard indexShard) {
+        logger.trace("Checkpoint received {}", () -> requestCheckpoint);
+        if (shouldProcessCheckpoint(requestCheckpoint, indexShard)) {
+            logger.trace("Processing new checkpoint {}", requestCheckpoint);
+            startReplication(requestCheckpoint, indexShard, new SegmentReplicationListener() {
+                @Override
+                public void onReplicationDone(SegmentReplicationState state) {
+                    logger.trace("Replication complete to {}", indexShard.getLatestReplicationCheckpoint());
+                }
+
+                @Override
+                public void onReplicationFailure(SegmentReplicationState state, OpenSearchException e, boolean sendShardFailure) {
+                    if (sendShardFailure == true) {
+                        indexShard.failShard("replication failure", e);
+                    }
+                }
+            });
+
+        }
+    }
+
     public void startReplication(
         final ReplicationCheckpoint checkpoint,
         final IndexShard indexShard,
@@ -98,6 +128,38 @@ public void startReplication(final SegmentReplicationTarget target) {
         threadPool.generic().execute(new ReplicationRunner(replicationId));
     }
 
+    /**
+     * Checks if checkpoint should be processed
+     *
+     * @param requestCheckpoint       received checkpoint that is checked for processing
+     * @param indexShard      replica shard on which checkpoint is received
+     * @return true if checkpoint should be processed
+     */
+    private boolean shouldProcessCheckpoint(ReplicationCheckpoint requestCheckpoint, IndexShard indexShard) {
+        if (indexShard.state().equals(IndexShardState.STARTED) == false) {
+            logger.trace("Ignoring new replication checkpoint - shard is not started {}", indexShard.state());
+            return false;
+        }
+        ReplicationCheckpoint localCheckpoint = indexShard.getLatestReplicationCheckpoint();
+        if (onGoingReplications.isShardReplicating(indexShard.shardId())) {
+            logger.trace("Ignoring new replication checkpoint - shard is currently replicating to a checkpoint");
+            return false;
+        }
+        if (localCheckpoint.isAheadOf(requestCheckpoint)) {
+            logger.trace(
+                "Ignoring new replication checkpoint - Shard is already on checkpoint {} that is ahead of {}",
+                localCheckpoint,
+                requestCheckpoint
+            );
+            return false;
+        }
+        if (localCheckpoint.equals(requestCheckpoint)) {
+            logger.trace("Ignoring new replication checkpoint - Shard is already on checkpoint {}", requestCheckpoint);
+            return false;
+        }
+        return true;
+    }
+
     /**
      * Listener that runs on changes in Replication state
      *

server/src/main/java/org/opensearch/indices/replication/checkpoint/PublishCheckpointAction.java

@@ -28,6 +28,7 @@
 import org.opensearch.index.shard.IndexShard;
 import org.opensearch.index.shard.IndexShardClosedException;
 import org.opensearch.indices.IndicesService;
+import org.opensearch.indices.replication.SegmentReplicationTargetService;
 import org.opensearch.node.NodeClosedException;
 import org.opensearch.tasks.Task;
 import org.opensearch.threadpool.ThreadPool;
@@ -52,6 +53,8 @@ public class PublishCheckpointAction extends TransportReplicationAction<
     public static final String ACTION_NAME = "indices:admin/publishCheckpoint";
     protected static Logger logger = LogManager.getLogger(PublishCheckpointAction.class);
 
+    private final SegmentReplicationTargetService replicationService;
+
     @Inject
     public PublishCheckpointAction(
         Settings settings,
@@ -60,7 +63,8 @@ public PublishCheckpointAction(
         IndicesService indicesService,
         ThreadPool threadPool,
         ShardStateAction shardStateAction,
-        ActionFilters actionFilters
+        ActionFilters actionFilters,
+        SegmentReplicationTargetService targetService
     ) {
         super(
             settings,
@@ -75,6 +79,7 @@ public PublishCheckpointAction(
             PublishCheckpointRequest::new,
             ThreadPool.Names.REFRESH
         );
+        this.replicationService = targetService;
     }
 
     @Override
@@ -165,7 +170,7 @@ protected void shardOperationOnReplica(PublishCheckpointRequest request, IndexSh
         ActionListener.completeWith(listener, () -> {
             logger.trace("Checkpoint received on replica {}", request);
             if (request.getCheckpoint().getShardId().equals(replica.shardId())) {
-                replica.onNewCheckpoint(request);
+                replicationService.onNewCheckpoint(request.getCheckpoint(), replica);
             }
             return new ReplicaResult();
         });

server/src/main/java/org/opensearch/indices/replication/checkpoint/ReplicationCheckpoint.java

@@ -115,7 +115,7 @@ public int hashCode() {
      * Checks if other is aheadof current replication point by comparing segmentInfosVersion. Returns true for null
      */
     public boolean isAheadOf(@Nullable ReplicationCheckpoint other) {
-        return other == null || segmentInfosVersion > other.getSegmentInfosVersion();
+        return other == null || segmentInfosVersion > other.getSegmentInfosVersion() || primaryTerm > other.getPrimaryTerm();
     }
 
     @Override

server/src/main/java/org/opensearch/indices/replication/common/ReplicationCollection.java

@@ -235,6 +235,16 @@ public boolean cancelForShard(ShardId shardId, String reason) {
         return cancelled;
     }
 
+    /**
+     * check if a shard is currently replicating
+     *
+     * @param shardId      shardId for which to check if replicating
+     * @return true if shard is currently replicating
+     */
+    public boolean isShardReplicating(ShardId shardId) {
+        return onGoingTargetEvents.values().stream().anyMatch(t -> t.indexShard.shardId().equals(shardId));
+    }
+
     /**
      * a reference to {@link ReplicationTarget}, which implements {@link AutoCloseable}. closing the reference
      * causes {@link ReplicationTarget#decRef()} to be called. This makes sure that the underlying resources

server/src/main/java/org/opensearch/node/Node.java

@@ -37,6 +37,8 @@
 import org.apache.lucene.util.Constants;
 import org.apache.lucene.util.SetOnce;
 import org.opensearch.index.IndexingPressureService;
+import org.opensearch.indices.replication.SegmentReplicationSourceFactory;
+import org.opensearch.indices.replication.SegmentReplicationTargetService;
 import org.opensearch.watcher.ResourceWatcherService;
 import org.opensearch.Assertions;
 import org.opensearch.Build;
@@ -932,6 +934,15 @@ protected Node(
                         .toInstance(new PeerRecoverySourceService(transportService, indicesService, recoverySettings));
                     b.bind(PeerRecoveryTargetService.class)
                         .toInstance(new PeerRecoveryTargetService(threadPool, transportService, recoverySettings, clusterService));
+                    b.bind(SegmentReplicationTargetService.class)
+                        .toInstance(
+                            new SegmentReplicationTargetService(
+                                threadPool,
+                                recoverySettings,
+                                transportService,
+                                new SegmentReplicationSourceFactory(transportService, recoverySettings, clusterService)
+                            )
+                        );
                 }
                 b.bind(HttpServerTransport.class).toInstance(httpServerTransport);
                 pluginComponents.stream().forEach(p -> b.bind((Class) p.getClass()).toInstance(p));

server/src/test/java/org/opensearch/indices/replication/SegmentReplicationTargetServiceTests.java

@@ -9,6 +9,7 @@
 package org.opensearch.indices.replication;
 
 import org.junit.Assert;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mockito;
 import org.opensearch.OpenSearchException;
 import org.opensearch.action.ActionListener;
@@ -39,7 +40,7 @@ public void setUp() throws Exception {
         final ClusterSettings clusterSettings = new ClusterSettings(settings, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
         final RecoverySettings recoverySettings = new RecoverySettings(settings, clusterSettings);
         final TransportService transportService = mock(TransportService.class);
-        indexShard = newShard(false, settings);
+        indexShard = newStartedShard(false, settings);
         checkpoint = new ReplicationCheckpoint(indexShard.shardId(), 0L, 0L, 0L, 0L);
         SegmentReplicationSourceFactory replicationSourceFactory = mock(SegmentReplicationSourceFactory.class);
         replicationSource = mock(SegmentReplicationSource.class);
@@ -48,12 +49,6 @@ public void setUp() throws Exception {
         sut = new SegmentReplicationTargetService(threadPool, recoverySettings, transportService, replicationSourceFactory);
     }
 
-    @Override
-    public void tearDown() throws Exception {
-        closeShards(indexShard);
-        super.tearDown();
-    }
-
     public void testTargetReturnsSuccess_listenerCompletes() throws IOException {
         final SegmentReplicationTarget target = new SegmentReplicationTarget(
             checkpoint,
@@ -111,6 +106,68 @@ public void onReplicationFailure(SegmentReplicationState state, OpenSearchExcept
         closeShards(indexShard);
     }
 
+    public void testAlreadyOnNewCheckpoint() throws IOException {
+        SegmentReplicationTargetService spy = spy(sut);
+        spy.onNewCheckpoint(indexShard.getLatestReplicationCheckpoint(), indexShard);
+        verify(spy, times(0)).startReplication(any(), any(), any());
+        closeShards(indexShard);
+    }
+
+    public void testShardAlreadyReplicating() throws IOException {
+        SegmentReplicationTargetService spy = spy(sut);
+        final SegmentReplicationTarget target = new SegmentReplicationTarget(
+            checkpoint,
+            indexShard,
+            replicationSource,
+            mock(SegmentReplicationTargetService.SegmentReplicationListener.class)
+        );
+        final SegmentReplicationTarget spyTarget = Mockito.spy(target);
+        spy.startReplication(spyTarget);
+        spy.onNewCheckpoint(checkpoint, indexShard);
+        verify(spy, times(0)).startReplication(any(), any(), any());
+        closeShards(indexShard);
+    }
+
+    public void testNewCheckpointBehindCurrentCheckpoint() throws IOException {
+        SegmentReplicationTargetService spy = spy(sut);
+        spy.onNewCheckpoint(checkpoint, indexShard);
+        verify(spy, times(0)).startReplication(any(), any(), any());
+        closeShards(indexShard);
+    }
+
+    public void testShardNotStarted() throws IOException {
+        SegmentReplicationTargetService spy = spy(sut);
+        IndexShard shard = newShard(false);
+        spy.onNewCheckpoint(checkpoint, shard);
+        verify(spy, times(0)).startReplication(any(), any(), any());
+        closeShards(shard);
+        closeShards(indexShard);
+    }
+
+    public void testShouldProcessCheckpoint() throws IOException {
+        allowShardFailures();
+        SegmentReplicationTargetService spy = spy(sut);
+        IndexShard spyShard = spy(indexShard);
+        ReplicationCheckpoint cp = indexShard.getLatestReplicationCheckpoint();
+        ReplicationCheckpoint newCheckpoint = new ReplicationCheckpoint(
+            cp.getShardId(),
+            cp.getPrimaryTerm(),
+            cp.getSegmentsGen(),
+            cp.getSeqNo(),
+            cp.getSegmentInfosVersion() + 1
+        );
+        spy.onNewCheckpoint(newCheckpoint, spyShard);
+        ArgumentCaptor<SegmentReplicationTargetService.SegmentReplicationListener> captor = ArgumentCaptor.forClass(
+            SegmentReplicationTargetService.SegmentReplicationListener.class
+        );
+        verify(spy, times(1)).startReplication(any(), any(), captor.capture());
+        SegmentReplicationTargetService.SegmentReplicationListener listener = captor.getValue();
+        listener.onFailure(new SegmentReplicationState(), new OpenSearchException("testing"), true);
+        verify(spyShard).failShard(any(), any());
+        closeShard(indexShard, false);
+
+    }
+
     public void testBeforeIndexShardClosed_CancelsOngoingReplications() throws IOException {
         final SegmentReplicationTarget target = new SegmentReplicationTarget(
             checkpoint,