software.amazon.awssdk.services.s3.internal.multipart.GenericMultipartHelper Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of s3 Show documentation
The AWS Java SDK for Amazon S3 module holds the client classes that are used for communicating with Amazon Simple Storage Service
There is a newer version: 2.30.1
Show newest version
/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License").
 * You may not use this file except in compliance with the License.
 * A copy of the License is located at
 *
 *  http://aws.amazon.com/apache2.0
 *
 * or in the "license" file accompanying this file. This file is distributed
 * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
 * express or implied. See the License for the specific language governing
 * permissions and limitations under the License.
 */

package software.amazon.awssdk.services.s3.internal.multipart;

import static software.amazon.awssdk.services.s3.internal.multipart.SdkPojoConversionUtils.toCompleteMultipartUploadRequest;
import static software.amazon.awssdk.services.s3.multipart.S3MultipartExecutionAttribute.JAVA_PROGRESS_LISTENER;

import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.function.BiFunction;
import java.util.function.Function;
import java.util.function.Supplier;
import software.amazon.awssdk.annotations.SdkInternalApi;
import software.amazon.awssdk.core.async.listener.PublisherListener;
import software.amazon.awssdk.core.exception.SdkClientException;
import software.amazon.awssdk.core.exception.SdkException;
import software.amazon.awssdk.services.s3.S3AsyncClient;
import software.amazon.awssdk.services.s3.model.AbortMultipartUploadRequest;
import software.amazon.awssdk.services.s3.model.CompleteMultipartUploadRequest;
import software.amazon.awssdk.services.s3.model.CompleteMultipartUploadResponse;
import software.amazon.awssdk.services.s3.model.CompletedPart;
import software.amazon.awssdk.services.s3.model.PutObjectRequest;
import software.amazon.awssdk.services.s3.model.S3Request;
import software.amazon.awssdk.services.s3.model.S3Response;
import software.amazon.awssdk.utils.Logger;

@SdkInternalApi
public final class GenericMultipartHelper {
    private static final Logger log = Logger.loggerFor(GenericMultipartHelper.class);
    /**
     * The max number of parts on S3 side is 10,000
     */
    private static final long MAX_UPLOAD_PARTS = 10_000;

    private final S3AsyncClient s3AsyncClient;
    private final Function abortMultipartUploadRequestConverter;
    private final Function responseConverter;

    public GenericMultipartHelper(S3AsyncClient s3AsyncClient,
                                  Function abortMultipartUploadRequestConverter,
                                  Function responseConverter) {
        this.s3AsyncClient = s3AsyncClient;
        this.abortMultipartUploadRequestConverter = abortMultipartUploadRequestConverter;
        this.responseConverter = responseConverter;
    }

    public void handleException(CompletableFuture returnFuture,
                                Supplier message,
                                Throwable throwable) {
        Throwable cause = throwable instanceof CompletionException ? throwable.getCause() : throwable;

        if (cause instanceof Error || cause instanceof SdkException) {
            cause.addSuppressed(SdkClientException.create(message.get()));
            returnFuture.completeExceptionally(cause);
        } else {
            SdkClientException exception = SdkClientException.create(message.get(), cause);
            returnFuture.completeExceptionally(exception);
        }
    }

    public long calculateOptimalPartSizeFor(long contentLengthOfSource, long partSizeInBytes) {
        double optimalPartSize = contentLengthOfSource / (double) MAX_UPLOAD_PARTS;

        optimalPartSize = Math.ceil(optimalPartSize);
        return (long) Math.max(optimalPartSize, partSizeInBytes);
    }

    public int determinePartCount(long contentLength, long partSize) {
        return (int) Math.ceil(contentLength / (double) partSize);
    }

    public CompletableFuture completeMultipartUpload(
        PutObjectRequest request, String uploadId, CompletedPart[] parts) {
        log.debug(() -> String.format("Sending completeMultipartUploadRequest, uploadId: %s",
                                      uploadId));

        CompleteMultipartUploadRequest completeMultipartUploadRequest = toCompleteMultipartUploadRequest(request, uploadId,
                                                                                                         parts);
        return s3AsyncClient.completeMultipartUpload(completeMultipartUploadRequest);
    }

    public BiFunction handleExceptionOrResponse(RequestT request,
        CompletableFuture returnFuture, String uploadId) {
        PublisherListener progressListener = request.overrideConfiguration()
                                                          .map(c -> c.executionAttributes().getAttribute(JAVA_PROGRESS_LISTENER))
                                                          .orElseGet(PublisherListener::noOp);

        return (completeMultipartUploadResponse, throwable) -> {
            if (throwable != null) {
                cleanUpParts(uploadId, abortMultipartUploadRequestConverter.apply(request));
                handleException(returnFuture, () -> "Failed to send multipart requests",
                                throwable);
            } else {
                returnFuture.complete(responseConverter.apply(
                    completeMultipartUploadResponse));
                progressListener.subscriberOnComplete();
            }

            return null;
        };
    }

    public void cleanUpParts(String uploadId, AbortMultipartUploadRequest.Builder abortMultipartUploadRequest) {
        log.debug(() -> "Aborting multipart upload: " + uploadId);
        AbortMultipartUploadRequest request = abortMultipartUploadRequest.uploadId(uploadId).build();
        s3AsyncClient.abortMultipartUpload(request)
                     .exceptionally(throwable -> {
                         log.warn(() -> String.format("Failed to abort previous multipart upload "
                                                      + "(id: %s)"
                                                      + ". You may need to call "
                                                      + "S3AsyncClient#abortMultiPartUpload to "
                                                      + "free all storage consumed by"
                                                      + " all parts. ",
                                                      uploadId), throwable);
                         return null;
                     });
    }
}