Remote merged segment warmer

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

@@ -59,7 +59,7 @@ public MergedSegmentWarmerFactory(TransportService transportService, RecoverySet
 
     public IndexWriter.IndexReaderWarmer get(IndexShard shard) {
         if (shard.indexSettings().isAssignedOnRemoteNode()) {
-            return new RemoteStoreMergedSegmentWarmer(transportService, recoverySettings, clusterService);
+            return new RemoteStoreMergedSegmentWarmer(transportService, recoverySettings, clusterService, shard);
         } else if (shard.indexSettings().isSegRepLocalEnabled()) {
             return new LocalMergedSegmentWarmer(transportService, recoverySettings, clusterService, shard);
         } else if (shard.indexSettings().isDocumentReplication()) {

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

@@ -32,9 +32,14 @@
 
 package org.opensearch.index.engine;
 
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import org.apache.lucene.index.IndexWriter;
 import org.apache.lucene.index.LeafReader;
+import org.apache.lucene.index.SegmentCommitInfo;
+import org.apache.lucene.index.SegmentReader;
 import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.index.shard.IndexShard;
 import org.opensearch.indices.recovery.RecoverySettings;
 import org.opensearch.transport.TransportService;
 
@@ -49,19 +54,33 @@ public class RemoteStoreMergedSegmentWarmer implements IndexWriter.IndexReaderWa
     private final TransportService transportService;
     private final RecoverySettings recoverySettings;
     private final ClusterService clusterService;
+    private final IndexShard indexShard;
+
+    private final Logger logger = LogManager.getLogger(RemoteStoreMergedSegmentWarmer.class);
 
     public RemoteStoreMergedSegmentWarmer(
         TransportService transportService,
         RecoverySettings recoverySettings,
-        ClusterService clusterService
+        ClusterService clusterService,
+        IndexShard indexShard
     ) {
         this.transportService = transportService;
         this.recoverySettings = recoverySettings;
         this.clusterService = clusterService;
+        this.indexShard = indexShard;
     }
 
     @Override
     public void warm(LeafReader leafReader) throws IOException {
-        // TODO: remote store merged segment warmer
+        // IndexWriter.IndexReaderWarmer#warm is called by IndexWriter#mergeMiddle. The type of leafReader should be SegmentReader.
+        assert leafReader instanceof SegmentReader;
+
+        SegmentCommitInfo segmentCommitInfo = ((SegmentReader) leafReader).getSegmentInfo();
+        if(logger.isTraceEnabled()) {
+            logger.trace("[ShardId {}] Warming segment: {}", indexShard.shardId(), segmentCommitInfo);
+        }
+        indexShard.publishMergedSegment(segmentCommitInfo);
+        logger.trace("Completed segment warming for {} on shard {}",
+            segmentCommitInfo.info.name, indexShard.shardId());
     }
 }

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

@@ -201,10 +201,7 @@
 import org.opensearch.indices.recovery.RecoverySettings;
 import org.opensearch.indices.recovery.RecoveryState;
 import org.opensearch.indices.recovery.RecoveryTarget;
-import org.opensearch.indices.replication.checkpoint.MergeSegmentCheckpoint;
-import org.opensearch.indices.replication.checkpoint.MergedSegmentPublisher;
-import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint;
-import org.opensearch.indices.replication.checkpoint.SegmentReplicationCheckpointPublisher;
+import org.opensearch.indices.replication.checkpoint.*;
 import org.opensearch.indices.replication.common.ReplicationTimer;
 import org.opensearch.repositories.RepositoriesService;
 import org.opensearch.repositories.Repository;
@@ -1864,19 +1861,30 @@ public void publishMergedSegment(SegmentCommitInfo segmentCommitInfo) throws IOE
     }
 
     /**
-     * Compute {@link MergeSegmentCheckpoint} from a SegmentCommitInfo.
+     * Compute {@link MergedSegmentCheckpoint} from a SegmentCommitInfo.
      * This function fetches a metadata snapshot from the store that comes with an IO cost.
      *
      * @param segmentCommitInfo {@link SegmentCommitInfo} segmentCommitInfo to use to compute.
-     * @return {@link MergeSegmentCheckpoint} Checkpoint computed from the segmentCommitInfo.
+     * @return {@link MergedSegmentCheckpoint} Checkpoint computed from the segmentCommitInfo.
      * @throws IOException When there is an error computing segment metadata from the store.
      */
-    public MergeSegmentCheckpoint computeMergeSegmentCheckpoint(SegmentCommitInfo segmentCommitInfo) throws IOException {
+    public ReplicationCheckpoint computeMergeSegmentCheckpoint(SegmentCommitInfo segmentCommitInfo) throws IOException {
         // Only need to get the file metadata information in segmentCommitInfo and reuse Store#getSegmentMetadataMap.
         SegmentInfos segmentInfos = new SegmentInfos(Version.LATEST.major);
         segmentInfos.add(segmentCommitInfo);
         Map<String, StoreFileMetadata> segmentMetadataMap = store.getSegmentMetadataMap(segmentInfos);
-        return new MergeSegmentCheckpoint(
+        if (indexSettings.isRemoteStoreEnabled()) {
+            return new RemoteStoreMergedSegmentCheckpoint(
+                shardId,
+                getOperationPrimaryTerm(),
+                segmentMetadataMap.values().stream().mapToLong(StoreFileMetadata::length).sum(),
+                getEngine().config().getCodec().getName(),
+                segmentMetadataMap,
+                segmentCommitInfo.info.name,
+                null
+            );
+        }
+        return new MergedSegmentCheckpoint(
             shardId,
             getOperationPrimaryTerm(),
             segmentMetadataMap.values().stream().mapToLong(StoreFileMetadata::length).sum(),

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

@@ -35,6 +35,7 @@
 import org.opensearch.index.store.remote.metadata.RemoteSegmentMetadata;
 import org.opensearch.index.translog.Translog;
 import org.opensearch.indices.RemoteStoreSettings;
+import org.opensearch.indices.replication.ActiveMergesSegmentRegistry;
 import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint;
 import org.opensearch.indices.replication.checkpoint.SegmentReplicationCheckpointPublisher;
 import org.opensearch.threadpool.ThreadPool;
@@ -84,6 +85,7 @@ public final class RemoteStoreRefreshListener extends ReleasableRetryableRefresh
 
     public static final Set<String> EXCLUDE_FILES = Set.of("write.lock");
 
+    private final ActiveMergesSegmentRegistry activeMergesSegmentRegistry = ActiveMergesSegmentRegistry.getInstance();
     private final IndexShard indexShard;
     private final Directory storeDirectory;
     private final RemoteSegmentStoreDirectory remoteDirectory;
@@ -258,6 +260,8 @@ private boolean syncSegments() {
                     long lastRefreshedCheckpoint = ((InternalEngine) indexShard.getEngine()).lastRefreshedCheckpoint();
                     Collection<String> localSegmentsPostRefresh = segmentInfos.files(true);
 
+                    remoteDirectory.syncSegmentsUploadedToRemoteStoreWithActiveMergesSegmentRegistry(storeDirectory, localSegmentsPostRefresh);
+
                     // Create a map of file name to size and update the refresh segment tracker
                     Map<String, Long> localSegmentsSizeMap = updateLocalSizeMapAndTracker(localSegmentsPostRefresh).entrySet()
                         .stream()

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

@@ -0,0 +1,29 @@
+/*
+ * 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.index.shard;
+
+import org.opensearch.common.util.UploadListener;
+import org.opensearch.core.action.ActionListener;
+
+import java.util.Collection;
+import java.util.Map;
+import java.util.function.Function;
+
+/**
+ * Interface to handle the functionality for upload data in the remote store
+ */
+public interface RemoteStoreUploader {
+
+    void uploadSegments(
+        Collection<String> localSegments,
+        Map<String, Long> localSegmentsSizeMap,
+        ActionListener<Void> listener,
+        Function<Map<String, Long>, UploadListener> uploadListenerFunction
+    );
+}

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

@@ -0,0 +1,104 @@
+/*
+ * 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.index.shard;
+
+import org.apache.logging.log4j.Logger;
+import org.apache.logging.log4j.message.ParameterizedMessage;
+import org.apache.lucene.index.CorruptIndexException;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.store.FilterDirectory;
+import org.apache.lucene.store.IOContext;
+import org.opensearch.action.support.GroupedActionListener;
+import org.opensearch.cluster.routing.RecoverySource;
+import org.opensearch.common.logging.Loggers;
+import org.opensearch.common.util.UploadListener;
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.index.store.CompositeDirectory;
+import org.opensearch.index.store.RemoteSegmentStoreDirectory;
+
+import java.util.Collection;
+import java.util.Map;
+import java.util.function.Function;
+
+/**
+ * The service essentially acts as a bridge between local segment storage and remote storage,
+ * ensuring efficient and reliable segment synchronization while providing comprehensive monitoring and error handling.
+ */
+public class RemoteStoreUploaderService implements RemoteStoreUploader {
+
+    private final Logger logger;
+
+    private final IndexShard indexShard;
+    private final Directory storeDirectory;
+    private final RemoteSegmentStoreDirectory remoteDirectory;
+
+    public RemoteStoreUploaderService(IndexShard indexShard, Directory storeDirectory, RemoteSegmentStoreDirectory remoteDirectory) {
+        logger = Loggers.getLogger(getClass(), indexShard.shardId());
+        this.indexShard = indexShard;
+        this.storeDirectory = storeDirectory;
+        this.remoteDirectory = remoteDirectory;
+    }
+
+    @Override
+    public void uploadSegments(
+        Collection<String> localSegments,
+        Map<String, Long> localSegmentsSizeMap,
+        ActionListener<Void> listener,
+        Function<Map<String, Long>, UploadListener> uploadListenerFunction
+    ) {
+        if (localSegments.isEmpty()) {
+            logger.debug("No new segments to upload in uploadNewSegments");
+            listener.onResponse(null);
+            return;
+        }
+
+        logger.debug("Effective new segments files to upload {}", localSegments);
+        ActionListener<Collection<Void>> mappedListener = ActionListener.map(listener, resp -> null);
+        GroupedActionListener<Void> batchUploadListener = new GroupedActionListener<>(mappedListener, localSegments.size());
+        Directory directory = ((FilterDirectory) (((FilterDirectory) storeDirectory).getDelegate())).getDelegate();
+
+        for (String localSegment : localSegments) {
+            // Initializing listener here to ensure that the stats increment operations are thread-safe
+            UploadListener statsListener = uploadListenerFunction.apply(localSegmentsSizeMap);
+            ActionListener<Void> aggregatedListener = ActionListener.wrap(resp -> {
+                statsListener.onSuccess(localSegment);
+                batchUploadListener.onResponse(resp);
+                // Once uploaded to Remote, local files become eligible for eviction from FileCache
+                if (directory instanceof CompositeDirectory) {
+                    ((CompositeDirectory) directory).afterSyncToRemote(localSegment);
+                }
+            }, ex -> {
+                logger.warn(() -> new ParameterizedMessage("Exception: [{}] while uploading segment files", ex), ex);
+                if (ex instanceof CorruptIndexException) {
+                    indexShard.failShard(ex.getMessage(), ex);
+                }
+                statsListener.onFailure(localSegment);
+                batchUploadListener.onFailure(ex);
+            });
+            statsListener.beforeUpload(localSegment);
+            // Place where the actual upload is happening
+            remoteDirectory.copyFrom(storeDirectory, localSegment, IOContext.DEFAULT, aggregatedListener, isLowPriorityUpload());
+        }
+    }
+
+    boolean isLowPriorityUpload() {
+        return isLocalOrSnapshotRecoveryOrSeeding();
+    }
+
+    boolean isLocalOrSnapshotRecoveryOrSeeding() {
+        // In this case when the primary mode is false, we need to upload segments to Remote Store
+        // This is required in case of remote migration seeding/snapshots/shrink/ split/clone where we need to durable persist
+        // all segments to remote before completing the recovery to ensure durability.
+        return (indexShard.state() == IndexShardState.RECOVERING && indexShard.shardRouting.primary())
+            && indexShard.recoveryState() != null
+            && (indexShard.recoveryState().getRecoverySource().getType() == RecoverySource.Type.LOCAL_SHARDS
+                || indexShard.recoveryState().getRecoverySource().getType() == RecoverySource.Type.SNAPSHOT
+                || indexShard.shouldSeedRemoteStore());
+    }
+}

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

@@ -32,6 +32,7 @@
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.unit.ByteSizeUnit;
 import org.opensearch.index.store.exception.ChecksumCombinationException;
+import org.opensearch.indices.replication.ActiveMergesSegmentRegistry;
 
 import java.io.FileNotFoundException;
 import java.io.IOException;
@@ -67,7 +68,7 @@ public class RemoteDirectory extends Directory {
 
     private final UnaryOperator<OffsetRangeInputStream> lowPriorityUploadRateLimiter;
 
-    private final UnaryOperator<InputStream> downloadRateLimiter;
+    private final DownloadRateLimiterProvider downloadRateLimiterProvider;
 
     /**
      * Number of bytes in the segment file to store checksum
@@ -79,19 +80,20 @@ public BlobContainer getBlobContainer() {
     }
 
     public RemoteDirectory(BlobContainer blobContainer) {
-        this(blobContainer, UnaryOperator.identity(), UnaryOperator.identity(), UnaryOperator.identity());
+        this(blobContainer, UnaryOperator.identity(), UnaryOperator.identity(), UnaryOperator.identity(), UnaryOperator.identity());
     }
 
     public RemoteDirectory(
         BlobContainer blobContainer,
         UnaryOperator<OffsetRangeInputStream> uploadRateLimiter,
         UnaryOperator<OffsetRangeInputStream> lowPriorityUploadRateLimiter,
-        UnaryOperator<InputStream> downloadRateLimiter
+        UnaryOperator<InputStream> downloadRateLimiter,
+        UnaryOperator<InputStream> lowPriorityDownloadRateLimiter
     ) {
         this.blobContainer = blobContainer;
         this.lowPriorityUploadRateLimiter = lowPriorityUploadRateLimiter;
         this.uploadRateLimiter = uploadRateLimiter;
-        this.downloadRateLimiter = downloadRateLimiter;
+        this.downloadRateLimiterProvider = new DownloadRateLimiterProvider(downloadRateLimiter, lowPriorityDownloadRateLimiter);
     }
 
     /**
@@ -236,7 +238,7 @@ public IndexInput openInput(String name, long fileLength, IOContext context) thr
         InputStream inputStream = null;
         try {
             inputStream = blobContainer.readBlob(name);
-            return new RemoteIndexInput(name, downloadRateLimiter.apply(inputStream), fileLength);
+            return new RemoteIndexInput(name, downloadRateLimiterProvider.get(name).apply(inputStream), fileLength);
         } catch (Exception e) {
             // In case the RemoteIndexInput creation fails, close the input stream to avoid file handler leak.
             if (inputStream != null) {
@@ -475,8 +477,27 @@ private IndexInput getBlockInput(String name, long position, long length, long f
         byte[] bytes;
         try (InputStream inputStream = blobContainer.readBlob(name, position, length)) {
             // TODO - Explore how we can buffer small chunks of data instead of having the whole 8MB block in memory
-            bytes = downloadRateLimiter.apply(inputStream).readAllBytes();
+            bytes = downloadRateLimiterProvider.get(name).apply(inputStream).readAllBytes();
         }
         return new ByteArrayIndexInput(name, bytes);
     }
+
+    private class DownloadRateLimiterProvider {
+        private final ActiveMergesSegmentRegistry activeMergesSegmentRegistry = ActiveMergesSegmentRegistry.getInstance();
+
+        private final UnaryOperator<InputStream> downloadRateLimiter;
+        private final UnaryOperator<InputStream> lowPriorityDownloadRateLimiter;
+
+        DownloadRateLimiterProvider(UnaryOperator<InputStream> downloadRateLimiter, UnaryOperator<InputStream> lowPriorityDownloadRateLimiter) {
+            this.downloadRateLimiter = downloadRateLimiter;
+            this.lowPriorityDownloadRateLimiter = lowPriorityDownloadRateLimiter;
+        }
+
+        public UnaryOperator<InputStream> get(final String filename){
+            if (activeMergesSegmentRegistry.contains(filename)) {
+                return lowPriorityDownloadRateLimiter;
+            }
+            return downloadRateLimiter;
+        }
+    }
 }