Add async segment file download support from remote store

Signed-off-by: Kunal Kotwani <[email protected]>

server/src/internalClusterTest/java/org/opensearch/remotestore/multipart/RemoteStoreMultipartIT.java

@@ -13,6 +13,7 @@
 import org.opensearch.cluster.node.DiscoveryNode;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.unit.TimeValue;
+import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.core.common.unit.ByteSizeUnit;
 import org.opensearch.plugins.Plugin;
 import org.opensearch.remotestore.RemoteStoreIT;
@@ -86,6 +87,15 @@ public RepositoryMetadata buildRepositoryMetadata(DiscoveryNode node, String nam
         } else {
             return super.buildRepositoryMetadata(node, name);
         }
+
+    }
+
+    @Override
+    protected Settings featureFlagSettings() {
+        return Settings.builder()
+            .put(super.featureFlagSettings())
+            .put(FeatureFlags.REMOTE_STORE_EXPERIMENTAL_SETTING.getKey(), Boolean.TRUE.toString())
+            .build();
     }
 
     public void testRateLimitedRemoteUploads() throws Exception {

server/src/main/java/org/opensearch/common/settings/FeatureFlagSettings.java

@@ -36,6 +36,7 @@ protected FeatureFlagSettings(
         new HashSet<>(
             Arrays.asList(
                 FeatureFlags.SEGMENT_REPLICATION_EXPERIMENTAL_SETTING,
+                FeatureFlags.REMOTE_STORE_EXPERIMENTAL_SETTING,
                 FeatureFlags.EXTENSIONS_SETTING,
                 FeatureFlags.IDENTITY_SETTING,
                 FeatureFlags.CONCURRENT_SEGMENT_SEARCH_SETTING,

server/src/main/java/org/opensearch/common/util/FeatureFlags.java

@@ -26,6 +26,11 @@ public class FeatureFlags {
     public static final String SEGMENT_REPLICATION_EXPERIMENTAL =
         "opensearch.experimental.feature.segment_replication_experimental.enabled";
 
+    /**
+     * Gates the visibility of the remote store experimental features that allows users to test unreleased beta features.
+     */
+    public static final String REMOTE_STORE_EXPERIMENTAL = "opensearch.experimental.feature.remote_store_experimental.enabled";
+
     /**
      * Gates the ability for Searchable Snapshots to read snapshots that are older than the
      * guaranteed backward compatibility for OpenSearch (one prior major version) on a best effort basis.
@@ -89,6 +94,12 @@ public static boolean isEnabled(String featureFlagName) {
         Property.NodeScope
     );
 
+    public static final Setting<Boolean> REMOTE_STORE_EXPERIMENTAL_SETTING = Setting.boolSetting(
+        REMOTE_STORE_EXPERIMENTAL,
+        false,
+        Property.NodeScope
+    );
+
     public static final Setting<Boolean> EXTENSIONS_SETTING = Setting.boolSetting(EXTENSIONS, false, Property.NodeScope);
 
     public static final Setting<Boolean> IDENTITY_SETTING = Setting.boolSetting(IDENTITY, false, Property.NodeScope);

server/src/main/java/org/opensearch/index/IndexService.java

@@ -479,7 +479,7 @@ public synchronized IndexShard createShard(
             Store remoteStore = null;
             if (this.indexSettings.isRemoteStoreEnabled()) {
                 Directory remoteDirectory = remoteDirectoryFactory.newDirectory(this.indexSettings, path);
-                remoteStore = new Store(shardId, this.indexSettings, remoteDirectory, lock, Store.OnClose.EMPTY);
+                remoteStore = new Store(shardId, this.indexSettings, remoteDirectory, lock, Store.OnClose.EMPTY, path);
             }
 
             Directory directory = directoryFactory.newDirectory(this.indexSettings, path);
@@ -488,7 +488,8 @@ public synchronized IndexShard createShard(
                 this.indexSettings,
                 directory,
                 lock,
-                new StoreCloseListener(shardId, () -> eventListener.onStoreClosed(shardId))
+                new StoreCloseListener(shardId, () -> eventListener.onStoreClosed(shardId)),
+                path
             );
             eventListener.onStoreCreated(shardId);
             indexShard = new IndexShard(

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

@@ -96,7 +96,6 @@
 import org.opensearch.common.util.concurrent.RunOnce;
 import org.opensearch.common.util.concurrent.ThreadContext;
 import org.opensearch.common.util.io.IOUtils;
-import org.opensearch.common.util.set.Sets;
 import org.opensearch.core.Assertions;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.unit.ByteSizeValue;
@@ -198,6 +197,7 @@
 import java.nio.channels.ClosedByInterruptException;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.NoSuchFileException;
+import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -4833,39 +4833,96 @@ private String copySegmentFiles(
         Map<String, RemoteSegmentStoreDirectory.UploadedSegmentMetadata> uploadedSegments,
         boolean overrideLocal
     ) throws IOException {
-        List<String> downloadedSegments = new ArrayList<>();
-        List<String> skippedSegments = new ArrayList<>();
+        Set<String> toDownloadSegments = new HashSet<>();
+        Set<String> skippedSegments = new HashSet<>();
         String segmentNFile = null;
+
         try {
-            Set<String> localSegmentFiles = Sets.newHashSet(storeDirectory.listAll());
             if (overrideLocal) {
-                for (String file : localSegmentFiles) {
+                for (String file : storeDirectory.listAll()) {
                     storeDirectory.deleteFile(file);
                 }
             }
+
             for (String file : uploadedSegments.keySet()) {
                 long checksum = Long.parseLong(uploadedSegments.get(file).getChecksum());
                 if (overrideLocal || localDirectoryContains(storeDirectory, file, checksum) == false) {
-                    storeDirectory.copyFrom(sourceRemoteDirectory, file, file, IOContext.DEFAULT);
-                    downloadedSegments.add(file);
+                    toDownloadSegments.add(file);
                 } else {
                     skippedSegments.add(file);
                 }
-                if (targetRemoteDirectory != null) {
-                    targetRemoteDirectory.copyFrom(storeDirectory, file, file, IOContext.DEFAULT);
-                }
+
                 if (file.startsWith(IndexFileNames.SEGMENTS)) {
                     assert segmentNFile == null : "There should be only one SegmentInfosSnapshot file";
                     segmentNFile = file;
                 }
             }
+
+            if (toDownloadSegments.isEmpty() == false) {
+                try {
+                    final CountDownLatch completionLatch = new CountDownLatch(toDownloadSegments.size());
+                    final AtomicBoolean anyFileFailed = new AtomicBoolean();
+
+                    downloadSegments(
+                        storeDirectory,
+                        sourceRemoteDirectory,
+                        targetRemoteDirectory,
+                        toDownloadSegments,
+                        completionLatch,
+                        anyFileFailed
+                    );
+
+                    completionLatch.await();
+                    if (anyFileFailed.get()) {
+                        throw new IOException("Error occurred in one or more segment file downloads");
+                    }
+                } catch (Exception e) {
+                    throw new IOException("Error occurred when downloading segments from remote store", e);
+                }
+            }
         } finally {
-            logger.trace("Downloaded segments here: {}", downloadedSegments);
+            logger.trace("Downloaded segments here: {}", toDownloadSegments);
             logger.trace("Skipped download for segments here: {}", skippedSegments);
         }
+
         return segmentNFile;
     }
 
+    private void downloadSegments(
+        Directory storeDirectory,
+        RemoteSegmentStoreDirectory sourceRemoteDirectory,
+        RemoteSegmentStoreDirectory targetRemoteDirectory,
+        Set<String> toDownloadSegments,
+        CountDownLatch completionLatch,
+        AtomicBoolean anyFileFailed
+    ) {
+        final Path indexPath = store.shardPath() == null ? null : store.shardPath().resolveIndex();
+
+        final ActionListener<String> segmentsDownloadListener = new ActionListener<>() {
+            @Override
+            public void onResponse(String fileName) {
+                try {
+                    if (targetRemoteDirectory != null) {
+                        targetRemoteDirectory.copyFrom(storeDirectory, fileName, fileName, IOContext.DEFAULT);
+                    }
+                    completionLatch.countDown();
+                } catch (IOException ex) {
+                    logger.error("Failed to download segment file {} from the remote store", fileName);
+                    onFailure(ex);
+                }
+            }
+
+            @Override
+            public void onFailure(Exception e) {
+                logger.error("Failed to download one of the segment files from the remote store", e);
+                anyFileFailed.set(true);
+                completionLatch.countDown();
+            }
+        };
+
+        toDownloadSegments.forEach(file -> sourceRemoteDirectory.copyTo(file, storeDirectory, indexPath, segmentsDownloadListener));
+    }
+
     private boolean localDirectoryContains(Directory localDirectory, String file, long checksum) {
         try (IndexInput indexInput = localDirectory.openInput(file, IOContext.DEFAULT)) {
             if (checksum == CodecUtil.retrieveChecksum(indexInput)) {

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

@@ -24,8 +24,10 @@
 import org.apache.lucene.store.IndexOutput;
 import org.apache.lucene.util.Version;
 import org.opensearch.common.UUIDs;
+import org.opensearch.common.blobstore.AsyncMultiStreamBlobContainer;
 import org.opensearch.common.io.VersionedCodecStreamWrapper;
 import org.opensearch.common.lucene.store.ByteArrayIndexInput;
+import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.index.remote.RemoteStoreUtils;
 import org.opensearch.index.store.lockmanager.FileLockInfo;
@@ -40,6 +42,7 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.file.NoSuchFileException;
+import java.nio.file.Path;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
@@ -434,6 +437,43 @@ public void copyFrom(Directory from, String src, IOContext context, ActionListen
         }
     }
 
+    /**
+     * Copies an existing {@code source} file from this directory to a non-existent file (also
+     * named {@code source}) in either {@code destinationDirectory} or {@code destinationPath}.
+     * If the blob container backing this directory supports multipart downloads, the {@code source}
+     * file will be downloaded (potentially in multiple concurrent parts) directly to
+     * {@code destinationPath}. This method will return immediately and {@code fileCompletionListener}
+     * will be notified upon completion.
+     * <p>
+     * If multipart downloads are not supported, then {@code source} file will be copied to a file named
+     * {@code source} in a single part to {@code destinationDirectory}. The download will happen on the
+     * calling thread and {@code fileCompletionListener} will be notified synchronously before this
+     * method returns.
+     *
+     * @param source The source file name
+     * @param destinationDirectory The destination directory (if multipart is not supported)
+     * @param destinationPath The destination path (if multipart is supported)
+     * @param fileCompletionListener The listener to notify of completion
+     */
+    public void copyTo(String source, Directory destinationDirectory, Path destinationPath, ActionListener<String> fileCompletionListener) {
+        final String blobName = getExistingRemoteFilename(source);
+        if (FeatureFlags.isEnabled(FeatureFlags.REMOTE_STORE_EXPERIMENTAL)
+            && destinationPath != null
+            && remoteDataDirectory.getBlobContainer() instanceof AsyncMultiStreamBlobContainer) {
+            final AsyncMultiStreamBlobContainer blobContainer = (AsyncMultiStreamBlobContainer) remoteDataDirectory.getBlobContainer();
+            final Path destinationFilePath = destinationPath.resolve(source);
+            blobContainer.asyncBlobDownload(blobName, destinationFilePath, threadPool, fileCompletionListener);
+        } else {
+            // Fallback to older mechanism of downloading the file
+            try {
+                destinationDirectory.copyFrom(this, source, source, IOContext.DEFAULT);
+                fileCompletionListener.onResponse(source);
+            } catch (IOException e) {
+                fileCompletionListener.onFailure(e);
+            }
+        }
+    }
+
     /**
      * This acquires a lock on a given commit by creating a lock file in lock directory using {@code FileLockInfo}
      *

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

@@ -66,6 +66,7 @@
 import org.apache.lucene.util.Version;
 import org.opensearch.ExceptionsHelper;
 import org.opensearch.common.UUIDs;
+import org.opensearch.common.annotation.InternalApi;
 import org.opensearch.common.io.stream.BytesStreamOutput;
 import org.opensearch.common.logging.Loggers;
 import org.opensearch.common.lucene.Lucene;
@@ -92,6 +93,7 @@
 import org.opensearch.index.seqno.SequenceNumbers;
 import org.opensearch.index.shard.AbstractIndexShardComponent;
 import org.opensearch.index.shard.IndexShard;
+import org.opensearch.index.shard.ShardPath;
 import org.opensearch.index.translog.Translog;
 
 import java.io.Closeable;
@@ -179,6 +181,7 @@ public class Store extends AbstractIndexShardComponent implements Closeable, Ref
     private final ReentrantReadWriteLock metadataLock = new ReentrantReadWriteLock();
     private final ShardLock shardLock;
     private final OnClose onClose;
+    private final ShardPath shardPath;
 
     // used to ref count files when a new Reader is opened for PIT/Scroll queries
     // prevents segment files deletion until the PIT/Scroll expires or is discarded
@@ -192,17 +195,25 @@ protected void closeInternal() {
     };
 
     public Store(ShardId shardId, IndexSettings indexSettings, Directory directory, ShardLock shardLock) {
-        this(shardId, indexSettings, directory, shardLock, OnClose.EMPTY);
+        this(shardId, indexSettings, directory, shardLock, OnClose.EMPTY, null);
     }
 
-    public Store(ShardId shardId, IndexSettings indexSettings, Directory directory, ShardLock shardLock, OnClose onClose) {
+    public Store(
+        ShardId shardId,
+        IndexSettings indexSettings,
+        Directory directory,
+        ShardLock shardLock,
+        OnClose onClose,
+        ShardPath shardPath
+    ) {
         super(shardId, indexSettings);
         final TimeValue refreshInterval = indexSettings.getValue(INDEX_STORE_STATS_REFRESH_INTERVAL_SETTING);
         logger.debug("store stats are refreshed with refresh_interval [{}]", refreshInterval);
         ByteSizeCachingDirectory sizeCachingDir = new ByteSizeCachingDirectory(directory, refreshInterval);
         this.directory = new StoreDirectory(sizeCachingDir, Loggers.getLogger("index.store.deletes", shardId));
         this.shardLock = shardLock;
         this.onClose = onClose;
+        this.shardPath = shardPath;
         assert onClose != null;
         assert shardLock != null;
         assert shardLock.getShardId().equals(shardId);
@@ -213,6 +224,11 @@ public Directory directory() {
         return directory;
     }
 
+    @InternalApi
+    public ShardPath shardPath() {
+        return shardPath;
+    }
+
     /**
      * Returns the last committed segments info for this store
      *