[segment replication] Fix segment replication infinite retry due to stale metadata checkpoint

CHANGELOG.md

@@ -15,6 +15,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Fix flaky test failures in ShardsLimitAllocationDeciderIT ([#20375](https://github.com/opensearch-project/OpenSearch/pull/20375))
 - Prevent criteria update for context aware indices ([#20250](https://github.com/opensearch-project/OpenSearch/pull/20250))
 - Update EncryptedBlobContainer to adhere limits while listing blobs in specific sort order if wrapped blob container supports ([#20514](https://github.com/opensearch-project/OpenSearch/pull/20514))
+- [segment replication] Fix segment replication infinite retry due to stale metadata checkpoint ([#20551](https://github.com/opensearch-project/OpenSearch/pull/20551))
 - Changing opensearch.cgroups.hierarchy.override causes java.lang.SecurityException exception ([#20565](https://github.com/opensearch-project/OpenSearch/pull/20565))
 
 ### Dependencies

server/src/internalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationIT.java

@@ -63,6 +63,8 @@
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.common.util.set.Sets;
+import org.opensearch.core.common.breaker.CircuitBreaker;
+import org.opensearch.core.common.breaker.CircuitBreakingException;
 import org.opensearch.core.common.io.stream.NamedWriteableRegistry;
 import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
 import org.opensearch.core.index.shard.ShardId;
@@ -77,6 +79,7 @@
 import org.opensearch.index.engine.EngineConfig;
 import org.opensearch.index.engine.NRTReplicationReaderManager;
 import org.opensearch.index.shard.IndexShard;
+import org.opensearch.index.store.StoreFileMetadata;
 import org.opensearch.indices.recovery.FileChunkRequest;
 import org.opensearch.indices.replication.checkpoint.PublishCheckpointAction;
 import org.opensearch.indices.replication.common.ReplicationType;
@@ -140,6 +143,87 @@ private static String indexOrAlias() {
         return randomBoolean() ? INDEX_NAME : "alias";
     }
 
+    public void testLocalSegmentReplicationWithException() throws Exception {
+        // this test stubs transport calls specific to node-node replication.
+        assumeFalse(
+            "Skipping the test as its not compatible with segment replication with remote store.",
+            segmentReplicationWithRemoteEnabled()
+        );
+
+        final String primaryNode = internalCluster().startDataOnlyNode();
+        createIndex(INDEX_NAME);
+        ensureYellowAndNoInitializingShards(INDEX_NAME);
+        final String replicaNode = internalCluster().startDataOnlyNode();
+        ensureGreen(INDEX_NAME);
+
+        MockTransportService primaryTransportService = ((MockTransportService) internalCluster().getInstance(
+            TransportService.class,
+            primaryNode
+        ));
+
+        AtomicBoolean mockException = new AtomicBoolean(true);
+        CountDownLatch latch1 = new CountDownLatch(1);
+        CountDownLatch latch2 = new CountDownLatch(1);
+
+        primaryTransportService.addRequestHandlingBehavior(
+            SegmentReplicationSourceService.Actions.GET_SEGMENT_FILES,
+            (handler, request, channel, task) -> {
+                logger.info(
+                    "replicationId {}, get segment files {}",
+                    ((GetSegmentFilesRequest) request).getReplicationId(),
+                    ((GetSegmentFilesRequest) request).getFilesToFetch().stream().map(StoreFileMetadata::name).collect(Collectors.toList())
+                );
+                if (mockException.get()) {
+                    mockException.set(false);
+                    latch1.countDown();
+                    latch2.await();
+                    throw new CircuitBreakingException("mock circuit break exception", CircuitBreaker.Durability.TRANSIENT);
+                } else {
+                    handler.messageReceived(request, channel, task);
+                }
+            }
+        );
+
+        // generate _0.si
+        client().prepareIndex(INDEX_NAME)
+            .setId(String.valueOf(1))
+            .setSource("foo", "bar")
+            .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
+            .get();
+
+        latch1.await();
+
+        MockTransportService replicaTransportService = ((MockTransportService) internalCluster().getInstance(
+            TransportService.class,
+            replicaNode
+        ));
+        replicaTransportService.addRequestHandlingBehavior(
+            PublishCheckpointAction.ACTION_NAME + TransportReplicationAction.REPLICA_ACTION_SUFFIX,
+            (handler, request, channel, task) -> {
+                logger.info("replica receive publish checkpoint request");
+                handler.messageReceived(request, channel, task);
+                latch2.countDown();
+            }
+        );
+
+        // generate _1.si
+        client().prepareIndex(INDEX_NAME)
+            .setId(String.valueOf(2))
+            .setSource("foo2", "bar2")
+            .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
+            .get();
+
+        waitForSearchableDocs(2, primaryNode, replicaNode);
+
+        client().prepareIndex(INDEX_NAME)
+            .setId(String.valueOf(3))
+            .setSource("foo3", "bar3")
+            .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
+            .get();
+
+        waitForSearchableDocs(3, primaryNode, replicaNode);
+    }
+
     public void testAcquireLastIndexCommit() throws Exception {
         final String primaryNode = internalCluster().startDataOnlyNode();
         createIndex(INDEX_NAME);

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

@@ -47,17 +47,20 @@ public abstract class AbstractSegmentReplicationTarget extends ReplicationTarget
     protected final SegmentReplicationSource source;
     protected final SegmentReplicationState state;
     protected final MultiFileWriter multiFileWriter;
+    protected final boolean isRetry;
 
     public AbstractSegmentReplicationTarget(
         String name,
         IndexShard indexShard,
         ReplicationCheckpoint checkpoint,
         SegmentReplicationSource source,
+        boolean isRetry,
         ReplicationListener listener
     ) {
         super(name, indexShard, new ReplicationLuceneIndex(), listener);
         this.checkpoint = checkpoint;
         this.source = source;
+        this.isRetry = isRetry;
         this.state = new SegmentReplicationState(
             indexShard.routingEntry(),
             stateIndex,
@@ -168,7 +171,10 @@ public void startReplication(ActionListener<Void> listener, BiConsumer<Replicati
             // from before a restart, and should accept the primary's current state even if it appears older.
             // See: https://github.com/opensearch-project/OpenSearch/issues/19234
             boolean isRecovering = indexShard.routingEntry().initializing() || indexShard.routingEntry().relocating();
-            if (indexShard.indexSettings().isSegRepLocalEnabled() && checkpoint.isAheadOf(getMetadataCheckpoint) && !isRecovering) {
+            if (indexShard.indexSettings().isSegRepLocalEnabled()
+                && checkpoint.isAheadOf(getMetadataCheckpoint)
+                && false == isRecovering
+                && false == isRetry) {
                 // Fixes https://github.com/opensearch-project/OpenSearch/issues/18490
                 listener.onFailure(
                     new ReplicationFailedException(

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

@@ -33,7 +33,7 @@ public MergedSegmentReplicationTarget(
         SegmentReplicationSource source,
         ReplicationListener listener
     ) {
-        super("merged_segment_replication_target", indexShard, checkpoint, source, listener);
+        super("merged_segment_replication_target", indexShard, checkpoint, source, false, listener);
     }
 
     @Override

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

@@ -42,7 +42,17 @@ public SegmentReplicationTarget(
         SegmentReplicationSource source,
         ReplicationListener listener
     ) {
-        super("replication_target", indexShard, checkpoint, source, listener);
+        this(indexShard, checkpoint, source, false, listener);
+    }
+
+    public SegmentReplicationTarget(
+        IndexShard indexShard,
+        ReplicationCheckpoint checkpoint,
+        SegmentReplicationSource source,
+        boolean isRetry,
+        ReplicationListener listener
+    ) {
+        super("replication_target", indexShard, checkpoint, source, isRetry, listener);
     }
 
     @Override
@@ -128,6 +138,6 @@ protected void finalizeReplication(CheckpointInfoResponse checkpointInfoResponse
 
     @Override
     public SegmentReplicationTarget retryCopy() {
-        return new SegmentReplicationTarget(indexShard, checkpoint, source, listener);
+        return new SegmentReplicationTarget(indexShard, checkpoint, source, isRetry, listener);
     }
 }

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

@@ -295,7 +295,23 @@ public SegmentReplicationTarget get(ShardId shardId) {
      * @param receivedCheckpoint received checkpoint that is checked for processing
      * @param replicaShard       replica shard on which checkpoint is received
      */
-    public synchronized void onNewCheckpoint(final ReplicationCheckpoint receivedCheckpoint, final IndexShard replicaShard) {
+    public void onNewCheckpoint(final ReplicationCheckpoint receivedCheckpoint, final IndexShard replicaShard) {
+        onNewCheckpoint(receivedCheckpoint, replicaShard, false);
+    }
+
+    /**
+     * 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 receivedCheckpoint received checkpoint that is checked for processing
+     * @param replicaShard       replica shard on which checkpoint is received
+     * @param isRetry            is it a retry after failure
+     */
+    public synchronized void onNewCheckpoint(
+        final ReplicationCheckpoint receivedCheckpoint,
+        final IndexShard replicaShard,
+        boolean isRetry
+    ) {
         logger.debug(() -> new ParameterizedMessage("Replica received new replication checkpoint from primary [{}]", receivedCheckpoint));
         // if the shard is in any state
         if (replicaShard.state().equals(IndexShardState.CLOSED)) {
@@ -332,7 +348,7 @@ public synchronized void onNewCheckpoint(final ReplicationCheckpoint receivedChe
             }
             final Thread thread = Thread.currentThread();
             if (replicaShard.shouldProcessCheckpoint(receivedCheckpoint)) {
-                startReplication(replicaShard, receivedCheckpoint, new SegmentReplicationListener() {
+                startReplication(replicaShard, receivedCheckpoint, isRetry, new SegmentReplicationListener() {
                     @Override
                     public void onReplicationDone(SegmentReplicationState state) {
                         logger.debug(
@@ -366,7 +382,7 @@ public void onReplicationFailure(
                         if (sendShardFailure == true) {
                             failShard(e, replicaShard);
                         } else {
-                            processLatestReceivedCheckpoint(replicaShard, thread);
+                            processLatestReceivedCheckpoint(replicaShard, thread, true);
                         }
                     }
                 });
@@ -481,6 +497,11 @@ private DiscoveryNode getPrimaryNode(ShardRouting primaryShard) {
 
     // visible to tests
     protected boolean processLatestReceivedCheckpoint(IndexShard replicaShard, Thread thread) {
+        return processLatestReceivedCheckpoint(replicaShard, thread, false);
+    }
+
+    // visible to tests
+    protected boolean processLatestReceivedCheckpoint(IndexShard replicaShard, Thread thread, boolean isRetry) {
         final ReplicationCheckpoint latestPublishedCheckpoint = replicator.getPrimaryCheckpoint(replicaShard.shardId());
         if (latestPublishedCheckpoint != null) {
             logger.trace(
@@ -494,7 +515,7 @@ protected boolean processLatestReceivedCheckpoint(IndexShard replicaShard, Threa
                 // if we retry ensure the shard is not in the process of being closed.
                 // it will be removed from indexService's collection before the shard is actually marked as closed.
                 if (indicesService.getShardOrNull(replicaShard.shardId()) != null) {
-                    onNewCheckpoint(replicator.getPrimaryCheckpoint(replicaShard.shardId()), replicaShard);
+                    onNewCheckpoint(replicator.getPrimaryCheckpoint(replicaShard.shardId()), replicaShard, isRetry);
                 }
             };
             // Checks if we are using same thread and forks if necessary.
@@ -525,7 +546,24 @@ public SegmentReplicationTarget startReplication(
         final ReplicationCheckpoint checkpoint,
         final SegmentReplicationListener listener
     ) {
-        return replicator.startReplication(indexShard, checkpoint, sourceFactory.get(indexShard), listener);
+        return startReplication(indexShard, checkpoint, false, listener);
+    }
+
+    /**
+     * Start a round of replication and sync to at least the given checkpoint.
+     * @param indexShard - {@link IndexShard} replica shard
+     * @param checkpoint - {@link ReplicationCheckpoint} checkpoint to sync to
+     * @param isRetry - is it a retry after failure
+     * @param listener - {@link ReplicationListener}
+     * @return {@link SegmentReplicationTarget} target event orchestrating the event.
+     */
+    public SegmentReplicationTarget startReplication(
+        final IndexShard indexShard,
+        final ReplicationCheckpoint checkpoint,
+        final boolean isRetry,
+        final SegmentReplicationListener listener
+    ) {
+        return replicator.startReplication(indexShard, checkpoint, sourceFactory.get(indexShard), isRetry, listener);
     }
 
     // pkg-private for integration tests

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

@@ -80,6 +80,7 @@ public void startReplication(IndexShard shard) {
             shard,
             shard.getLatestReplicationCheckpoint(),
             sourceFactory.get().get(shard),
+            false,
             new SegmentReplicationTargetService.SegmentReplicationListener() {
                 @Override
                 public void onReplicationDone(SegmentReplicationState state) {
@@ -105,16 +106,19 @@ void setSourceFactory(SegmentReplicationSourceFactory sourceFactory) {
      * Start a round of replication and sync to at least the given checkpoint.
      * @param indexShard - {@link IndexShard} replica shard
      * @param checkpoint - {@link ReplicationCheckpoint} checkpoint to sync to
+     * @param source - {@link SegmentReplicationSource} segment replication source
+     * @param isRetry - is it a retry after failure
      * @param listener - {@link ReplicationListener}
      * @return {@link SegmentReplicationTarget} target event orchestrating the event.
      */
     SegmentReplicationTarget startReplication(
         final IndexShard indexShard,
         final ReplicationCheckpoint checkpoint,
         final SegmentReplicationSource source,
+        final boolean isRetry,
         final SegmentReplicationTargetService.SegmentReplicationListener listener
     ) {
-        final SegmentReplicationTarget target = new SegmentReplicationTarget(indexShard, checkpoint, source, listener);
+        final SegmentReplicationTarget target = new SegmentReplicationTarget(indexShard, checkpoint, source, isRetry, listener);
         startReplication(target, indexShard.getRecoverySettings().activityTimeout());
         return target;
     }

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

@@ -68,6 +68,7 @@
 
 import static org.opensearch.index.seqno.SequenceNumbers.NO_OPS_PERFORMED;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.doAnswer;
@@ -409,6 +410,7 @@ public void testShardAlreadyReplicating_HigherPrimaryTermReceived() throws Inter
         // skip post replication actions so we can assert execution counts. This will continue to process bc replica's pterm is not advanced
         // post replication.
         doReturn(true).when(serviceSpy).processLatestReceivedCheckpoint(any(), any());
+        doReturn(true).when(serviceSpy).processLatestReceivedCheckpoint(any(), any(), anyBoolean());
         // Create a Mockito spy of target to stub response of few method calls.
 
         CountDownLatch latch = new CountDownLatch(1);
@@ -476,7 +478,7 @@ public void cancel() {
 
         // ensure the old target is cancelled. and new iteration kicks off.
         verify(targetSpy, times(1)).cancel("Cancelling stuck target after new primary");
-        verify(serviceSpy, times(1)).startReplication(eq(replicaShard), any(), any());
+        verify(serviceSpy, times(1)).startReplication(eq(replicaShard), any(), anyBoolean(), any());
     }
 
     public void testMergedSegmentReplicating_HigherPrimaryTermReceived() throws IOException {
@@ -628,10 +630,10 @@ public void testStartReplicationListenerSuccess() throws InterruptedException {
         SegmentReplicationTargetService spy = spy(sut);
         CountDownLatch latch = new CountDownLatch(1);
         doAnswer(i -> {
-            ((SegmentReplicationTargetService.SegmentReplicationListener) i.getArgument(2)).onReplicationDone(state);
+            ((SegmentReplicationTargetService.SegmentReplicationListener) i.getArgument(3)).onReplicationDone(state);
             latch.countDown();
             return null;
-        }).when(spy).startReplication(any(), any(), any());
+        }).when(spy).startReplication(any(), any(), anyBoolean(), any());
         doNothing().when(spy).updateVisibleCheckpoint(eq(0L), any());
         spy.afterIndexShardStarted(replicaShard);
 
@@ -675,7 +677,7 @@ public void testProcessLatestCheckpointIfCheckpointAhead() {
         doReturn(mock(SegmentReplicationTarget.class)).when(service).startReplication(any(), any(), any());
         service.updateLatestReceivedCheckpoint(aheadCheckpoint, replicaShard);
         service.processLatestReceivedCheckpoint(replicaShard, null);
-        verify(service, times(1)).startReplication(eq(replicaShard), eq(aheadCheckpoint), any());
+        verify(service, times(1)).startReplication(eq(replicaShard), eq(aheadCheckpoint), anyBoolean(), any());
     }
 
     public void testOnNewCheckpointInvokedOnClosedShardDoesNothing() throws IOException {