[Remote Store] Integrate S3 async upload with S3BlobContainer (#7218)

Signed-off-by: Raghuvansh Raj <[email protected]>

plugins/repository-s3/src/internalClusterTest/java/org/opensearch/repositories/s3/S3BlobStoreRepositoryTests.java

@@ -172,7 +172,7 @@ protected S3Repository createRepository(
  ClusterService clusterService,
  RecoverySettings recoverySettings
  ) {
- return new S3Repository(metadata, registry, service, clusterService, recoverySettings) {
+ return new S3Repository(metadata, registry, service, clusterService, recoverySettings, null, null, null, null, false) {
 
  @Override
  public BlobStore blobStore() {

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

@@ -39,11 +39,15 @@
 import org.opensearch.action.ActionListener;
 import org.opensearch.common.Nullable;
 import org.opensearch.common.SetOnce;
+import org.opensearch.common.StreamContext;
 import org.opensearch.common.blobstore.BlobContainer;
 import org.opensearch.common.blobstore.BlobMetadata;
 import org.opensearch.common.blobstore.BlobPath;
 import org.opensearch.common.blobstore.BlobStoreException;
 import org.opensearch.common.blobstore.DeleteResult;
+import org.opensearch.common.blobstore.VerifyingMultiStreamBlobContainer;
+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;
@@ -72,6 +76,8 @@
 import software.amazon.awssdk.services.s3.model.UploadPartResponse;
 import software.amazon.awssdk.services.s3.paginators.ListObjectsV2Iterable;
 import org.opensearch.core.common.Strings;
+import org.opensearch.repositories.s3.async.UploadRequest;
+import software.amazon.awssdk.services.s3.S3AsyncClient;
 
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
@@ -82,6 +88,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -90,12 +97,13 @@
 import static org.opensearch.repositories.s3.S3Repository.MAX_FILE_SIZE_USING_MULTIPART;
 import static org.opensearch.repositories.s3.S3Repository.MIN_PART_SIZE_USING_MULTIPART;
 
-class S3BlobContainer extends AbstractBlobContainer {
+class S3BlobContainer extends AbstractBlobContainer implements VerifyingMultiStreamBlobContainer {
 
  private static final Logger logger = LogManager.getLogger(S3BlobContainer.class);
 
  /**
  * Maximum number of deletes in a {@link DeleteObjectsRequest}.
+ *
  * @see <a href="https://docs.aws.amazon.com/AmazonS3/latest/API/multiobjectdeleteapi.html">S3 Documentation</a>.
  */
  private static final int MAX_BULK_DELETES = 1000;
@@ -166,6 +174,42 @@ public void writeBlob(String blobName, InputStream inputStream, long blobSize, b
  });
  }
 
+ @Override
+ public void asyncBlobUpload(WriteContext writeContext, ActionListener<Void> completionListener) throws IOException {
+ UploadRequest uploadRequest = new UploadRequest(
+ blobStore.bucket(),
+ buildKey(writeContext.getFileName()),
+ writeContext.getFileSize(),
+ writeContext.getWritePriority(),
+ writeContext.getUploadFinalizer(),
+ writeContext.doRemoteDataIntegrityCheck(),
+ writeContext.getExpectedChecksum()
+ );
+ try {
+ long partSize = blobStore.getAsyncTransferManager().calculateOptimalPartSize(writeContext.getFileSize());
+ StreamContext streamContext = SocketAccess.doPrivileged(() -> writeContext.getStreamProvider(partSize));
+ try (AmazonAsyncS3Reference amazonS3Reference = SocketAccess.doPrivileged(blobStore::asyncClientReference)) {
+
+ S3AsyncClient s3AsyncClient = writeContext.getWritePriority() == WritePriority.HIGH
+ ? amazonS3Reference.get().priorityClient()
+ : amazonS3Reference.get().client();
+ CompletableFuture<Void> completableFuture = blobStore.getAsyncTransferManager()
+ .uploadObject(s3AsyncClient, uploadRequest, streamContext);
+ completableFuture.whenComplete((response, throwable) -> {
+ if (throwable == null) {
+ completionListener.onResponse(response);
+ } else {
+ Exception ex = throwable instanceof Error ? new Exception(throwable) : (Exception) throwable;
+ completionListener.onFailure(ex);
+ }
+ });
+ }
+ } catch (Exception e) {
+ logger.info("exception error from blob container for file {}", writeContext.getFileName());
+ throw new IOException(e);
+ }
+ }
+
  // package private for testing
  long getLargeBlobThresholdInBytes() {
  return blobStore.bufferSizeInBytes();

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

@@ -42,6 +42,8 @@
 import org.opensearch.common.unit.ByteSizeValue;
 import software.amazon.awssdk.services.s3.model.ObjectCannedACL;
 import software.amazon.awssdk.services.s3.model.StorageClass;
+import org.opensearch.repositories.s3.async.AsyncExecutorContainer;
+import org.opensearch.repositories.s3.async.AsyncTransferManager;
 
 import java.io.IOException;
 import java.util.Locale;
@@ -53,6 +55,8 @@ class S3BlobStore implements BlobStore {
 
  private final S3Service service;
 
+ private final S3AsyncService s3AsyncService;
+
  private final String bucket;
 
  private final ByteSizeValue bufferSize;
@@ -67,22 +71,41 @@ class S3BlobStore implements BlobStore {
 
  private final StatsMetricPublisher statsMetricPublisher = new StatsMetricPublisher();
 
+ private final AsyncTransferManager asyncTransferManager;
+ private final AsyncExecutorContainer priorityExecutorBuilder;
+ private final AsyncExecutorContainer normalExecutorBuilder;
+ private final boolean multipartUploadEnabled;
+
  S3BlobStore(
  S3Service service,
+ S3AsyncService s3AsyncService,
+ boolean multipartUploadEnabled,
  String bucket,
  boolean serverSideEncryption,
  ByteSizeValue bufferSize,
  String cannedACL,
  String storageClass,
- RepositoryMetadata repositoryMetadata
+ RepositoryMetadata repositoryMetadata,
+ AsyncTransferManager asyncTransferManager,
+ AsyncExecutorContainer priorityExecutorBuilder,
+ AsyncExecutorContainer normalExecutorBuilder
  ) {
  this.service = service;
+ this.s3AsyncService = s3AsyncService;
+ this.multipartUploadEnabled = multipartUploadEnabled;
  this.bucket = bucket;
  this.serverSideEncryption = serverSideEncryption;
  this.bufferSize = bufferSize;
  this.cannedACL = initCannedACL(cannedACL);
  this.storageClass = initStorageClass(storageClass);
  this.repositoryMetadata = repositoryMetadata;
+ this.asyncTransferManager = asyncTransferManager;
+ this.normalExecutorBuilder = normalExecutorBuilder;
+ this.priorityExecutorBuilder = priorityExecutorBuilder;
+ }
+
+ public boolean isMultipartUploadEnabled() {
+ return multipartUploadEnabled;
  }
 
  @Override
@@ -94,6 +117,10 @@ public AmazonS3Reference clientReference() {
  return service.client(repositoryMetadata);
  }
 
+ public AmazonAsyncS3Reference asyncClientReference() {
+ return s3AsyncService.client(repositoryMetadata, priorityExecutorBuilder, normalExecutorBuilder);
+ }
+
  int getMaxRetries() {
  return service.settings(repositoryMetadata).maxRetries;
  }
@@ -117,7 +144,12 @@ public BlobContainer blobContainer(BlobPath path) {
 
  @Override
  public void close() throws IOException {
- this.service.close();
+ if (service != null) {
+ this.service.close();
+ }
+ if (s3AsyncService != null) {
+ this.s3AsyncService.close();
+ }
  }
 
  @Override
@@ -170,4 +202,8 @@ public static ObjectCannedACL initCannedACL(String cannedACL) {
 
  throw new BlobStoreException("cannedACL is not valid: [" + cannedACL + "]");
  }
+
+ public AsyncTransferManager getAsyncTransferManager() {
+ return asyncTransferManager;
+ }
 }

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

@@ -34,7 +34,6 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
-
 import org.opensearch.Version;
 import org.opensearch.action.ActionListener;
 import org.opensearch.cluster.ClusterState;
@@ -57,6 +56,8 @@
 import org.opensearch.repositories.RepositoryException;
 import org.opensearch.repositories.ShardGenerations;
 import org.opensearch.repositories.blobstore.MeteredBlobStoreRepository;
+import org.opensearch.repositories.s3.async.AsyncExecutorContainer;
+import org.opensearch.repositories.s3.async.AsyncTransferManager;
 import org.opensearch.snapshots.SnapshotId;
 import org.opensearch.snapshots.SnapshotInfo;
 import org.opensearch.threadpool.Scheduler;
@@ -103,6 +104,11 @@ class S3Repository extends MeteredBlobStoreRepository {
  ByteSizeUnit.BYTES
  );
 
+ private static final ByteSizeValue DEFAULT_MULTIPART_UPLOAD_MINIMUM_PART_SIZE = new ByteSizeValue(
+ ByteSizeUnit.MB.toBytes(16),
+ ByteSizeUnit.BYTES
+ );
+
  static final Setting<String> BUCKET_SETTING = Setting.simpleString("bucket");
 
  /**
@@ -146,6 +152,26 @@ class S3Repository extends MeteredBlobStoreRepository {
  MAX_PART_SIZE_USING_MULTIPART
  );
 
+ /**
+ * Minimum part size for parallel multipart uploads
+ */
+ static final Setting<ByteSizeValue> PARALLEL_MULTIPART_UPLOAD_MINIMUM_PART_SIZE_SETTING = Setting.byteSizeSetting(
+ "parallel_multipart_upload.minimum_part_size",
+ DEFAULT_MULTIPART_UPLOAD_MINIMUM_PART_SIZE,
+ MIN_PART_SIZE_USING_MULTIPART,
+ MAX_PART_SIZE_USING_MULTIPART,
+ Setting.Property.NodeScope
+ );
+
+ /**
+ * This setting controls whether parallel multipart uploads will be used when calling S3 or not
+ */
+ public static Setting<Boolean> PARALLEL_MULTIPART_UPLOAD_ENABLED_SETTING = Setting.boolSetting(
+ "parallel_multipart_upload.enabled",
+ true,
+ Setting.Property.NodeScope
+ );
+
  /**
  * Big files can be broken down into chunks during snapshotting if needed. Defaults to 1g.
  */
@@ -193,6 +219,12 @@ class S3Repository extends MeteredBlobStoreRepository {
 
  private final RepositoryMetadata repositoryMetadata;
 
+ private final AsyncTransferManager asyncUploadUtils;
+ private final S3AsyncService s3AsyncService;
+ private final boolean multipartUploadEnabled;
+ private final AsyncExecutorContainer priorityExecutorBuilder;
+ private final AsyncExecutorContainer normalExecutorBuilder;
+
  /**
  * Constructs an s3 backed repository
  */
@@ -201,7 +233,12 @@ class S3Repository extends MeteredBlobStoreRepository {
  final NamedXContentRegistry namedXContentRegistry,
  final S3Service service,
  final ClusterService clusterService,
- final RecoverySettings recoverySettings
+ final RecoverySettings recoverySettings,
+ final AsyncTransferManager asyncUploadUtils,
+ final AsyncExecutorContainer priorityExecutorBuilder,
+ final AsyncExecutorContainer normalExecutorBuilder,
+ final S3AsyncService s3AsyncService,
+ final boolean multipartUploadEnabled
  ) {
  super(
  metadata,
@@ -212,8 +249,13 @@ class S3Repository extends MeteredBlobStoreRepository {
  buildLocation(metadata)
  );
  this.service = service;
+ this.s3AsyncService = s3AsyncService;
+ this.multipartUploadEnabled = multipartUploadEnabled;
 
  this.repositoryMetadata = metadata;
+ this.asyncUploadUtils = asyncUploadUtils;
+ this.priorityExecutorBuilder = priorityExecutorBuilder;
+ this.normalExecutorBuilder = normalExecutorBuilder;
 
  // Parse and validate the user's S3 Storage Class setting
  this.bucket = BUCKET_SETTING.get(metadata.settings());
@@ -314,7 +356,20 @@ public void deleteSnapshots(
 
  @Override
  protected S3BlobStore createBlobStore() {
- return new S3BlobStore(service, bucket, serverSideEncryption, bufferSize, cannedACL, storageClass, repositoryMetadata);
+ return new S3BlobStore(
+ service,
+ s3AsyncService,
+ multipartUploadEnabled,
+ bucket,
+ serverSideEncryption,
+ bufferSize,
+ cannedACL,
+ storageClass,
+ repositoryMetadata,
+ asyncUploadUtils,
+ priorityExecutorBuilder,
+ normalExecutorBuilder
+ );
  }
 
  // only use for testing