This is an automated email from the ASF dual-hosted git repository.
gianm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/druid.git
The following commit(s) were added to refs/heads/master by this push:
new cb0e388b2ef Only do multipart uploads in RetryableS3OutputStream if
the bytes to push expand beyond the threshold to do so (#19893)
cb0e388b2ef is described below
commit cb0e388b2ef9a139ac524b537f3dc82ae4e2fc91
Author: Lucas Capistrant <[email protected]>
AuthorDate: Thu Aug 6 01:10:09 2026 -0500
Only do multipart uploads in RetryableS3OutputStream if the bytes to push
expand beyond the threshold to do so (#19893)
---
.../storage/s3/output/RetryableS3OutputStream.java | 98 +++++++++++++++++-----
.../s3/output/RetryableS3OutputStreamTest.java | 79 +++++++++++++++++
2 files changed, 156 insertions(+), 21 deletions(-)
diff --git
a/extensions-core/s3-extensions/src/main/java/org/apache/druid/storage/s3/output/RetryableS3OutputStream.java
b/extensions-core/s3-extensions/src/main/java/org/apache/druid/storage/s3/output/RetryableS3OutputStream.java
index f3e2075cb53..075f97ed1cc 100644
---
a/extensions-core/s3-extensions/src/main/java/org/apache/druid/storage/s3/output/RetryableS3OutputStream.java
+++
b/extensions-core/s3-extensions/src/main/java/org/apache/druid/storage/s3/output/RetryableS3OutputStream.java
@@ -24,6 +24,7 @@ import com.google.common.io.CountingOutputStream;
import it.unimi.dsi.fastutil.io.FastBufferedOutputStream;
import org.apache.druid.java.util.common.FileUtils;
import org.apache.druid.java.util.common.RetryUtils;
+import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.java.util.common.io.Closer;
import org.apache.druid.java.util.common.logger.Logger;
import org.apache.druid.java.util.emitter.service.ServiceMetricEvent;
@@ -37,6 +38,7 @@ import
software.amazon.awssdk.services.s3.model.CreateMultipartUploadRequest;
import software.amazon.awssdk.services.s3.model.CreateMultipartUploadResponse;
import software.amazon.awssdk.services.s3.model.UploadPartResponse;
+import javax.annotation.Nullable;
import java.io.Closeable;
import java.io.File;
import java.io.FileNotFoundException;
@@ -56,12 +58,14 @@ import java.util.concurrent.TimeUnit;
* <ol>
* <li>When new data is written, it first creates a chunk in local disk.</li>
* <li>New data is written to the local chunk until it is full.</li>
- * <li>When the chunk is full, it uploads the chunk to s3 using the multipart
upload API.
- * Since this happens synchronously, {@link #write(byte[], int, int)} can be
blocked until the upload is done.
+ * <li>When the chunk is full, a multipart upload is started if one is not
already in progress, and the chunk is
+ * queued for upload as a part. {@link #write(byte[], int, int)} can be
blocked while the upload queue is saturated.
* The upload can be retried when it fails with transient errors.</li>
- * <li>Once the upload succeeds, it creates a new chunk and continue.</li>
- * <li>When the stream is closed, it uploads the last chunk and finalize the
multipart upload.
- * {@link #close()} can be blocked until upload is done.</li>
+ * <li>Once the chunk is queued, it creates a new chunk and continue.</li>
+ * <li>When the stream is closed, what happens depends on whether any part was
uploaded. A stream that never filled
+ * a chunk is uploaded as a single object, costing one request instead of the
three a multipart upload needs;
+ * otherwise the last chunk is uploaded and the multipart upload is finalized.
{@link #close()} can be blocked
+ * until upload is done.</li>
* </ol>
* For compression format support, this output stream supports compression
formats if they are <i>concatenatable</i>,
* such as ZIP or GZIP.
@@ -81,10 +85,19 @@ public class RetryableS3OutputStream extends OutputStream
private final S3OutputConfig config;
private final ServerSideEncryptingAmazonS3 s3;
private final String s3Key;
- private final String uploadId;
private final File chunkStorePath;
private final long chunkSize;
+ /**
+ * Multipart upload ID, or null while the stream still fits in {@link
#currentChunk}. A multipart upload costs a
+ * create and a complete request on top of the part uploads themselves, so
it is only started once a chunk actually
+ * needs pushing; a stream that never fills a chunk is uploaded by {@link
#close()} as a single putObject.
+ *
+ * @see #initiateMultipartUploadIfNeeded()
+ */
+ @Nullable
+ private String uploadId;
+
private final byte[] singularBuffer = new byte[1];
// metric
@@ -123,18 +136,7 @@ public class RetryableS3OutputStream extends OutputStream
this.s3Key = s3Key;
this.uploadManager = uploadManager;
- final CreateMultipartUploadResponse result;
- try {
- CreateMultipartUploadRequest.Builder requestBuilder =
CreateMultipartUploadRequest.builder()
- .bucket(config.getBucket())
- .key(s3Key);
- result = S3Utils.retryS3Operation(() ->
s3.createMultipartUpload(requestBuilder), config.getMaxRetry());
- }
- catch (Exception e) {
- throw new IOException("Unable to start multipart upload", e);
- }
- this.uploadId = result.uploadId();
- this.chunkStorePath = new File(config.getTempDir(), uploadId +
UUID.randomUUID());
+ this.chunkStorePath = new File(config.getTempDir(),
UUID.randomUUID().toString());
FileUtils.mkdirp(this.chunkStorePath);
this.chunkSize = config.getChunkSize();
this.pushStopwatch = Stopwatch.createStarted();
@@ -196,12 +198,37 @@ public class RetryableS3OutputStream extends OutputStream
currentChunk.close();
final Chunk chunk = currentChunk;
if (chunk.length() > 0) {
+ initiateMultipartUploadIfNeeded();
futures.add(
uploadManager.queueChunkForUpload(s3, s3Key, chunk.id, chunk.file,
uploadId, config)
);
}
}
+ /**
+ * Starts the multipart upload the first time a chunk actually needs
pushing. Called only from
+ * {@link #pushCurrentChunk()}, so a stream that never fills a chunk issues
no create request at all.
+ */
+ private void initiateMultipartUploadIfNeeded() throws IOException
+ {
+ if (uploadId != null) {
+ return;
+ }
+ try {
+ final CreateMultipartUploadRequest.Builder requestBuilder =
CreateMultipartUploadRequest.builder()
+ .bucket(config.getBucket())
+ .key(s3Key);
+ final CreateMultipartUploadResponse result = S3Utils.retryS3Operation(
+ () -> s3.createMultipartUpload(requestBuilder),
+ config.getMaxRetry()
+ );
+ uploadId = result.uploadId();
+ }
+ catch (Exception e) {
+ throw new IOException("Unable to start multipart upload", e);
+ }
+ }
+
@Override
public void close() throws IOException
{
@@ -226,19 +253,48 @@ public class RetryableS3OutputStream extends OutputStream
uploadId
);
- final ServiceMetricEvent.Builder builder = new
ServiceMetricEvent.Builder().setDimension("uploadId", uploadId);
+ final ServiceMetricEvent.Builder builder =
+ new ServiceMetricEvent.Builder().setDimension("uploadId", uploadId
== null ? "none" : uploadId);
uploadManager.emitMetric(builder.setMetric(METRIC_TOTAL_UPLOAD_TIME,
totalUploadTimeMillis));
uploadManager.emitMetric(builder.setMetric(METRIC_TOTAL_UPLOAD_BYTES,
totalBytesUploaded));
});
try (Closer ignored = closer) {
if (!error) {
- pushCurrentChunk();
- completeMultipartUpload();
+ if (uploadId == null) {
+ // Everything written fits in the first chunk, so a single putObject
does the job that a create, an upload
+ // part and a complete would otherwise take.
+ putCurrentChunkAsWholeObject();
+ } else {
+ pushCurrentChunk();
+ completeMultipartUpload();
+ }
}
}
}
+ /**
+ * Uploads {@link #currentChunk} as a complete object. Only valid while
{@link #uploadId} is null, i.e. when no part
+ * has ever been pushed and the chunk therefore holds the entire stream.
+ */
+ private void putCurrentChunkAsWholeObject() throws IOException
+ {
+ currentChunk.close();
+ if (currentChunk.length() == 0) {
+ // Nothing was written, so there is no object to create.
+ return;
+ }
+ try {
+ S3Utils.retryS3Operation(
+ () -> s3.putObject(config.getBucket(), s3Key, currentChunk.file),
+ config.getMaxRetry()
+ );
+ }
+ catch (Exception e) {
+ throw new IOException(StringUtils.format("Unable to upload s3Key[%s]",
s3Key), e);
+ }
+ }
+
private void completeMultipartUpload()
{
final List<CompletedPart> pushResults = new ArrayList<>();
diff --git
a/extensions-core/s3-extensions/src/test/java/org/apache/druid/storage/s3/output/RetryableS3OutputStreamTest.java
b/extensions-core/s3-extensions/src/test/java/org/apache/druid/storage/s3/output/RetryableS3OutputStreamTest.java
index 07274757f09..48d943c2a7d 100644
---
a/extensions-core/s3-extensions/src/test/java/org/apache/druid/storage/s3/output/RetryableS3OutputStreamTest.java
+++
b/extensions-core/s3-extensions/src/test/java/org/apache/druid/storage/s3/output/RetryableS3OutputStreamTest.java
@@ -19,6 +19,7 @@
package org.apache.druid.storage.s3.output;
+import com.google.common.collect.ImmutableList;
import org.apache.druid.java.util.common.FileUtils;
import org.apache.druid.java.util.common.HumanReadableBytes;
import org.apache.druid.java.util.common.IOE;
@@ -42,6 +43,7 @@ import
software.amazon.awssdk.services.s3.model.CompleteMultipartUploadResponse;
import software.amazon.awssdk.services.s3.model.CompletedPart;
import software.amazon.awssdk.services.s3.model.CreateMultipartUploadRequest;
import software.amazon.awssdk.services.s3.model.CreateMultipartUploadResponse;
+import software.amazon.awssdk.services.s3.model.PutObjectResponse;
import software.amazon.awssdk.services.s3.model.UploadPartRequest;
import software.amazon.awssdk.services.s3.model.UploadPartResponse;
@@ -217,11 +219,71 @@ public class RetryableS3OutputStreamTest
s3.assertCancelled();
}
+ /**
+ * A multipart upload costs three S3 requests at minimum — create, upload
part, complete — so a stream small enough
+ * to need only one part should not use one. A task writes one object per
output partition, and every request lands
+ * on the same key prefix, so the per-object floor sets the burst rate
against that prefix.
+ */
+ @Test
+ public void testStreamFittingInOneChunkIsUploadedWithASinglePut() throws
IOException
+ {
+ chunkSize = 10;
+ ByteBuffer bb = ByteBuffer.allocate(Integer.BYTES);
+ try (RetryableS3OutputStream out =
+ new RetryableS3OutputStream(config, s3, path, s3UploadManager)) {
+ bb.putInt(1);
+ out.write(bb.array());
+ }
+
+ Assertions.assertEquals(0, s3.createMultipartUploadCount);
+ Assertions.assertEquals(0, s3.partRequests.size());
+ Assertions.assertNull(s3.completeRequest);
+ Assertions.assertEquals(ImmutableList.of((long) Integer.BYTES),
s3.putObjectContentLengths);
+ }
+
+ /**
+ * Once a stream outgrows a single chunk it must use multipart, costing one
create, one request per part, and one
+ * complete.
+ */
+ @Test
+ public void testStreamSpanningMultipleChunksUsesMultipartUpload() throws
IOException
+ {
+ chunkSize = 10;
+ ByteBuffer bb = ByteBuffer.allocate(Integer.BYTES);
+ try (RetryableS3OutputStream out =
+ new RetryableS3OutputStream(config, s3, path, s3UploadManager)) {
+ for (int i = 0; i < 25; i++) {
+ bb.clear();
+ bb.putInt(i);
+ out.write(bb.array());
+ }
+ }
+
+ Assertions.assertEquals(1, s3.createMultipartUploadCount);
+ Assertions.assertEquals(10, s3.partRequests.size());
+ Assertions.assertEquals(0, s3.putObjectContentLengths.size());
+ s3.assertCompleted(chunkSize, Integer.BYTES * 25);
+ }
+
+ /**
+ * A stream closed without any bytes written produces no object, and should
reach S3 not at all to do so.
+ */
+ @Test
+ public void testClosingWithoutWritingCreatesNoObject() throws IOException
+ {
+ chunkSize = 10;
+ new RetryableS3OutputStream(config, s3, path, s3UploadManager).close();
+
+ s3.assertNoRequestsIssued();
+ }
+
private static class TestAmazonS3 extends ServerSideEncryptingAmazonS3
{
private final List<UploadPartRequest> partRequests = new ArrayList<>();
+ private final List<Long> putObjectContentLengths = new ArrayList<>();
private int uploadFailuresLeft;
+ private int createMultipartUploadCount = 0;
private boolean cancelled = false;
@Nullable
private CompleteMultipartUploadRequest completeRequest;
@@ -236,11 +298,28 @@ public class RetryableS3OutputStreamTest
public CreateMultipartUploadResponse
createMultipartUpload(CreateMultipartUploadRequest.Builder requestBuilder)
throws SdkClientException
{
+ ++createMultipartUploadCount;
return CreateMultipartUploadResponse.builder()
.uploadId("uploadId")
.build();
}
+ @Override
+ public PutObjectResponse putObject(String bucket, String key, File file)
+ {
+ putObjectContentLengths.add(file.length());
+ return PutObjectResponse.builder().build();
+ }
+
+ private void assertNoRequestsIssued()
+ {
+ Assertions.assertEquals(0, createMultipartUploadCount);
+ Assertions.assertEquals(0, putObjectContentLengths.size());
+ Assertions.assertEquals(0, partRequests.size());
+ Assertions.assertNull(completeRequest);
+ Assertions.assertFalse(cancelled);
+ }
+
@Override
public UploadPartResponse uploadPart(UploadPartRequest.Builder
requestBuilder, RequestBody requestBody)
throws SdkClientException
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
