[Backport 2.x] [Segment Replication] Backport PR's : #3525 #3533 #3540 …

…#3943 #3963 From main branch (#4181)

* Resolving import conflict in Node.java and mergining PR #3525.

Signed-off-by: Rishikesh1159 <rishireddy1159@gmail.com>

* Resolving conflicts and merging PR #3533.

Signed-off-by: Rishikesh1159 <rishireddy1159@gmail.com>

* Resolving conflicts and Merging PR #3540.

Signed-off-by: Rishikesh1159 <rishireddy1159@gmail.com>

* Applying spotlesscheck and fixing wildcard imports.

Signed-off-by: Rishikesh1159 <rishireddy1159@gmail.com>

* [Segment Replication] Fixing flaky test failure happening for testShardAlreadyReplicating() (#3943)

* Fixing flaky test failure happening for testShardAlreadyReplicating()

Signed-off-by: Rishikesh1159 <rishireddy1159@gmail.com>

* Fix possible flaky test for testBeforeIndexShardClosed_CancelsOngoingReplications() (#3963)

* Fixing flaky test failure happening for testShardAlreadyReplicating()

Signed-off-by: Rishikesh1159 <rishireddy1159@gmail.com>

* Removing assert segrep() in getProcessedLocalCheckpoint() of Index.shard class.

Signed-off-by: Rishikesh1159 <rishireddy1159@gmail.com>

* Adding back assert statement and make index setting to segment replication in SegmentReplicationSourceHandlerTests and SegmentReplicationTargetServiceTests.

Signed-off-by: Rishikesh1159 <rishireddy1159@gmail.com>

* Revert "Adding back assert statement and make index setting to segment replication in SegmentReplicationSourceHandlerTests and SegmentReplicationTargetServiceTests."
Signed-off-by: Rishikesh1159 <rishireddy1159@gmail.com>
This reverts commit 8c5753b.

Signed-off-by: Rishikesh1159 <rishireddy1159@gmail.com>
Co-authored-by: Marc Handalian <handalm@amazon.com>
Co-authored-by: Poojita Raj <poojiraj@amazon.com>

server/src/main/java/org/opensearch/OpenSearchException.java

@@ -67,6 +67,7 @@
 import static java.util.Collections.emptyMap;
 import static java.util.Collections.singletonMap;
 import static java.util.Collections.unmodifiableMap;
+import static org.opensearch.Version.V_2_1_0;
 import static org.opensearch.cluster.metadata.IndexMetadata.INDEX_UUID_NA_VALUE;
 import static org.opensearch.common.xcontent.XContentParserUtils.ensureExpectedToken;
 import static org.opensearch.common.xcontent.XContentParserUtils.ensureFieldName;
@@ -1594,6 +1595,12 @@ private enum OpenSearchExceptionHandle {
             org.opensearch.transport.NoSeedNodeLeftException::new,
             160,
             LegacyESVersion.V_7_10_0
+        ),
+        REPLICATION_FAILED_EXCEPTION(
+            org.opensearch.indices.replication.common.ReplicationFailedException.class,
+            org.opensearch.indices.replication.common.ReplicationFailedException::new,
+            161,
+            V_2_1_0
         );
 
         final Class<? extends OpenSearchException> exceptionClass;

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

@@ -84,7 +84,6 @@
 import org.opensearch.index.translog.Translog;
 import org.opensearch.index.translog.TranslogStats;
 import org.opensearch.index.translog.TranslogDeletionPolicy;
-import org.opensearch.index.translog.TranslogStats;
 import org.opensearch.search.suggest.completion.CompletionStats;
 
 import java.io.Closeable;

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

@@ -2285,7 +2285,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

@@ -162,9 +162,9 @@
 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.indices.replication.checkpoint.SegmentReplicationCheckpointPublisher;
 import org.opensearch.repositories.RepositoriesService;
 import org.opensearch.repositories.Repository;
 import org.opensearch.rest.RestStatus;
@@ -1361,6 +1361,12 @@ public GatedCloseable<IndexCommit> acquireLastIndexCommit(boolean flushFirst) th
         }
     }
 
+    public void finalizeReplication(SegmentInfos infos, long seqNo) throws IOException {
+        if (getReplicationEngine().isPresent()) {
+            getReplicationEngine().get().updateSegments(infos, seqNo);
+        }
+    }
+
     /**
      * Snapshots the most recent safe index commit from the currently running engine.
      * All index files referenced by this index commit won't be freed until the commit/snapshot is closed.
@@ -1379,15 +1385,60 @@ public GatedCloseable<IndexCommit> acquireSafeIndexCommit() throws EngineExcepti
      * Returns the lastest Replication Checkpoint that shard received
      */
     public ReplicationCheckpoint getLatestReplicationCheckpoint() {
-        return new ReplicationCheckpoint(shardId, 0, 0, 0, 0);
+        try (final GatedCloseable<SegmentInfos> snapshot = getSegmentInfosSnapshot()) {
+            return Optional.ofNullable(snapshot.get())
+                .map(
+                    segmentInfos -> new ReplicationCheckpoint(
+                        this.shardId,
+                        getOperationPrimaryTerm(),
+                        segmentInfos.getGeneration(),
+                        getProcessedLocalCheckpoint(),
+                        segmentInfos.getVersion()
+                    )
+                )
+                .orElse(
+                    new ReplicationCheckpoint(
+                        shardId,
+                        getOperationPrimaryTerm(),
+                        SequenceNumbers.NO_OPS_PERFORMED,
+                        getProcessedLocalCheckpoint(),
+                        SequenceNumbers.NO_OPS_PERFORMED
+                    )
+                );
+        } catch (IOException ex) {
+            throw new OpenSearchException("Error Closing SegmentInfos Snapshot", ex);
+        }
     }
 
     /**
-     * Invoked when a new checkpoint is received from a primary shard.  Starts the copy process.
+     * Checks if checkpoint should be processed
+     *
+     * @param requestCheckpoint       received checkpoint that is checked for processing
+     * @return true if checkpoint should be processed
      */
-    public synchronized void onNewCheckpoint(final PublishCheckpointRequest request) {
-        assert shardRouting.primary() == false;
-        // TODO
+    public final boolean shouldProcessCheckpoint(ReplicationCheckpoint requestCheckpoint) {
+        if (state().equals(IndexShardState.STARTED) == false) {
+            logger.trace(() -> new ParameterizedMessage("Ignoring new replication checkpoint - shard is not started {}", state()));
+            return false;
+        }
+        ReplicationCheckpoint localCheckpoint = getLatestReplicationCheckpoint();
+        if (localCheckpoint.isAheadOf(requestCheckpoint)) {
+            logger.trace(
+                () -> new ParameterizedMessage(
+                    "Ignoring new replication checkpoint - Shard is already on checkpoint {} that is ahead of {}",
+                    localCheckpoint,
+                    requestCheckpoint
+                )
+            );
+            return false;
+        }
+        if (localCheckpoint.equals(requestCheckpoint)) {
+            logger.trace(
+                () -> new ParameterizedMessage("Ignoring new replication checkpoint - Shard is already on checkpoint {}", requestCheckpoint)
+            );
+            return false;
+        }
+        return true;
     }
 
     /**
@@ -2657,7 +2708,6 @@ public long getLocalCheckpoint() {
      * Also see {@link #getLocalCheckpoint()}.
      */
     public long getProcessedLocalCheckpoint() {
-        assert indexSettings.isSegRepEnabled();
         // Returns checkpoint only if the current engine is an instance of NRTReplicationEngine or InternalEngine
         return getReplicationEngine().map(NRTReplicationEngine::getProcessedLocalCheckpoint).orElseGet(() -> {
             final Engine engine = getEngine();

server/src/main/java/org/opensearch/index/store/Store.java

@@ -64,6 +64,7 @@
 import org.apache.lucene.util.BytesRefBuilder;
 import org.apache.lucene.util.Version;
 import org.opensearch.ExceptionsHelper;
+import org.opensearch.common.Nullable;
 import org.opensearch.common.UUIDs;
 import org.opensearch.common.bytes.BytesReference;
 import org.opensearch.common.io.Streams;
@@ -706,6 +707,51 @@ public void cleanupAndVerify(String reason, MetadataSnapshot sourceMetadata) thr
         }
     }
 
+    /**
+     * This method deletes every file in this store that is not contained in either the remote or local metadata snapshots.
+     * This method is used for segment replication when the in memory SegmentInfos can be ahead of the on disk segment file.
+     * In this case files from both snapshots must be preserved. Verification has been done that all files are present on disk.
+     * @param reason         the reason for this cleanup operation logged for each deleted file
+     * @param localSnapshot  The local snapshot from in memory SegmentInfos.
+     * @throws IllegalStateException if the latest snapshot in this store differs from the given one after the cleanup.
+     */
+    public void cleanupAndPreserveLatestCommitPoint(String reason, MetadataSnapshot localSnapshot) throws IOException {
+        // fetch a snapshot from the latest on disk Segments_N file. This can be behind
+        // the passed in local in memory snapshot, so we want to ensure files it references are not removed.
+        metadataLock.writeLock().lock();
+        try (Lock writeLock = directory.obtainLock(IndexWriter.WRITE_LOCK_NAME)) {
+            cleanupFiles(reason, localSnapshot, getMetadata(readLastCommittedSegmentsInfo()));
+        } finally {
+            metadataLock.writeLock().unlock();
+        }
+    }
+
+    private void cleanupFiles(String reason, MetadataSnapshot localSnapshot, @Nullable MetadataSnapshot additionalSnapshot)
+        throws IOException {
+        assert metadataLock.isWriteLockedByCurrentThread();
+        for (String existingFile : directory.listAll()) {
+            if (Store.isAutogenerated(existingFile)
+                || localSnapshot.contains(existingFile)
+                || (additionalSnapshot != null && additionalSnapshot.contains(existingFile))) {
+                // don't delete snapshot file, or the checksums file (note, this is extra protection since the Store won't delete
+                // checksum)
+                continue;
+            }
+            try {
+                directory.deleteFile(reason, existingFile);
+            } catch (IOException ex) {
+                if (existingFile.startsWith(IndexFileNames.SEGMENTS) || existingFile.startsWith(CORRUPTED_MARKER_NAME_PREFIX)) {
+                    // TODO do we need to also fail this if we can't delete the pending commit file?
+                    // if one of those files can't be deleted we better fail the cleanup otherwise we might leave an old commit
+                    // point around?
+                    throw new IllegalStateException("Can't delete " + existingFile + " - cleanup failed", ex);
+                }
+                logger.debug(() -> new ParameterizedMessage("failed to delete file [{}]", existingFile), ex);
+                // ignore, we don't really care, will get deleted later on
+            }
+        }
+    }
+
     // pkg private for testing
     final void verifyAfterCleanup(MetadataSnapshot sourceMetadata, MetadataSnapshot targetMetadata) {
         final RecoveryDiff recoveryDiff = targetMetadata.recoveryDiff(sourceMetadata);

server/src/main/java/org/opensearch/indices/IndicesModule.java

@@ -284,6 +284,8 @@ protected void configure() {
         bind(RetentionLeaseSyncer.class).asEagerSingleton();
         if (FeatureFlags.isEnabled(FeatureFlags.REPLICATION_TYPE)) {
             bind(SegmentReplicationCheckpointPublisher.class).asEagerSingleton();
+        } else {
+            bind(SegmentReplicationCheckpointPublisher.class).toInstance(SegmentReplicationCheckpointPublisher.EMPTY);
         }
     }
 

server/src/main/java/org/opensearch/indices/RunUnderPrimaryPermit.java

@@ -0,0 +1,72 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.indices;
+
+import org.apache.logging.log4j.Logger;
+import org.opensearch.action.ActionListener;
+import org.opensearch.common.lease.Releasable;
+import org.opensearch.common.util.CancellableThreads;
+import org.opensearch.common.util.concurrent.FutureUtils;
+import org.opensearch.index.shard.IndexShard;
+import org.opensearch.index.shard.IndexShardRelocatedException;
+import org.opensearch.threadpool.ThreadPool;
+
+import java.util.concurrent.CompletableFuture;
+
+/**
+ * Execute a Runnable after acquiring the primary's operation permit.
+ *
+ * @opensearch.internal
+ */
+public final class RunUnderPrimaryPermit {
+
+    public static void run(
+        CancellableThreads.Interruptible runnable,
+        String reason,
+        IndexShard primary,
+        CancellableThreads cancellableThreads,
+        Logger logger
+    ) {
+        cancellableThreads.execute(() -> {
+            CompletableFuture<Releasable> permit = new CompletableFuture<>();
+            final ActionListener<Releasable> onAcquired = new ActionListener<>() {
+                @Override
+                public void onResponse(Releasable releasable) {
+                    if (permit.complete(releasable) == false) {
+                        releasable.close();
+                    }
+                }
+
+                @Override
+                public void onFailure(Exception e) {
+                    permit.completeExceptionally(e);
+                }
+            };
+            primary.acquirePrimaryOperationPermit(onAcquired, ThreadPool.Names.SAME, reason);
+            try (Releasable ignored = FutureUtils.get(permit)) {
+                // check that the IndexShard still has the primary authority. This needs to be checked under operation permit to prevent
+                // races, as IndexShard will switch its authority only when it holds all operation permits, see IndexShard.relocated()
+                if (primary.isRelocatedPrimary()) {
+                    throw new IndexShardRelocatedException(primary.shardId());
+                }
+                runnable.run();
+            } finally {
+                // just in case we got an exception (likely interrupted) while waiting for the get
+                permit.whenComplete((r, e) -> {
+                    if (r != null) {
+                        r.close();
+                    }
+                    if (e != null) {
+                        logger.trace("suppressing exception on completion (it was already bubbled up or the operation was aborted)", e);
+                    }
+                });
+            }
+        });
+    }
+}

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

@@ -0,0 +1,31 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.indices.recovery;
+
+import org.opensearch.action.ActionListener;
+import org.opensearch.common.bytes.BytesReference;
+import org.opensearch.index.store.StoreFileMetadata;
+
+/**
+ * Writes a partial file chunk to the target store.
+ *
+ * @opensearch.internal
+ */
+@FunctionalInterface
+public interface FileChunkWriter {
+
+    void writeFileChunk(
+        StoreFileMetadata fileMetadata,
+        long position,
+        BytesReference content,
+        boolean lastChunk,
+        int totalTranslogOps,
+        ActionListener<Void> listener
+    );
+}