Tell ByteArrayOutputStream required capacity

Author: rtyley

src/main/java/com/madgag/aws/sdk/async/responsebytes/Main.java

@@ -2,7 +2,7 @@
 
 import com.google.common.hash.HashCode;
 import com.google.common.hash.Hashing;
-import com.madgag.aws.sdk.async.responsebytes.awssdk.core.async.AsyncResponseTransformerAlternative;
+import com.madgag.aws.sdk.async.responsebytes.awssdk.services.s3.AsyncS3ResponseTransformer;
 import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider;
 import software.amazon.awssdk.core.ResponseBytes;
 import software.amazon.awssdk.http.crt.AwsCrtAsyncHttpClient;
@@ -23,7 +23,7 @@ public static void main(String[] args) throws ExecutionException, InterruptedExc
             KnownS3Object knownS3Object = KnownS3Object.BIG;
             ResponseBytes<GetObjectResponse> responseBytes = s3Client.getObject(
                     GetObjectRequest.builder().bucket(knownS3Object.bucket()).key(knownS3Object.key()).build(),
-                    AsyncResponseTransformerAlternative.toBytes()
+                    AsyncS3ResponseTransformer.toBytes()
             ).get();
             System.out.println(responseBytes.response().contentLength());
             HashCode hashOfDownloadedData = Hashing.sha256().hashBytes(responseBytes.asByteBuffer());

src/main/java/com/madgag/aws/sdk/async/responsebytes/awssdk/core/async/AsyncResponseTransformerAlternative.java

@@ -4,8 +4,15 @@
 import software.amazon.awssdk.core.ResponseBytes;
 import software.amazon.awssdk.core.async.AsyncResponseTransformer;
 
+import java.util.Optional;
+import java.util.function.Function;
+
 public class AsyncResponseTransformerAlternative {
     public static <ResponseT> AsyncResponseTransformer<ResponseT, ResponseBytes<ResponseT>> toBytes() {
-        return new ByteArrayAsyncResponseTransformerAlternative<>();
+        return new ByteArrayAsyncResponseTransformerAlternative<>(Optional.empty());
+    }
+
+    public static <ResponseT> AsyncResponseTransformer<ResponseT, ResponseBytes<ResponseT>> toBytes(Function<ResponseT, Integer> f) {
+        return new ByteArrayAsyncResponseTransformerAlternative<>(Optional.of(f));
     }
 }

src/main/java/com/madgag/aws/sdk/async/responsebytes/awssdk/core/internal/async/ByteArrayAsyncResponseTransformerAlternative.java

@@ -15,11 +15,6 @@
 
 package com.madgag.aws.sdk.async.responsebytes.awssdk.core.internal.async;
 
-import static software.amazon.awssdk.utils.FunctionalUtils.invokeSafely;
-
-import java.io.ByteArrayOutputStream;
-import java.nio.ByteBuffer;
-import java.util.concurrent.CompletableFuture;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
 import software.amazon.awssdk.annotations.SdkInternalApi;
@@ -28,6 +23,14 @@
 import software.amazon.awssdk.core.async.SdkPublisher;
 import software.amazon.awssdk.utils.BinaryUtils;
 
+import java.io.ByteArrayOutputStream;
+import java.nio.ByteBuffer;
+import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import java.util.function.Function;
+
+import static software.amazon.awssdk.utils.FunctionalUtils.invokeSafely;
+
 /**
  * Implementation of {@link AsyncResponseTransformer} that dumps content into a byte array and supports further
  * conversions into types, like strings.
@@ -41,9 +44,14 @@
 public final class ByteArrayAsyncResponseTransformerAlternative<ResponseT> implements
         AsyncResponseTransformer<ResponseT, ResponseBytes<ResponseT>> {
 
+    private final Optional<Function<ResponseT, Integer>> knownSize;
     private volatile CompletableFuture<byte[]> cf;
     private volatile ResponseT response;
 
+    public ByteArrayAsyncResponseTransformerAlternative(Optional<Function<ResponseT, Integer>> knownSize) {
+        this.knownSize = knownSize;
+    }
+
     @Override
     public CompletableFuture<ResponseBytes<ResponseT>> prepare() {
         cf = new CompletableFuture<>();
@@ -57,7 +65,9 @@ public void onResponse(ResponseT response) {
 
     @Override
     public void onStream(SdkPublisher<ByteBuffer> publisher) {
-        publisher.subscribe(new BaosSubscriber(cf));
+        ByteArrayOutputStream baos =
+                knownSize.map(f -> new ByteArrayOutputStream(f.apply(response))).orElse(new ByteArrayOutputStream());
+        publisher.subscribe(new BaosSubscriber(cf, baos));
     }
 
     @Override
@@ -68,12 +78,13 @@ public void exceptionOccurred(Throwable throwable) {
     static class BaosSubscriber implements Subscriber<ByteBuffer> {
         private final CompletableFuture<byte[]> resultFuture;
 
-        private ByteArrayOutputStream baos = new ByteArrayOutputStream();
+        private ByteArrayOutputStream baos;
 
         private Subscription subscription;
 
-        BaosSubscriber(CompletableFuture<byte[]> resultFuture) {
+        BaosSubscriber(CompletableFuture<byte[]> resultFuture, ByteArrayOutputStream baos) {
             this.resultFuture = resultFuture;
+            this.baos = baos;
         }
 
         @Override

src/main/java/com/madgag/aws/sdk/async/responsebytes/awssdk/services/s3/AsyncS3ResponseTransformer.java

@@ -0,0 +1,13 @@
+package com.madgag.aws.sdk.async.responsebytes.awssdk.services.s3;
+
+import com.madgag.aws.sdk.async.responsebytes.awssdk.core.async.AsyncResponseTransformerAlternative;
+import software.amazon.awssdk.core.ResponseBytes;
+import software.amazon.awssdk.core.async.AsyncResponseTransformer;
+import software.amazon.awssdk.services.s3.model.GetObjectResponse;
+
+public class AsyncS3ResponseTransformer {
+
+    public static AsyncResponseTransformer<GetObjectResponse, ResponseBytes<GetObjectResponse>> toBytes() {
+        return AsyncResponseTransformerAlternative.toBytes(r -> r.contentLength().intValue());
+    }
+}