com.netflix.genie.web.services.impl.S3AttachmentServiceImpl Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of genie-web Show documentation
genie-web
There is a newer version: 4.3.20
/*
 *
 *  Copyright 2020 Netflix, Inc.
 *
 *     Licensed under the Apache License, Version 2.0 (the "License");
 *     you may not use this file except in compliance with the License.
 *     You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0
 *
 *     Unless required by applicable law or agreed to in writing, software
 *     distributed under the License 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 com.netflix.genie.web.services.impl;

import com.amazonaws.SdkClientException;
import com.amazonaws.services.s3.AmazonS3;
import com.amazonaws.services.s3.AmazonS3URI;
import com.amazonaws.services.s3.model.ObjectMetadata;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Sets;
import com.netflix.genie.common.internal.aws.s3.S3ClientFactory;
import com.netflix.genie.web.exceptions.checked.AttachmentTooLargeException;
import com.netflix.genie.web.exceptions.checked.SaveAttachmentException;
import com.netflix.genie.web.properties.AttachmentServiceProperties;
import com.netflix.genie.web.services.AttachmentService;
import com.netflix.genie.web.util.MetricsUtils;
import io.micrometer.core.instrument.MeterRegistry;
import io.micrometer.core.instrument.Tag;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.core.io.Resource;

import javax.annotation.Nullable;
import java.io.IOException;
import java.io.InputStream;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.TimeUnit;

/**
 * Implementation of the AttachmentService interface which saves attachments to AWS S3.
 *
 * @author mprimi
 * @since 4.0.0
 */
@Slf4j
public class S3AttachmentServiceImpl implements AttachmentService {

    private static final String METRICS_PREFIX = "genie.jobs.attachments.s3";
    private static final String COUNT_DISTRIBUTION = METRICS_PREFIX + ".count.distribution";
    private static final String LARGEST_SIZE_DISTRIBUTION = METRICS_PREFIX + ".largest.distribution";
    private static final String TOTAL_SIZE_DISTRIBUTION = METRICS_PREFIX + ".totalSize.distribution";
    private static final String SAVE_TIMER = METRICS_PREFIX + ".upload.timer";
    private static final Set EMPTY_SET = ImmutableSet.of();
    private static final String SLASH = "/";
    private static final String S3 = "s3";
    private final S3ClientFactory s3ClientFactory;
    private final AttachmentServiceProperties properties;
    private final MeterRegistry meterRegistry;
    private final AmazonS3URI s3BaseURI;

    /**
     * Constructor.
     *
     * @param s3ClientFactory             the s3 client factory
     * @param attachmentServiceProperties the service properties
     * @param meterRegistry               the meter registry
     */
    public S3AttachmentServiceImpl(
        final S3ClientFactory s3ClientFactory,
        final AttachmentServiceProperties attachmentServiceProperties,
        final MeterRegistry meterRegistry
    ) {
        this.s3ClientFactory = s3ClientFactory;
        this.properties = attachmentServiceProperties;
        this.meterRegistry = meterRegistry;
        this.s3BaseURI = new AmazonS3URI(attachmentServiceProperties.getLocationPrefix());
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public Set saveAttachments(
        @Nullable final String jobId,
        final Set attachments
    ) throws SaveAttachmentException {

        // Track number of attachments, including zeroes
        this.meterRegistry.summary(COUNT_DISTRIBUTION).record(attachments.size());

        log.debug("Saving {} attachments for job request with id: {}", attachments.size(), jobId);

        if (attachments.size() == 0) {
            return EMPTY_SET;
        }

        // Check for attachment size limits
        this.checkLimits(attachments);

        final long start = System.nanoTime();
        final Set tags = Sets.newHashSet();
        try {
            // Upload all to S3
            final Set attachmentURIs = this.uploadAllAttachments(jobId, attachments);
            MetricsUtils.addSuccessTags(tags);
            return attachmentURIs;
        } catch (SaveAttachmentException e) {
            log.error("Failed to save attachments (requested job id: {}): {}", jobId, e.getMessage(), e);
            MetricsUtils.addFailureTagsWithException(tags, e);
            throw e;
        } finally {
            this.meterRegistry
                .timer(SAVE_TIMER, tags)
                .record(System.nanoTime() - start, TimeUnit.NANOSECONDS);
        }
    }

    private void checkLimits(final Set attachments) throws SaveAttachmentException {

        final long singleSizeLimit = this.properties.getMaxSize().toBytes();
        final long totalSizeLimit = this.properties.getMaxTotalSize().toBytes();

        long totalSize = 0;
        long largestSize = 0;
        for (final Resource attachment : attachments) {
            final String filename = attachment.getFilename();

            final long attachmentSize;

            try {
                attachmentSize = attachment.contentLength();
            } catch (IOException e) {
                throw new SaveAttachmentException(
                    "Failed to get size of attachment: " + filename + ": " + e.getMessage(),
                    e
                );
            }

            if (attachmentSize > largestSize) {
                largestSize = attachmentSize;
            }
            totalSize += attachmentSize;
        }


        if (largestSize > singleSizeLimit) {
            throw new AttachmentTooLargeException(
                "Size of attachment exceeds the maximum allowed"
                    + " (" + largestSize + " > " + singleSizeLimit + ")"
            );
        }

        if (totalSize > totalSizeLimit) {
            throw new AttachmentTooLargeException(
                "Total size of attachments exceeds the maximum allowed"
                    + " (" + totalSize + " > " + totalSizeLimit + ")"
            );
        }

        this.meterRegistry.summary(LARGEST_SIZE_DISTRIBUTION).record(largestSize);
        this.meterRegistry.summary(TOTAL_SIZE_DISTRIBUTION).record(totalSize);
    }

    private Set uploadAllAttachments(
        @Nullable final String jobId,
        final Set attachments
    ) throws SaveAttachmentException {
        final AmazonS3 s3Client = this.s3ClientFactory.getClient(this.s3BaseURI);
        final String bundleId = UUID.randomUUID().toString();
        final String commonPrefix = this.s3BaseURI.getKey() + SLASH + bundleId + SLASH;

        log.debug(
            "Uploading {} attachments for job request with id {} to: {}",
            attachments.size(),
            jobId,
            commonPrefix
        );

        final Set attachmentURIs = Sets.newHashSet();

        for (final Resource attachment : attachments) {
            final String filename = attachment.getFilename();
            if (StringUtils.isBlank(filename)) {
                throw new SaveAttachmentException("Attachment filename is missing");
            }
            final String objectBucket = this.s3BaseURI.getBucket();
            final String objectKey = commonPrefix + filename;

            final ObjectMetadata metadata = new ObjectMetadata();
            URI attachmentURI = null;

            try (InputStream inputStream = attachment.getInputStream()) {
                // Prepare object
                metadata.setContentLength(attachment.contentLength());
                attachmentURI = new URI(S3, objectBucket, SLASH + objectKey, null);
                // Upload
                s3Client.putObject(
                    objectBucket,
                    objectKey,
                    inputStream,
                    metadata
                );

                // Add attachment URI to the set
                attachmentURIs.add(attachmentURI);

            } catch (IOException | SdkClientException | URISyntaxException e) {
                throw new SaveAttachmentException(
                    "Failed to upload attachment: " + attachmentURI + " - " + e.getMessage(),
                    e
                );
            }
        }

        return attachmentURIs;
    }
}