aws / aws/aws-sdk-java-v2

S3 getObject combined with AsyncResponseTransformer.toBytes() copies too much data

Open
#3,193 3 comments 2 reactions 0 assignees View on GitHub
feature-request p2
Dominant language
Java
Stars
2.6k
Forks
1k
Avg merge
2d 9h
Merged PRs (30d)
51

Description

### Describe the bug

Given the following piece of code:

```
client.getObject(
builder -> builder.bucket(bucket),
AsyncResponseTransformer.toBytes())
```

Several data copies do take place:

- Netty ByteBuf get's copied into a heap ByteBuffer. I understand this copy is done to present a common abstraction level for all transports. Yet maybe `ByteBuf::nioBuffers` could be relied on to minimize copies. I did not expore this option myself.
- `AsyncResponseTransformer.toBytes()` relies internally on an unsized ByteArrayOutputStream (that thus will expend many times) and the result will be copied when calling `toByteArray()` and another defensive copy is carried other when transforming the future.

### Expected Behavior

Minimize memory copies in such a scenario to minimize heap pressure.

### Current Behavior

![Screenshot from 2022-05-16 10-20-07](https://user-images.githubusercontent.com/6928740/168514448-cdee7863-a5c8-45f0-8171-5f101648982a.png)

### Reproduction Steps

```
client.getObject(
builder -> builder.bucket(bucket),
AsyncResponseTransformer.toBytes())
```

I used https://github.com/jvm-profiling-tools/async-profiler for the flame graphs.

### Possible Solution

- Remove needless defensive copies inside `ByteArrayAsyncResponseTransformer`. The byte array is passed to the caller, who then becomes responsible of it, and nobody else references the old byte array once the publisher completes. This can be an instant win coming at a very low price.
- Rely on `GetResponse::contentLength` to size a byte array and copy incoming buffers to it in place. This requires knowledge about response type... Thus this might be hardly doable in a generic fashion.

I tested this successfully:

```
package org.apache.james.blob.objectstorage.aws;

import java.nio.ByteBuffer;
import java.util.concurrent.CompletableFuture;

import org.reactivestreams.Subscriber;
import org.reactivestreams.Subscription;

import software.amazon.awssdk.core.ResponseBytes;
import software.amazon.awssdk.core.async.AsyncResponseTransformer;
import software.amazon.awssdk.core.async.SdkPublisher;
import software.amazon.awssdk.services.s3.model.GetObjectResponse;

/**
* Class copied for {@link software.amazon.awssdk.core.internal.async.ByteArrayAsyncResponseTransformer}
*
* Modified to take advantage of the content length of the get response in order to use a sized array
* upon content copy. This avoids the usage of a ByteArrayOutputStream that yields additional copies
* (resizing upon copy, copy of the resulting byte array).
*
* A defensive copy upon returning the result is also removed (responsibility transfered to the caller, no other usages)
*/
public class MinimalCopyBytesResponseTransformer implements AsyncResponseTransformer> {
private volatile CompletableFuture cf;
private volatile GetObjectResponse response;

public MinimalCopyBytesResponseTransformer() {

}

public CompletableFuture> prepare() {
this.cf = new CompletableFuture();
// Modifcation: Remove a defensive copy of the buffer upon completion: the caller is now the sole user of the array
return this.cf.thenApply(arr -> ResponseBytes.fromByteArrayUnsafe(response, arr));
}

public void onResponse(GetObjectResponse response) {
this.response = response;
}

public void onStream(SdkPublisher publisher) {
publisher.subscribe(new BaosSubscriber(this.cf, response.contentLength().intValue()));
}

public void exceptionOccurred(Throwable throwable) {
this.cf.completeExceptionally(throwable);
}

static class BaosSubscriber implements Subscriber {
private final CompletableFuture resultFuture;
// Modification: use a byte array instead of the ByteArrayInputStream and track position
private final byte[] buffer;
private int pos = 0;
private Subscription subscription;

BaosSubscriber(CompletableFuture resultFuture, int size) {
this.resultFuture = resultFuture;
this.buffer = new byte[size];
}

public void onSubscribe(Subscription s) {
if (this.subscription != null) {
s.cancel();
} else {
this.subscription = s;
this.subscription.request(9223372036854775807L);
}
}

public void onNext(ByteBuffer byteBuffer) {
// Modification: copy the response part in place into the result buffer and track position
int written = byteBuffer.remaining();
byteBuffer.get(buffer, pos, written);
pos += written;
this.subscription.request(1L);
}

public void onError(Throwable throwable) {
this.resultFuture.completeExceptionally(throwable);
}

public void onComplete() {
this.resultFuture.complete(this.buffer);
}
}
}
```

### Additional Information/Context

[Apache James](https://james.apache.org) relies on the S3 driver for email content, thus speed and heap impact is important for us as S3 getObject will be called on each email content read.

### AWS Java SDK version used

2.17.189

### JDK version used

openjdk version "11.0.15" 2022-04-19

### Operating System and version

Ubuntu 20.04.4 LTS

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.