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

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,81 @@ private static String indexOrAlias() {
         return randomBoolean() ? INDEX_NAME : "alias";
     }
 
+    public void testSegmentReplicationWithException() throws Exception {
+        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");
+                latch2.countDown();
+                handler.messageReceived(request, channel, task);
+            }
+        );
+
+        // 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

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

@@ -288,14 +288,23 @@ public SegmentReplicationTarget get(ShardId shardId) {
         return replicator.get(shardId);
     }
 
+    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) {
+    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 +341,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 +375,7 @@ public void onReplicationFailure(
                         if (sendShardFailure == true) {
                             failShard(e, replicaShard);
                         } else {
-                            processLatestReceivedCheckpoint(replicaShard, thread);
+                            processLatestReceivedCheckpoint(replicaShard, thread, true);
                         }
                     }
                 });
@@ -479,8 +488,12 @@ private DiscoveryNode getPrimaryNode(ShardRouting primaryShard) {
         return clusterService.state().nodes().get(primaryShard.currentNodeId());
     }
 
-    // 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 +507,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.
@@ -513,19 +526,29 @@ protected void updateLatestReceivedCheckpoint(ReplicationCheckpoint receivedChec
         replicator.updateReplicationCheckpointStats(receivedCheckpoint, replicaShard);
     }
 
+    public SegmentReplicationTarget startReplication(
+        final IndexShard indexShard,
+        final ReplicationCheckpoint checkpoint,
+        final SegmentReplicationListener 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), 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;
     }