[Segment Replication] Added source-side classes for orchestrating replication events

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

@@ -175,6 +175,19 @@ public final EngineConfig config() {
      */
     protected abstract SegmentInfos getLatestSegmentInfos();
 
+    /**
+     * Fetch a snapshot of the latest SegmentInfos from the engine. Using this method
+     * ensures that segment files are retained in the directory until the reference is closed.
+     *
+     * @return {@link GatedCloseable} - A wrapper around a {@link SegmentInfos} instance that
+     * must be closed for segment files to be deleted.
+     * @throws EngineException - When segment infos cannot be safely retrieved
+     */
+    public GatedCloseable<SegmentInfos> getSegmentInfosSnapshot() {
+        // default implementation
+        return new GatedCloseable<>(getLatestSegmentInfos(), () -> {});
+    }
+
     public MergeStats getMergeStats() {
         return new MergeStats();
     }
@@ -846,6 +859,15 @@ public final CommitStats commitStats() {
      */
     public abstract long getPersistedLocalCheckpoint();
 
+    /**
+     * @return the latest checkpoint that has been processed but not necessarily persisted.
+     * Also see {@link #getPersistedLocalCheckpoint()}
+     */
+    public long getProcessedLocalCheckpoint() {
+        // default implementation
+        return 0L;
+    }
+
     /**
      * @return a {@link SeqNoStats} object, using local state and the supplied global checkpoint
      */

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

@@ -2305,6 +2305,17 @@ public SegmentInfos getLatestSegmentInfos() {
         }
     }
 
+    @Override
+    public GatedCloseable<SegmentInfos> getSegmentInfosSnapshot() {
+        final SegmentInfos segmentInfos = getLatestSegmentInfos();
+        try {
+            indexWriter.incRefDeleter(segmentInfos);
+        } catch (IOException e) {
+            throw new EngineException(shardId, e.getMessage(), e);
+        }
+        return new GatedCloseable<>(segmentInfos, () -> indexWriter.decRefDeleter(segmentInfos));
+    }
+
     @Override
     protected final void writerSegmentStats(SegmentsStats stats) {
         stats.addVersionMapMemoryInBytes(versionMap.ramBytesUsed());
@@ -2724,6 +2735,7 @@ public long getLastSyncedGlobalCheckpoint() {
         return getTranslog().getLastSyncedGlobalCheckpoint();
     }
 
+    @Override
     public long getProcessedLocalCheckpoint() {
         return localCheckpointTracker.getProcessedCheckpoint();
     }

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

@@ -2638,6 +2638,14 @@ public long getLocalCheckpoint() {
         return getEngine().getPersistedLocalCheckpoint();
     }
 
+    /**
+     * Fetch the latest checkpoint that has been processed but not necessarily persisted.
+     * Also see {@link #getLocalCheckpoint()}.
+     */
+    public long getProcessedLocalCheckpoint() {
+        return getEngine().getProcessedLocalCheckpoint();
+    }
+
     /**
      * Returns the global checkpoint for the shard.
      *
@@ -4005,4 +4013,14 @@ public void verifyShardBeforeIndexClosing() throws IllegalStateException {
     RetentionLeaseSyncer getRetentionLeaseSyncer() {
         return retentionLeaseSyncer;
     }
+
+    /**
+     * Fetch the latest SegmentInfos held by the shard's underlying Engine, wrapped
+     * by a a {@link GatedCloseable} to ensure files are not deleted/merged away.
+     *
+     * @throws EngineException - When segment infos cannot be safely retrieved
+     */
+    public GatedCloseable<SegmentInfos> getSegmentInfosSnapshot() {
+        return getEngine().getSegmentInfosSnapshot();
+    }
 }

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

@@ -274,6 +274,13 @@ public MetadataSnapshot getMetadata(IndexCommit commit) throws IOException {
         return getMetadata(commit, false);
     }
 
+    /**
+     * Conveience wrapper around the {@link #getMetadata(IndexCommit)} method for null input.
+     */
+    public MetadataSnapshot getMetadata() throws IOException {
+        return getMetadata(null, false);
+    }
+
     /**
      * Returns a new MetadataSnapshot for the given commit. If the given commit is <code>null</code>
      * the latest commit point is used.
@@ -315,6 +322,10 @@ public MetadataSnapshot getMetadata(IndexCommit commit, boolean lockDirectory) t
         }
     }
 
+    public MetadataSnapshot getMetadata(SegmentInfos segmentInfos) throws IOException {
+        return new MetadataSnapshot(segmentInfos, directory, logger);
+    }
+
     /**
      * Renames all the given files from the key of the map to the
      * value of the map. All successfully renamed files are removed from the map in-place.
@@ -477,7 +488,7 @@ public static MetadataSnapshot readMetadataSnapshot(
             Directory dir = new NIOFSDirectory(indexLocation)
         ) {
             failIfCorrupted(dir);
-            return new MetadataSnapshot(null, dir, logger);
+            return new MetadataSnapshot((IndexCommit) null, dir, logger);
         } catch (IndexNotFoundException ex) {
             // that's fine - happens all the time no need to log
         } catch (FileNotFoundException | NoSuchFileException ex) {
@@ -682,7 +693,7 @@ public void cleanupAndVerify(String reason, MetadataSnapshot sourceMetadata) thr
                 }
             }
             directory.syncMetaData();
-            final Store.MetadataSnapshot metadataOrEmpty = getMetadata(null);
+            final Store.MetadataSnapshot metadataOrEmpty = getMetadata();
             verifyAfterCleanup(sourceMetadata, metadataOrEmpty);
         } finally {
             metadataLock.writeLock().unlock();
@@ -822,7 +833,14 @@ public MetadataSnapshot(Map<String, StoreFileMetadata> metadata, Map<String, Str
         }
 
         MetadataSnapshot(IndexCommit commit, Directory directory, Logger logger) throws IOException {
-            LoadedMetadata loadedMetadata = loadMetadata(commit, directory, logger);
+            this(loadMetadata(commit, directory, logger));
+        }
+
+        MetadataSnapshot(SegmentInfos segmentInfos, Directory directory, Logger logger) throws IOException {
+            this(loadMetadata(segmentInfos, directory, logger));
+        }
+
+        private MetadataSnapshot(LoadedMetadata loadedMetadata) {
             metadata = loadedMetadata.fileMetadata;
             commitUserData = loadedMetadata.userData;
             numDocs = loadedMetadata.numDocs;
@@ -890,40 +908,9 @@ static class LoadedMetadata {
         }
 
         static LoadedMetadata loadMetadata(IndexCommit commit, Directory directory, Logger logger) throws IOException {
-            long numDocs;
-            Map<String, StoreFileMetadata> builder = new HashMap<>();
-            Map<String, String> commitUserDataBuilder = new HashMap<>();
             try {
                 final SegmentInfos segmentCommitInfos = Store.readSegmentsInfo(commit, directory);
-                numDocs = Lucene.getNumDocs(segmentCommitInfos);
-                commitUserDataBuilder.putAll(segmentCommitInfos.getUserData());
-                // we don't know which version was used to write so we take the max version.
-                Version maxVersion = segmentCommitInfos.getMinSegmentLuceneVersion();
-                for (SegmentCommitInfo info : segmentCommitInfos) {
-                    final Version version = info.info.getVersion();
-                    if (version == null) {
-                        // version is written since 3.1+: we should have already hit IndexFormatTooOld.
-                        throw new IllegalArgumentException("expected valid version value: " + info.info.toString());
-                    }
-                    if (version.onOrAfter(maxVersion)) {
-                        maxVersion = version;
-                    }
-                    for (String file : info.files()) {
-                        checksumFromLuceneFile(
-                            directory,
-                            file,
-                            builder,
-                            logger,
-                            version,
-                            SEGMENT_INFO_EXTENSION.equals(IndexFileNames.getExtension(file))
-                        );
-                    }
-                }
-                if (maxVersion == null) {
-                    maxVersion = org.opensearch.Version.CURRENT.minimumIndexCompatibilityVersion().luceneVersion;
-                }
-                final String segmentsFile = segmentCommitInfos.getSegmentsFileName();
-                checksumFromLuceneFile(directory, segmentsFile, builder, logger, maxVersion, true);
+                return loadMetadata(segmentCommitInfos, directory, logger);
             } catch (CorruptIndexException | IndexNotFoundException | IndexFormatTooOldException | IndexFormatTooNewException ex) {
                 // we either know the index is corrupted or it's just not there
                 throw ex;
@@ -949,6 +936,40 @@ static LoadedMetadata loadMetadata(IndexCommit commit, Directory directory, Logg
                 }
                 throw ex;
             }
+        }
+
+        static LoadedMetadata loadMetadata(SegmentInfos segmentInfos, Directory directory, Logger logger) throws IOException {
+            long numDocs = Lucene.getNumDocs(segmentInfos);
+            Map<String, String> commitUserDataBuilder = new HashMap<>();
+            commitUserDataBuilder.putAll(segmentInfos.getUserData());
+            Map<String, StoreFileMetadata> builder = new HashMap<>();
+            // we don't know which version was used to write so we take the max version.
+            Version maxVersion = segmentInfos.getMinSegmentLuceneVersion();
+            for (SegmentCommitInfo info : segmentInfos) {
+                final Version version = info.info.getVersion();
+                if (version == null) {
+                    // version is written since 3.1+: we should have already hit IndexFormatTooOld.
+                    throw new IllegalArgumentException("expected valid version value: " + info.info.toString());
+                }
+                if (version.onOrAfter(maxVersion)) {
+                    maxVersion = version;
+                }
+                for (String file : info.files()) {
+                    checksumFromLuceneFile(
+                        directory,
+                        file,
+                        builder,
+                        logger,
+                        version,
+                        SEGMENT_INFO_EXTENSION.equals(IndexFileNames.getExtension(file))
+                    );
+                }
+            }
+            if (maxVersion == null) {
+                maxVersion = org.opensearch.Version.CURRENT.minimumIndexCompatibilityVersion().luceneVersion;
+            }
+            final String segmentsFile = segmentInfos.getSegmentsFileName();
+            checksumFromLuceneFile(directory, segmentsFile, builder, logger, maxVersion, true);
             return new LoadedMetadata(unmodifiableMap(builder), unmodifiableMap(commitUserDataBuilder), numDocs);
         }
 

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

@@ -0,0 +1,54 @@
+/*
+ * 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.replication;
+
+import org.opensearch.cluster.node.DiscoveryNode;
+import org.opensearch.common.io.stream.StreamInput;
+import org.opensearch.common.io.stream.StreamOutput;
+import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint;
+import org.opensearch.indices.replication.common.SegmentReplicationTransportRequest;
+
+import java.io.IOException;
+
+/**
+ * Request object for fetching segment metadata for a {@link ReplicationCheckpoint} from
+ * a {@link SegmentReplicationSource}. This object is created by the target node and sent
+ * to the source node.
+ *
+ * @opensearch.internal
+ */
+public class CheckpointInfoRequest extends SegmentReplicationTransportRequest {
+
+    private final ReplicationCheckpoint checkpoint;
+
+    public CheckpointInfoRequest(StreamInput in) throws IOException {
+        super(in);
+        checkpoint = new ReplicationCheckpoint(in);
+    }
+
+    public CheckpointInfoRequest(
+        long replicationId,
+        String targetAllocationId,
+        DiscoveryNode targetNode,
+        ReplicationCheckpoint checkpoint
+    ) {
+        super(replicationId, targetAllocationId, targetNode);
+        this.checkpoint = checkpoint;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        super.writeTo(out);
+        checkpoint.writeTo(out);
+    }
+
+    public ReplicationCheckpoint getCheckpoint() {
+        return checkpoint;
+    }
+}

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

@@ -0,0 +1,60 @@
+/*
+ * 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.replication;
+
+import org.opensearch.cluster.node.DiscoveryNode;
+import org.opensearch.common.io.stream.StreamInput;
+import org.opensearch.common.io.stream.StreamOutput;
+import org.opensearch.index.store.StoreFileMetadata;
+import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint;
+import org.opensearch.indices.replication.common.SegmentReplicationTransportRequest;
+
+import java.io.IOException;
+import java.util.List;
+
+/**
+ * Request object for fetching a list of segment files metadata from a {@link SegmentReplicationSource}.
+ * This object is created by the target node and sent to the source node.
+ *
+ * @opensearch.internal
+ */
+public class GetSegmentFilesRequest extends SegmentReplicationTransportRequest {
+
+    private final List<StoreFileMetadata> filesToFetch;
+    private final ReplicationCheckpoint checkpoint;
+
+    public GetSegmentFilesRequest(StreamInput in) throws IOException {
+        super(in);
+        this.filesToFetch = in.readList(StoreFileMetadata::new);
+        this.checkpoint = new ReplicationCheckpoint(in);
+    }
+
+    public GetSegmentFilesRequest(
+        long replicationId,
+        String targetAllocationId,
+        DiscoveryNode targetNode,
+        List<StoreFileMetadata> filesToFetch,
+        ReplicationCheckpoint checkpoint
+    ) {
+        super(replicationId, targetAllocationId, targetNode);
+        this.filesToFetch = filesToFetch;
+        this.checkpoint = checkpoint;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        super.writeTo(out);
+        out.writeList(filesToFetch);
+        checkpoint.writeTo(out);
+    }
+
+    public ReplicationCheckpoint getCheckpoint() {
+        return checkpoint;
+    }
+}