Enabling Support for Conditional Multi-Part Upload

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

@@ -40,6 +40,7 @@
 import software.amazon.awssdk.services.s3.model.AbortMultipartUploadRequest;
 import software.amazon.awssdk.services.s3.model.CommonPrefix;
 import software.amazon.awssdk.services.s3.model.CompleteMultipartUploadRequest;
+import software.amazon.awssdk.services.s3.model.CompleteMultipartUploadResponse;
 import software.amazon.awssdk.services.s3.model.CompletedMultipartUpload;
 import software.amazon.awssdk.services.s3.model.CompletedPart;
 import software.amazon.awssdk.services.s3.model.CreateMultipartUploadRequest;
@@ -53,6 +54,7 @@
 import software.amazon.awssdk.services.s3.model.NoSuchKeyException;
 import software.amazon.awssdk.services.s3.model.ObjectAttributes;
 import software.amazon.awssdk.services.s3.model.PutObjectRequest;
+import software.amazon.awssdk.services.s3.model.S3Exception;
 import software.amazon.awssdk.services.s3.model.ServerSideEncryption;
 import software.amazon.awssdk.services.s3.model.UploadPartRequest;
 import software.amazon.awssdk.services.s3.model.UploadPartResponse;
@@ -63,6 +65,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.apache.logging.log4j.message.ParameterizedMessage;
+import org.opensearch.OpenSearchException;
 import org.opensearch.action.support.PlainActionFuture;
 import org.opensearch.common.Nullable;
 import org.opensearch.common.SetOnce;
@@ -97,6 +100,7 @@
 import java.io.InputStream;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
@@ -117,6 +121,7 @@ class S3BlobContainer extends AbstractBlobContainer implements AsyncMultiStreamB
 
     private final S3BlobStore blobStore;
     private final String keyPath;
+    public static final int HTTP_STATUS_PRECONDITION_FAILED = 412;
 
     S3BlobContainer(BlobPath path, S3BlobStore blobStore) {
         super(path);
@@ -509,6 +514,169 @@ private String buildKey(String blobName) {
         return keyPath + blobName;
     }
 
+    public void executeMultipartUploadIfEtagMatches(
+        final S3BlobStore blobStore,
+        final String blobName,
+        final InputStream input,
+        final long blobSize,
+        final Map<String, String> metadata,
+        final String eTag,
+        final ActionListener<String> etagListener
+    ) throws IOException {
+
+        ensureMultiPartUploadSize(blobSize);
+
+        final long partSize = blobStore.bufferSizeInBytes();
+        final Tuple<Long, Long> multiparts = numberOfMultiparts(blobSize, partSize);
+        if (multiparts.v1() > Integer.MAX_VALUE) {
+            throw new IllegalArgumentException("Too many multipart upload parts; consider a larger buffer size.");
+        }
+        final int nbParts = multiparts.v1().intValue();
+        final long lastPartSize = multiparts.v2();
+        assert blobSize == (((nbParts - 1) * partSize) + lastPartSize) : "blobSize does not match multipart sizes";
+        // test
+        CreateMultipartUploadRequest.Builder createRequestBuilder = CreateMultipartUploadRequest.builder()
+            .bucket(blobStore.bucket())
+            .key(blobName)
+            .storageClass(blobStore.getStorageClass())
+            .acl(blobStore.getCannedACL())
+            .overrideConfiguration(o -> o.addMetricPublisher(blobStore.getStatsMetricPublisher().multipartUploadMetricCollector));
+
+        if (metadata != null && !metadata.isEmpty()) {
+            createRequestBuilder.metadata(metadata);
+        }
+        if (blobStore.serverSideEncryption()) {
+            createRequestBuilder.serverSideEncryption(ServerSideEncryption.AES256);
+        }
+
+        final CreateMultipartUploadRequest createMultipartUploadRequest = createRequestBuilder.build();
+        final SetOnce<String> uploadId = new SetOnce<>();
+        final String bucketName = blobStore.bucket();
+        boolean success = false;
+
+        final InputStream requestInputStream = blobStore.isUploadRetryEnabled()
+            ? new BufferedInputStream(input, (int) (partSize + 1))
+            : input;
+
+        try (AmazonS3Reference clientReference = blobStore.clientReference()) {
+            uploadId.set(
+                SocketAccess.doPrivileged(() -> clientReference.get().createMultipartUpload(createMultipartUploadRequest).uploadId())
+            );
+            if (Strings.isEmpty(uploadId.get())) {
+                IOException exception = new IOException("Failed to initialize multipart upload for " + blobName);
+                etagListener.onFailure(exception);
+                throw exception;
+            }
+
+            final List<CompletedPart> parts = new ArrayList<>(nbParts);
+            long bytesCount = 0;
+
+            for (int i = 1; i <= nbParts; i++) {
+                long currentPartSize = (i < nbParts) ? partSize : lastPartSize;
+                final UploadPartRequest uploadPartRequest = UploadPartRequest.builder()
+                    .bucket(bucketName)
+                    .key(blobName)
+                    .uploadId(uploadId.get())
+                    .partNumber(i)
+                    .contentLength(currentPartSize)
+                    .overrideConfiguration(o -> o.addMetricPublisher(blobStore.getStatsMetricPublisher().multipartUploadMetricCollector))
+                    .build();
+
+                bytesCount += currentPartSize;
+
+                final UploadPartResponse uploadResponse = SocketAccess.doPrivileged(
+                    () -> clientReference.get()
+                        .uploadPart(uploadPartRequest, RequestBody.fromInputStream(requestInputStream, currentPartSize))
+                );
+
+                String partETag = uploadResponse.eTag();
+                if (partETag == null) {
+                    IOException exception = new IOException(
+                        String.format(Locale.ROOT, "S3 part upload for [%s] part [%d] returned null ETag", blobName, i)
+                    );
+                    etagListener.onFailure(exception);
+                    throw exception;
+                }
+
+                parts.add(CompletedPart.builder().partNumber(i).eTag(partETag).build());
+            }
+
+            if (bytesCount != blobSize) {
+                IOException exception = new IOException(
+                    String.format(Locale.ROOT, "Multipart upload for [%s] sent %d bytes; expected %d bytes", blobName, bytesCount, blobSize)
+                );
+                etagListener.onFailure(exception);
+                throw exception;
+            }
+
+            CompleteMultipartUploadRequest completeRequest = CompleteMultipartUploadRequest.builder()
+                .bucket(bucketName)
+                .key(blobName)
+                .uploadId(uploadId.get())
+                .multipartUpload(CompletedMultipartUpload.builder().parts(parts).build())
+                .ifMatch(eTag)
+                .overrideConfiguration(o -> o.addMetricPublisher(blobStore.getStatsMetricPublisher().multipartUploadMetricCollector))
+                .build();
+
+            CompleteMultipartUploadResponse completeResponse = SocketAccess.doPrivileged(
+                () -> clientReference.get().completeMultipartUpload(completeRequest)
+            );
+
+            if (completeResponse.eTag() != null) {
+                success = true;
+                etagListener.onResponse(completeResponse.eTag());
+            } else {
+                IOException exception = new IOException(
+                    "S3 multipart upload for [" + blobName + "] returned null ETag, violating data integrity expectations"
+                );
+                etagListener.onFailure(exception);
+                throw exception;
+            }
+
+        } catch (S3Exception e) {
+            if (e.statusCode() == HTTP_STATUS_PRECONDITION_FAILED) {
+                etagListener.onFailure(new OpenSearchException("stale_primary_shard", e, "Precondition Failed : Etag Mismatch", blobName));
+                throw new IOException("Unable to upload object [" + blobName + "] due to ETag mismatch", e);
+            } else {
+                IOException exception = new IOException(
+                    String.format(Locale.ROOT, "S3 error during multipart upload [%s]: %s", blobName, e.getMessage()),
+                    e
+                );
+                etagListener.onFailure(exception);
+                throw exception;
+            }
+        } catch (SdkException e) {
+            IOException exception = new IOException(String.format(Locale.ROOT, "S3 multipart upload failed for [%s]", blobName), e);
+            etagListener.onFailure(exception);
+            throw exception;
+        } catch (Exception e) {
+            IOException exception = new IOException(
+                String.format(Locale.ROOT, "Unexpected error during multipart upload [%s]: %s", blobName, e.getMessage()),
+                e
+            );
+            etagListener.onFailure(exception);
+            throw exception;
+        } finally {
+            if (!success && Strings.hasLength(uploadId.get())) {
+                AbortMultipartUploadRequest abortRequest = AbortMultipartUploadRequest.builder()
+                    .bucket(bucketName)
+                    .key(blobName)
+                    .uploadId(uploadId.get())
+                    .build();
+                try (AmazonS3Reference abortClient = blobStore.clientReference()) {
+                    SocketAccess.doPrivilegedVoid(() -> abortClient.get().abortMultipartUpload(abortRequest));
+                } catch (Exception abortException) {
+                    logger.warn(
+                        "Failed to abort incomplete multipart upload [{}] with ID [{}]. "
+                            + "This may result in orphaned S3 data and charges.",
+                        new Object[] { blobName, uploadId.get() },
+                        abortException
+                    );
+                }
+            }
+        }
+    }
+
     /**
      * Uploads a blob using a single upload request
      */