Initial commit for POC

plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java

@@ -46,13 +46,15 @@
 import org.opensearch.common.blobstore.BlobStoreException;
 import org.opensearch.common.blobstore.DeleteResult;
 import org.opensearch.common.blobstore.VerifyingMultiStreamBlobContainer;
+import org.opensearch.common.blobstore.stream.read.ReadContext;
 import org.opensearch.common.blobstore.stream.write.WriteContext;
 import org.opensearch.common.blobstore.stream.write.WritePriority;
 import org.opensearch.common.blobstore.support.AbstractBlobContainer;
 import org.opensearch.common.blobstore.support.PlainBlobMetadata;
 import org.opensearch.common.collect.Tuple;
 import org.opensearch.common.unit.ByteSizeUnit;
 import org.opensearch.common.unit.ByteSizeValue;
+import org.opensearch.repositories.s3.async.DownloadRequest;
 import software.amazon.awssdk.core.exception.SdkException;
 import software.amazon.awssdk.core.sync.RequestBody;
 import software.amazon.awssdk.services.s3.model.AbortMultipartUploadRequest;
@@ -64,10 +66,13 @@
 import software.amazon.awssdk.services.s3.model.Delete;
 import software.amazon.awssdk.services.s3.model.DeleteObjectsRequest;
 import software.amazon.awssdk.services.s3.model.DeleteObjectsResponse;
+import software.amazon.awssdk.services.s3.model.GetObjectAttributesRequest;
+import software.amazon.awssdk.services.s3.model.GetObjectAttributesResponse;
 import software.amazon.awssdk.services.s3.model.HeadObjectRequest;
 import software.amazon.awssdk.services.s3.model.ListObjectsV2Request;
 import software.amazon.awssdk.services.s3.model.ListObjectsV2Response;
 import software.amazon.awssdk.services.s3.model.NoSuchKeyException;
+import software.amazon.awssdk.services.s3.model.ObjectAttributes;
 import software.amazon.awssdk.services.s3.model.ObjectIdentifier;
 import software.amazon.awssdk.services.s3.model.PutObjectRequest;
 import software.amazon.awssdk.services.s3.model.S3Error;
@@ -89,6 +94,7 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Future;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -210,6 +216,50 @@ public void asyncBlobUpload(WriteContext writeContext, ActionListener<Void> comp
         }
     }
 
+    @Override
+    public CompletableFuture<ReadContext> asyncBlobDownload(String blobName, boolean forceSingleStream) throws IOException {
+        try (AmazonAsyncS3Reference amazonS3Reference = SocketAccess.doPrivileged(blobStore::asyncClientReference)) {
+            S3AsyncClient s3AsyncClient = amazonS3Reference.get().client();
+
+            GetObjectAttributesRequest getObjectAttributesRequest = GetObjectAttributesRequest.builder()
+                .bucket(blobStore.bucket())
+                .key(blobName)
+                .objectAttributes(ObjectAttributes.CHECKSUM, ObjectAttributes.OBJECT_SIZE, ObjectAttributes.OBJECT_PARTS)
+                .build();
+
+            GetObjectAttributesResponse getObject = s3AsyncClient.getObjectAttributes(getObjectAttributesRequest).join();
+
+            final long blobSize = getObject.objectSize();
+            final List<CompletableFuture<InputStream>> blobInputStreamsFuture = new ArrayList<>();
+            final List<InputStream> blobInputStreams = new ArrayList<>();
+            final int numStreams;
+
+            if (forceSingleStream) {
+                numStreams = 1;
+                DownloadRequest downloadRequest = new DownloadRequest(blobStore.bucket(), blobName, blobSize);
+                blobInputStreamsFuture.add(blobStore.getAsyncTransferManager().downloadObjectFutureStream(s3AsyncClient, downloadRequest)
+                    // TODO: Add error handling
+                    .whenComplete((data, error) -> blobInputStreams.add(data)));
+            } else {
+                final long optimalStreamSize = 8 * 1024 * 1024; // TODO: Replace this with configurable value
+                numStreams = (int) Math.ceil(blobSize * 1.0 / optimalStreamSize);
+
+                for (int streamNumber = 0; streamNumber < numStreams; streamNumber++) {
+                    long start = streamNumber * optimalStreamSize;
+                    long end = Math.min(blobSize, (streamNumber + 1) * optimalStreamSize) - 1;
+                    DownloadRequest downloadRequest = new DownloadRequest(blobStore.bucket(), blobName, blobSize, start, end);
+                    blobInputStreamsFuture.add(blobStore.getAsyncTransferManager().downloadObjectFutureStream(s3AsyncClient, downloadRequest)
+                        // TODO: Add error handling
+                        .whenComplete((data, error) -> blobInputStreams.add(data)));
+                }
+            }
+
+            // TODO: Add error handling
+            return CompletableFuture.allOf(blobInputStreamsFuture.toArray(new CompletableFuture[0]))
+                .thenCompose(ignored -> CompletableFuture.supplyAsync(() -> new ReadContext(blobInputStreams, null, numStreams, blobSize)));
+        }
+    }
+
     // package private for testing
     long getLargeBlobThresholdInBytes() {
         return blobStore.bufferSizeInBytes();
@@ -392,16 +442,16 @@ public Map<String, BlobMetadata> listBlobs() throws IOException {
     public Map<String, BlobContainer> children() throws IOException {
         try (AmazonS3Reference clientReference = blobStore.clientReference()) {
             return executeListing(clientReference, listObjectsRequest(keyPath)).stream().flatMap(listObjectsResponse -> {
-                assert listObjectsResponse.contents().stream().noneMatch(s -> {
-                    for (CommonPrefix commonPrefix : listObjectsResponse.commonPrefixes()) {
-                        if (s.key().substring(keyPath.length()).startsWith(commonPrefix.prefix())) {
-                            return true;
+                    assert listObjectsResponse.contents().stream().noneMatch(s -> {
+                        for (CommonPrefix commonPrefix : listObjectsResponse.commonPrefixes()) {
+                            if (s.key().substring(keyPath.length()).startsWith(commonPrefix.prefix())) {
+                                return true;
+                            }
                         }
-                    }
-                    return false;
-                }) : "Response contained children for listed common prefixes.";
-                return listObjectsResponse.commonPrefixes().stream();
-            })
+                        return false;
+                    }) : "Response contained children for listed common prefixes.";
+                    return listObjectsResponse.commonPrefixes().stream();
+                })
                 .map(commonPrefix -> commonPrefix.prefix().substring(keyPath.length()))
                 .filter(name -> name.isEmpty() == false)
                 // Stripping the trailing slash off of the common prefix

plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/AsyncTransferManager.java

@@ -21,7 +21,9 @@
 import org.opensearch.common.util.ByteUtils;
 import org.opensearch.repositories.s3.io.CheckedContainer;
 import org.opensearch.repositories.s3.SocketAccess;
+import org.opensearch.repositories.s3.utils.HttpRangeUtils;
 import software.amazon.awssdk.core.async.AsyncRequestBody;
+import software.amazon.awssdk.core.async.AsyncResponseTransformer;
 import software.amazon.awssdk.core.exception.SdkClientException;
 import software.amazon.awssdk.http.HttpStatusCode;
 import software.amazon.awssdk.services.s3.S3AsyncClient;
@@ -33,11 +35,13 @@
 import software.amazon.awssdk.services.s3.model.CreateMultipartUploadRequest;
 import software.amazon.awssdk.services.s3.model.CreateMultipartUploadResponse;
 import software.amazon.awssdk.services.s3.model.DeleteObjectRequest;
+import software.amazon.awssdk.services.s3.model.GetObjectRequest;
 import software.amazon.awssdk.services.s3.model.PutObjectRequest;
 import software.amazon.awssdk.services.s3.model.S3Exception;
 import software.amazon.awssdk.utils.CompletableFutureUtils;
 
 import java.io.IOException;
+import java.io.InputStream;
 import java.util.Arrays;
 import java.util.Base64;
 import java.util.List;
@@ -102,6 +106,22 @@ public CompletableFuture<Void> uploadObject(S3AsyncClient s3AsyncClient, UploadR
         return returnFuture;
     }
 
+    public CompletableFuture<InputStream> downloadObjectFutureStream(S3AsyncClient s3AsyncClient, DownloadRequest downloadRequest) {
+        GetObjectRequest.Builder getObjectRequestBuilder = GetObjectRequest.builder()
+            .bucket(downloadRequest.getBucket())
+            .key(downloadRequest.getKey());
+
+        if (downloadRequest.getStart() != -1 && downloadRequest.getEnd() != -1) {
+            getObjectRequestBuilder.range(HttpRangeUtils.toHttpRangeHeader(downloadRequest.getStart(), downloadRequest.getEnd()));
+        }
+
+        CompletableFuture<InputStream> objectFutureStream =
+            s3AsyncClient.getObject(getObjectRequestBuilder.build(), AsyncResponseTransformer.toBlockingInputStream())
+                .thenCompose(x -> CompletableFuture.supplyAsync(() -> x));
+
+        return objectFutureStream;
+    }
+
     private void uploadInParts(
         S3AsyncClient s3AsyncClient,
         UploadRequest uploadRequest,

plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/DownloadRequest.java

@@ -0,0 +1,52 @@
+/*
+ * 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.repositories.s3.async;
+
+/**
+ * A model encapsulating all details for a download to S3
+ */
+public class DownloadRequest {
+    private final String bucket;
+    private final String key;
+    private final long contentLength;
+    private final long start;
+    private final long end;
+
+    public DownloadRequest(String bucket, String key, long contentLength, long start, long end) {
+        this.bucket = bucket;
+        this.key = key;
+        this.contentLength = contentLength;
+        this.start = start;
+        this.end = end;
+    }
+
+    public DownloadRequest(String bucket, String key, long contentLength) {
+        this(bucket, key, contentLength, -1, -1);
+    }
+
+    public String getBucket() {
+        return bucket;
+    }
+
+    public String getKey() {
+        return key;
+    }
+
+    public long getContentLength() {
+        return contentLength;
+    }
+
+    public long getStart() {
+        return start;
+    }
+
+    public long getEnd() {
+        return end;
+    }
+}

server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreIT.java

@@ -85,7 +85,7 @@ private Map<String, Long> indexData(int numberOfIterations, boolean invokeFlush)
             }
             maxSeqNoRefreshedOrFlushed = maxSeqNo;
             refreshedOrFlushedOperations = totalOperations;
-            int numberOfOperations = randomIntBetween(20, 50);
+            int numberOfOperations = randomIntBetween(1, 5);
             for (int j = 0; j < numberOfOperations; j++) {
                 IndexResponse response = indexSingleDoc();
                 maxSeqNo = response.getSeqNo();
@@ -112,7 +112,7 @@ private void verifyRestoredData(Map<String, Long> indexStats, boolean checkTotal
         assertHitCount(client().prepareSearch(INDEX_NAME).setSize(0).get(), indexStats.get(statsGranularity) + 1);
     }
 
-    private void testRestoreFlow(boolean remoteTranslog, int numberOfIterations, boolean invokeFlush) throws IOException {
+    protected void testRestoreFlow(boolean remoteTranslog, int numberOfIterations, boolean invokeFlush) throws IOException {
         internalCluster().startDataOnlyNodes(3);
         if (remoteTranslog) {
             createIndex(INDEX_NAME, remoteTranslogIndexSettings(0));
@@ -129,12 +129,10 @@ private void testRestoreFlow(boolean remoteTranslog, int numberOfIterations, boo
 
         client().admin().cluster().restoreRemoteStore(new RestoreRemoteStoreRequest().indices(INDEX_NAME), PlainActionFuture.newFuture());
         ensureGreen(INDEX_NAME);
+        // assertEquals(indexStats.get(TOTAL_OPERATIONS).longValue(),
+        // client().prepareSearch(INDEX_NAME).setSize(0).get().getInternalResponse().hits().getTotalHits().value);
+        verifyRestoredData(indexStats, remoteTranslog);
 
-        if (remoteTranslog) {
-            verifyRestoredData(indexStats, true);
-        } else {
-            verifyRestoredData(indexStats, false);
-        }
     }
 
     public void testRemoteSegmentStoreRestoreWithNoDataPostCommit() throws IOException {

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

@@ -13,6 +13,7 @@
 import org.opensearch.remotestore.RemoteStoreIT;
 import org.opensearch.remotestore.multipart.mocks.MockFsRepositoryPlugin;
 
+import java.io.IOException;
 import java.nio.file.Path;
 import java.util.Collection;
 import java.util.stream.Collectors;
@@ -29,10 +30,16 @@ protected Collection<Class<? extends Plugin>> nodePlugins() {
 
     @Override
     protected void putRepository(Path path) {
+        logger.error("Repo Path: {}", path);
         assertAcked(
             clusterAdmin().preparePutRepository(REPOSITORY_NAME)
                 .setType(MockFsRepositoryPlugin.TYPE)
                 .setSettings(Settings.builder().put("location", path))
         );
     }
+
+    @Override
+    public void testRemoteSegmentStoreRestoreWithNoDataPostCommit() throws IOException {
+        testRestoreFlow(true, 2, true);
+    }
 }