Adding changes to support async fetch blob for transfer manager

Signed-off-by: Ajay Kumar Movva <movvaam@amazon.com>

Author: Ajay Kumar Movva

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

@@ -712,7 +712,14 @@ protected void closeInternal() {
             if (FeatureFlags.isEnabled(FeatureFlags.WRITABLE_WARM_INDEX_SETTING) &&
             // TODO : Need to remove this check after support for hot indices is added in Composite Directory
                 this.indexSettings.isWarmIndex()) {
-                directory = compositeDirectoryFactory.newDirectory(this.indexSettings, path, directoryFactory, remoteDirectory, fileCache);
+                directory = compositeDirectoryFactory.newDirectory(
+                    this.indexSettings,
+                    path,
+                    directoryFactory,
+                    remoteDirectory,
+                    fileCache,
+                    threadPool
+                );
             } else {
                 directory = directoryFactory.newDirectory(this.indexSettings, path);
             }

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

@@ -28,6 +28,7 @@
 import org.opensearch.index.store.remote.filecache.FileCache.RestoredCachedIndexInput;
 import org.opensearch.index.store.remote.utils.FileTypeUtils;
 import org.opensearch.index.store.remote.utils.TransferManager;
+import org.opensearch.threadpool.ThreadPool;
 
 import java.io.FileNotFoundException;
 import java.io.IOException;
@@ -59,19 +60,21 @@ public class CompositeDirectory extends FilterDirectory {
     protected final RemoteSegmentStoreDirectory remoteDirectory;
     protected final FileCache fileCache;
     protected final TransferManager transferManager;
+    private final ThreadPool threadPool;
 
     /**
      * Constructor to initialise the composite directory
      * @param localDirectory corresponding to the local FSDirectory
      * @param remoteDirectory corresponding to the remote directory
      * @param fileCache used to cache the remote files locally
      */
-    public CompositeDirectory(Directory localDirectory, Directory remoteDirectory, FileCache fileCache) {
+    public CompositeDirectory(Directory localDirectory, Directory remoteDirectory, FileCache fileCache, ThreadPool threadPool) {
         super(localDirectory);
         validate(localDirectory, remoteDirectory, fileCache);
         this.localDirectory = (FSDirectory) localDirectory;
         this.remoteDirectory = (RemoteSegmentStoreDirectory) remoteDirectory;
         this.fileCache = fileCache;
+        this.threadPool = threadPool;
         transferManager = new TransferManager(
             (name, position, length) -> new InputStreamIndexInput(
                 CompositeDirectory.this.remoteDirectory.openBlockInput(name, position, length, IOContext.DEFAULT),

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

@@ -15,6 +15,7 @@
 import org.opensearch.index.shard.ShardPath;
 import org.opensearch.index.store.remote.filecache.FileCache;
 import org.opensearch.plugins.IndexStorePlugin;
+import org.opensearch.threadpool.ThreadPool;
 
 import java.io.IOException;
 
@@ -31,10 +32,11 @@ public Directory newDirectory(
         ShardPath shardPath,
         IndexStorePlugin.DirectoryFactory localDirectoryFactory,
         Directory remoteDirectory,
-        FileCache fileCache
+        FileCache fileCache,
+        ThreadPool threadPool
     ) throws IOException {
         logger.trace("Creating composite directory from core - Default CompositeDirectoryFactory");
         Directory localDirectory = localDirectoryFactory.newDirectory(indexSettings, shardPath);
-        return new CompositeDirectory(localDirectory, remoteDirectory, fileCache);
+        return new CompositeDirectory(localDirectory, remoteDirectory, fileCache, threadPool);
     }
 }

server/src/main/java/org/opensearch/index/store/remote/file/OnDemandBlockIndexInput.java

@@ -35,7 +35,7 @@
  *
  * @opensearch.internal
  */
-abstract class OnDemandBlockIndexInput extends IndexInput implements RandomAccessInput {
+public abstract class OnDemandBlockIndexInput extends IndexInput implements RandomAccessInput {
     private static final Logger logger = LogManager.getLogger(OnDemandBlockIndexInput.class);
 
     public static final String CLEANER_THREAD_NAME_PREFIX = "index-input-cleaner";
@@ -372,6 +372,10 @@ public static Builder builder() {
         return new Builder();
     }
 
+    /**
+     * Builder for {@link OnDemandBlockIndexInput}. The default block size is 8MiB
+     * (see {@link Builder#DEFAULT_BLOCK_SIZE_SHIFT}).
+     */
     public static class Builder {
         // Block size shift (default value is 23 == 2^23 == 8MiB)
         public static final int DEFAULT_BLOCK_SIZE_SHIFT = 23;
@@ -407,7 +411,7 @@ public Builder length(long length) {
             return this;
         }
 
-        Builder blockSizeShift(int blockSizeShift) {
+        public Builder blockSizeShift(int blockSizeShift) {
             assert blockSizeShift < 31 : "blockSizeShift must be < 31";
             this.blockSizeShift = blockSizeShift;
             this.blockSize = 1 << blockSizeShift;

server/src/main/java/org/opensearch/index/store/remote/file/OnDemandBlockSnapshotIndexInput.java

@@ -32,7 +32,7 @@ public class OnDemandBlockSnapshotIndexInput extends OnDemandBlockIndexInput {
     /**
      * Where this class fetches IndexInput parts from
      */
-    final TransferManager transferManager;
+    protected final TransferManager transferManager;
 
     /**
      * FileInfo contains snapshot metadata references for this IndexInput

server/src/main/java/org/opensearch/index/store/remote/filecache/CachedIndexInput.java

@@ -9,9 +9,12 @@
 package org.opensearch.index.store.remote.filecache;
 
 import org.apache.lucene.store.IndexInput;
+import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.common.annotation.PublicApi;
 
 import java.io.IOException;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executor;
 
 /**
  * Interface for an entry in the {@link FileCache} that can return an
@@ -29,6 +32,15 @@ public interface CachedIndexInput extends AutoCloseable {
      */
     IndexInput getIndexInput() throws IOException;
 
+    /**
+     *  Trigger and get the completable future of index input
+     *  @return CompletableFuture of IndexInput
+     */
+    @ExperimentalApi
+    default CompletableFuture<IndexInput> asyncLoadIndexInput(Executor executor) {
+        return null;
+    }
+
     /**
      * @return length in bytes
      */

server/src/main/java/org/opensearch/index/store/remote/utils/TransferManager.java

@@ -12,6 +12,7 @@
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.store.IOContext;
 import org.apache.lucene.store.IndexInput;
+import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.index.store.remote.filecache.CachedIndexInput;
 import org.opensearch.index.store.remote.filecache.FileCache;
 import org.opensearch.index.store.remote.filecache.FileCachedIndexInput;
@@ -28,6 +29,7 @@
 import java.security.PrivilegedExceptionAction;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
+import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
@@ -107,6 +109,35 @@ public IndexInput fetchBlob(BlobFetchRequest blobFetchRequest) throws IOExceptio
         }
     }
 
+    @ExperimentalApi
+    public CompletableFuture<IndexInput> fetchBlobAsync(BlobFetchRequest blobFetchRequest, Executor executor) throws IOException {
+        final Path key = blobFetchRequest.getFilePath();
+        logger.trace("Asynchronous fetchBlob called for {}", key.toString());
+        try {
+            CachedIndexInput cacheEntry = fileCache.compute(key, (path, cachedIndexInput) -> {
+                if (cachedIndexInput == null || cachedIndexInput.isClosed()) {
+                    logger.trace("Transfer Manager - IndexInput closed or not in cache");
+                    // Doesn't exist or is closed, either way create a new one
+                    return new DelayedCreationCachedIndexInput(fileCache, streamReader, blobFetchRequest);
+                } else {
+                    logger.trace("Transfer Manager - Required blob Already in cache: {}", blobFetchRequest.toString());
+                    // already in the cache and ready to be used (open)
+                    return cachedIndexInput;
+                }
+            });
+            // Cache entry was either retrieved from the cache or newly added, either
+            // way the reference count has been incremented by one. We can only
+            // decrement this reference _after_ creating the clone to be returned.
+            try {
+                return cacheEntry.asyncLoadIndexInput(executor);
+            } finally {
+                fileCache.decRef(key);
+            }
+        } catch (Exception cause) {
+            throw (RuntimeException) cause;
+        }
+    }
+
     private static FileCachedIndexInput createIndexInput(FileCache fileCache, StreamReader streamReader, BlobFetchRequest request) {
         try {
             // This local file cache is ref counted and may not strictly enforce configured capacity.
@@ -195,6 +226,32 @@ public IndexInput getIndexInput() throws IOException {
             }
         }
 
+        @ExperimentalApi
+        public CompletableFuture<IndexInput> asyncLoadIndexInput(Executor executor) {
+            if (isClosed.get()) {
+                return CompletableFuture.failedFuture(new IllegalStateException("Already closed"));
+            }
+            if (isStarted.getAndSet(true) == false) {
+                // Create new future and set it as the result
+                CompletableFuture.supplyAsync(() -> {
+                    try {
+                        return createIndexInput(fileCache, streamReader, request);
+                    } catch (Exception e) {
+                        fileCache.remove(request.getFilePath());
+                        throw new CompletionException(e);
+                    }
+                }, executor).handle((indexInput, throwable) -> {
+                    if (throwable != null) {
+                        result.completeExceptionally(throwable);
+                    } else {
+                        result.complete(indexInput);
+                    }
+                    return null;
+                });
+            }
+            return result;
+        }
+
         @Override
         public long length() {
             return request.getBlobLength();

server/src/main/java/org/opensearch/plugins/IndexStorePlugin.java

@@ -43,6 +43,7 @@
 import org.opensearch.index.store.IndexStoreListener;
 import org.opensearch.index.store.remote.filecache.FileCache;
 import org.opensearch.indices.recovery.RecoveryState;
+import org.opensearch.threadpool.ThreadPool;
 
 import java.io.IOException;
 import java.util.Collections;
@@ -104,7 +105,8 @@ Directory newDirectory(
             ShardPath shardPath,
             DirectoryFactory localDirectoryFactory,
             Directory remoteDirectory,
-            FileCache fileCache
+            FileCache fileCache,
+            ThreadPool threadPool
         ) throws IOException;
     }
 

server/src/test/java/org/opensearch/index/store/CompositeDirectoryTests.java

@@ -66,7 +66,7 @@ public void setup() throws IOException {
         localDirectory = FSDirectory.open(createTempDir());
         removeExtraFSFiles();
         fileCache = FileCacheFactory.createConcurrentLRUFileCache(FILE_CACHE_CAPACITY, new NoopCircuitBreaker(CircuitBreaker.REQUEST));
-        compositeDirectory = new CompositeDirectory(localDirectory, remoteSegmentStoreDirectory, fileCache);
+        compositeDirectory = new CompositeDirectory(localDirectory, remoteSegmentStoreDirectory, fileCache, threadPool);
         addFilesToDirectory(LOCAL_FILES);
     }
 

server/src/test/java/org/opensearch/index/store/DefaultCompositeDirectoryFactoryTests.java

@@ -59,7 +59,8 @@ public void testNewDirectory() throws IOException {
             shardPath,
             localDirectoryFactory,
             remoteSegmentStoreDirectory,
-            fileCache
+            fileCache,
+            threadPool
         );
         assertNotNull(directory);
         assert (directory instanceof CompositeDirectory);