|
| 1 | +/* |
| 2 | + * Licensed to Elasticsearch under one or more contributor |
| 3 | + * license agreements. See the NOTICE file distributed with |
| 4 | + * this work for additional information regarding copyright |
| 5 | + * ownership. Elasticsearch licenses this file to you under |
| 6 | + * the Apache License, Version 2.0 (the "License"); you may |
| 7 | + * not use this file except in compliance with the License. |
| 8 | + * You may obtain a copy of the License at |
| 9 | + * |
| 10 | + * http://www.apache.org/licenses/LICENSE-2.0 |
| 11 | + * |
| 12 | + * Unless required by applicable law or agreed to in writing, |
| 13 | + * software distributed under the License is distributed on an |
| 14 | + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| 15 | + * KIND, either express or implied. See the License for the |
| 16 | + * specific language governing permissions and limitations |
| 17 | + * under the License. |
| 18 | + */ |
| 19 | + |
| 20 | +package org.elasticsearch.cloud.aws.blobstore; |
| 21 | + |
| 22 | +import com.amazonaws.services.s3.model.*; |
| 23 | +import org.elasticsearch.common.unit.ByteSizeUnit; |
| 24 | +import org.elasticsearch.common.unit.ByteSizeValue; |
| 25 | + |
| 26 | +import java.io.ByteArrayInputStream; |
| 27 | +import java.io.IOException; |
| 28 | +import java.io.InputStream; |
| 29 | +import java.util.ArrayList; |
| 30 | +import java.util.List; |
| 31 | + |
| 32 | +/** |
| 33 | + * DefaultS3OutputStream uploads data to the AWS S3 service using 2 modes: single and multi part. |
| 34 | + * <p/> |
| 35 | + * When the length of the chunk is lower than buffer_size, the chunk is uploaded with a single request. |
| 36 | + * Otherwise multiple requests are made, each of buffer_size (except the last one which can be lower than buffer_size). |
| 37 | + * <p/> |
| 38 | + * Quick facts about S3: |
| 39 | + * <p/> |
| 40 | + * Maximum object size: 5 TB |
| 41 | + * Maximum number of parts per upload: 10,000 |
| 42 | + * Part numbers: 1 to 10,000 (inclusive) |
| 43 | + * Part size: 5 MB to 5 GB, last part can be < 5 MB |
| 44 | + * <p/> |
| 45 | + * See http://docs.aws.amazon.com/AmazonS3/latest/dev/qfacts.html |
| 46 | + * See http://docs.aws.amazon.com/AmazonS3/latest/dev/uploadobjusingmpu.html |
| 47 | + */ |
| 48 | +public class DefaultS3OutputStream extends S3OutputStream { |
| 49 | + |
| 50 | + private static final ByteSizeValue MULTIPART_MAX_SIZE = new ByteSizeValue(5, ByteSizeUnit.GB); |
| 51 | + |
| 52 | + /** |
| 53 | + * Multipart Upload API data |
| 54 | + */ |
| 55 | + private String multipartId; |
| 56 | + private int multipartChunks; |
| 57 | + private List<PartETag> multiparts; |
| 58 | + |
| 59 | + public DefaultS3OutputStream(S3BlobStore blobStore, String bucketName, String blobName, int bufferSizeInBytes, int numberOfRetries, boolean serverSideEncryption) { |
| 60 | + super(blobStore, bucketName, blobName, bufferSizeInBytes, numberOfRetries, serverSideEncryption); |
| 61 | + } |
| 62 | + |
| 63 | + @Override |
| 64 | + public void flush(byte[] bytes, int off, int len, boolean closing) throws IOException { |
| 65 | + if (len > MULTIPART_MAX_SIZE.getBytes()) { |
| 66 | + throw new IOException("Unable to upload files larger than " + MULTIPART_MAX_SIZE + " to Amazon S3"); |
| 67 | + } |
| 68 | + |
| 69 | + if (!closing) { |
| 70 | + if (len < getBufferSize()) { |
| 71 | + upload(bytes, off, len); |
| 72 | + } else { |
| 73 | + if (getFlushCount() == 0) { |
| 74 | + initializeMultipart(); |
| 75 | + } |
| 76 | + uploadMultipart(bytes, off, len, false); |
| 77 | + } |
| 78 | + } else { |
| 79 | + if (multipartId != null) { |
| 80 | + uploadMultipart(bytes, off, len, true); |
| 81 | + completeMultipart(); |
| 82 | + } else { |
| 83 | + upload(bytes, off, len); |
| 84 | + } |
| 85 | + } |
| 86 | + } |
| 87 | + |
| 88 | + /** |
| 89 | + * Upload data using a single request. |
| 90 | + * |
| 91 | + * @param bytes |
| 92 | + * @param off |
| 93 | + * @param len |
| 94 | + * @throws IOException |
| 95 | + */ |
| 96 | + private void upload(byte[] bytes, int off, int len) throws IOException { |
| 97 | + try (ByteArrayInputStream is = new ByteArrayInputStream(bytes, off, len)) { |
| 98 | + int retry = 0; |
| 99 | + while (retry < getNumberOfRetries()) { |
| 100 | + try { |
| 101 | + doUpload(getBlobStore(), getBucketName(), getBlobName(), is, len, isServerSideEncryption()); |
| 102 | + break; |
| 103 | + } catch (AmazonS3Exception e) { |
| 104 | + if (shouldRetry(e)) { |
| 105 | + is.reset(); |
| 106 | + retry++; |
| 107 | + } else { |
| 108 | + throw new IOException("Unable to upload object " + getBlobName(), e); |
| 109 | + } |
| 110 | + } |
| 111 | + } |
| 112 | + } |
| 113 | + } |
| 114 | + |
| 115 | + protected void doUpload(S3BlobStore blobStore, String bucketName, String blobName, InputStream is, int length, |
| 116 | + boolean serverSideEncryption) throws AmazonS3Exception { |
| 117 | + ObjectMetadata md = new ObjectMetadata(); |
| 118 | + if (serverSideEncryption) { |
| 119 | + md.setSSEAlgorithm(ObjectMetadata.AES_256_SERVER_SIDE_ENCRYPTION); |
| 120 | + } |
| 121 | + md.setContentLength(length); |
| 122 | + blobStore.client().putObject(bucketName, blobName, is, md); |
| 123 | + } |
| 124 | + |
| 125 | + private void initializeMultipart() { |
| 126 | + if (multipartId == null) { |
| 127 | + multipartId = doInitialize(getBlobStore(), getBucketName(), getBlobName(), isServerSideEncryption()); |
| 128 | + if (multipartId != null) { |
| 129 | + multipartChunks = 1; |
| 130 | + multiparts = new ArrayList<>(); |
| 131 | + } |
| 132 | + } |
| 133 | + } |
| 134 | + |
| 135 | + protected String doInitialize(S3BlobStore blobStore, String bucketName, String blobName, boolean serverSideEncryption) { |
| 136 | + InitiateMultipartUploadRequest request = new InitiateMultipartUploadRequest(bucketName, blobName); |
| 137 | + if (serverSideEncryption) { |
| 138 | + ObjectMetadata md = new ObjectMetadata(); |
| 139 | + md.setSSEAlgorithm(ObjectMetadata.AES_256_SERVER_SIDE_ENCRYPTION); |
| 140 | + request.setObjectMetadata(md); |
| 141 | + } |
| 142 | + return blobStore.client().initiateMultipartUpload(request).getUploadId(); |
| 143 | + } |
| 144 | + |
| 145 | + private void uploadMultipart(byte[] bytes, int off, int len, boolean lastPart) throws IOException { |
| 146 | + try (ByteArrayInputStream is = new ByteArrayInputStream(bytes, off, len)) { |
| 147 | + int retry = 0; |
| 148 | + while (retry < getNumberOfRetries()) { |
| 149 | + try { |
| 150 | + PartETag partETag = doUploadMultipart(getBlobStore(), getBucketName(), getBlobName(), multipartId, is, len, lastPart); |
| 151 | + multiparts.add(partETag); |
| 152 | + multipartChunks++; |
| 153 | + return; |
| 154 | + } catch (AmazonS3Exception e) { |
| 155 | + if (shouldRetry(e) && retry < getNumberOfRetries()) { |
| 156 | + is.reset(); |
| 157 | + retry++; |
| 158 | + } else { |
| 159 | + abortMultipart(); |
| 160 | + throw e; |
| 161 | + } |
| 162 | + } |
| 163 | + } |
| 164 | + } |
| 165 | + } |
| 166 | + |
| 167 | + protected PartETag doUploadMultipart(S3BlobStore blobStore, String bucketName, String blobName, String uploadId, InputStream is, |
| 168 | + int length, boolean lastPart) throws AmazonS3Exception { |
| 169 | + UploadPartRequest request = new UploadPartRequest() |
| 170 | + .withBucketName(bucketName) |
| 171 | + .withKey(blobName) |
| 172 | + .withUploadId(uploadId) |
| 173 | + .withPartNumber(multipartChunks) |
| 174 | + .withInputStream(is) |
| 175 | + .withPartSize(length) |
| 176 | + .withLastPart(lastPart); |
| 177 | + |
| 178 | + UploadPartResult response = blobStore.client().uploadPart(request); |
| 179 | + return response.getPartETag(); |
| 180 | + |
| 181 | + } |
| 182 | + |
| 183 | + private void completeMultipart() { |
| 184 | + int retry = 0; |
| 185 | + while (retry < getNumberOfRetries()) { |
| 186 | + try { |
| 187 | + doCompleteMultipart(getBlobStore(), getBucketName(), getBlobName(), multipartId, multiparts); |
| 188 | + multipartId = null; |
| 189 | + return; |
| 190 | + } catch (AmazonS3Exception e) { |
| 191 | + if (shouldRetry(e) && retry < getNumberOfRetries()) { |
| 192 | + retry++; |
| 193 | + } else { |
| 194 | + abortMultipart(); |
| 195 | + throw e; |
| 196 | + } |
| 197 | + } |
| 198 | + } |
| 199 | + } |
| 200 | + |
| 201 | + protected void doCompleteMultipart(S3BlobStore blobStore, String bucketName, String blobName, String uploadId, List<PartETag> parts) |
| 202 | + throws AmazonS3Exception { |
| 203 | + CompleteMultipartUploadRequest request = new CompleteMultipartUploadRequest(bucketName, blobName, uploadId, parts); |
| 204 | + blobStore.client().completeMultipartUpload(request); |
| 205 | + } |
| 206 | + |
| 207 | + private void abortMultipart() { |
| 208 | + if (multipartId != null) { |
| 209 | + try { |
| 210 | + doAbortMultipart(getBlobStore(), getBucketName(), getBlobName(), multipartId); |
| 211 | + } finally { |
| 212 | + multipartId = null; |
| 213 | + } |
| 214 | + } |
| 215 | + } |
| 216 | + |
| 217 | + protected void doAbortMultipart(S3BlobStore blobStore, String bucketName, String blobName, String uploadId) |
| 218 | + throws AmazonS3Exception { |
| 219 | + blobStore.client().abortMultipartUpload(new AbortMultipartUploadRequest(bucketName, blobName, uploadId)); |
| 220 | + } |
| 221 | + |
| 222 | + protected boolean shouldRetry(AmazonS3Exception e) { |
| 223 | + return e.getStatusCode() == 400 && "RequestTimeout".equals(e.getErrorCode()); |
| 224 | + } |
| 225 | +} |
0 commit comments