Update tests and refactor

Author: davidh44

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ByteArrayAsyncResponseTransformer.java

@@ -20,7 +20,6 @@
 import java.io.ByteArrayOutputStream;
 import java.nio.ByteBuffer;
 import java.util.concurrent.CompletableFuture;
-import java.util.function.Consumer;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
 import software.amazon.awssdk.annotations.SdkInternalApi;
@@ -69,9 +68,11 @@ public void exceptionOccurred(Throwable throwable) {
 
     @Override
     public SplitResult<ResponseT, ResponseBytes<ResponseT>> split(SplittingTransformerConfiguration splitConfig) {
+        // TODO - splitConfig not used - support this or log warning/update javdocs
+
         CompletableFuture<ResponseBytes<ResponseT>> future = new CompletableFuture<>();
         SdkPublisher<AsyncResponseTransformer<ResponseT, ResponseT>> transformer =
-            new ByteArraySplittingTransformer(this, future);
+            new ByteArraySplittingTransformer<>(this, future);
         return AsyncResponseTransformer.SplitResult.<ResponseT, ResponseBytes<ResponseT>>builder()
                                                    .publisher(transformer)
                                                    .resultFuture(future)
@@ -121,5 +122,4 @@ public void onComplete() {
             resultFuture.complete(baos.toByteArray());
         }
     }
-
 }

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ByteArraySplittingTransformer.java

@@ -27,6 +27,7 @@
 import java.util.concurrent.atomic.AtomicReference;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
+import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.core.ResponseBytes;
 import software.amazon.awssdk.core.async.AsyncResponseTransformer;
 import software.amazon.awssdk.core.async.SdkPublisher;
@@ -36,13 +37,13 @@
 import software.amazon.awssdk.utils.async.DelegatingBufferingSubscriber;
 import software.amazon.awssdk.utils.async.SimplePublisher;
 
+@SdkInternalApi
 public class ByteArraySplittingTransformer<ResponseT> implements SdkPublisher<AsyncResponseTransformer<ResponseT, ResponseT>> {
     private static final Logger log = Logger.loggerFor(ByteArraySplittingTransformer.class);
     private final AsyncResponseTransformer<ResponseT, ResponseBytes<ResponseT>> upstreamResponseTransformer;
-    private final SimplePublisher<AsyncResponseTransformer<ResponseT, ResponseT>> simplePublisher = new SimplePublisher<>();
     private final CompletableFuture<ResponseBytes<ResponseT>> resultFuture;
     private Subscriber<? super AsyncResponseTransformer<ResponseT, ResponseT>> downstreamSubscriber;
-    private final AtomicInteger partNumber = new AtomicInteger(0);
+    private final AtomicInteger onNextSignalsSent = new AtomicInteger(0);
     private final AtomicReference<ResponseT> responseT = new AtomicReference<>();
 
     /**
@@ -77,7 +78,8 @@ public class ByteArraySplittingTransformer<ResponseT> implements SdkPublisher<As
 
     private final Map<Integer, ByteBuffer> buffers;
 
-    public ByteArraySplittingTransformer(AsyncResponseTransformer<ResponseT, ResponseBytes<ResponseT>> upstreamResponseTransformer,
+    public ByteArraySplittingTransformer(AsyncResponseTransformer<ResponseT, ResponseBytes<ResponseT>>
+                                             upstreamResponseTransformer,
                                          CompletableFuture<ResponseBytes<ResponseT>> resultFuture) {
         this.upstreamResponseTransformer = upstreamResponseTransformer;
         this.resultFuture = resultFuture;
@@ -141,7 +143,7 @@ private boolean doEmit() {
             }
             if (outstandingDemand.get() > 0) {
                 demand = outstandingDemand.decrementAndGet();
-                downstreamSubscriber.onNext(new IndividualTransformer(partNumber.incrementAndGet()));
+                downstreamSubscriber.onNext(new IndividualTransformer(onNextSignalsSent.incrementAndGet()));
             }
         }
         return false;
@@ -195,15 +197,15 @@ private void handleSubscriptionCancel() {
         }
     }
 
-    private class IndividualTransformer implements AsyncResponseTransformer<ResponseT, ResponseT> {
-        private final int partNumber;
-        private ByteArrayAsyncResponseTransformer<ResponseT> delegate = new ByteArrayAsyncResponseTransformer<>();
+    private final class IndividualTransformer implements AsyncResponseTransformer<ResponseT, ResponseT> {
+        private final int onNextCount;
+        private final ByteArrayAsyncResponseTransformer<ResponseT> delegate = new ByteArrayAsyncResponseTransformer<>();
 
         private CompletableFuture<ResponseT> future;
-        private List<CompletableFuture<ResponseBytes<ResponseT>>> delegatePrepareFutures = new ArrayList<>();
+        private final List<CompletableFuture<ResponseBytes<ResponseT>>> delegatePrepareFutures = new ArrayList<>();
 
-        private IndividualTransformer(int partNumber) {
-            this.partNumber = partNumber;
+        private IndividualTransformer(int onNextCount) {
+            this.onNextCount = onNextCount;
         }
 
         @Override
@@ -213,7 +215,7 @@ public CompletableFuture<ResponseT> prepare() {
             CompletableFutureUtils.forwardExceptionTo(prepare, future);
             delegatePrepareFutures.add(prepare);
             return prepare.thenApply(responseTResponseBytes -> {
-                    buffers.put(partNumber, responseTResponseBytes.asByteBuffer());
+                buffers.put(onNextCount, responseTResponseBytes.asByteBuffer());
                 return responseTResponseBytes.response();
             });
         }

services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/S3MultipartClientGetObjectToBytesWiremockTest.java

@@ -29,14 +29,14 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
-import static org.junit.jupiter.api.Assertions.assertTrue;
+import static software.amazon.awssdk.services.s3.internal.multipart.utils.MultipartDownloadTestUtils.internalErrorBody;
+import static software.amazon.awssdk.services.s3.internal.multipart.utils.MultipartDownloadTestUtils.slowdownErrorBody;
 
 import com.github.tomakehurst.wiremock.http.Fault;
 import com.github.tomakehurst.wiremock.junit5.WireMockRuntimeInfo;
 import com.github.tomakehurst.wiremock.junit5.WireMockTest;
 import com.github.tomakehurst.wiremock.stubbing.Scenario;
 import java.net.URI;
-import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.ArrayList;
@@ -46,19 +46,15 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicBoolean;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
-import org.reactivestreams.Subscriber;
-import org.reactivestreams.Subscription;
 import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
 import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
 import software.amazon.awssdk.awscore.retry.AwsRetryStrategy;
 import software.amazon.awssdk.core.ResponseBytes;
 import software.amazon.awssdk.core.async.AsyncResponseTransformer;
-import software.amazon.awssdk.core.async.SdkPublisher;
 import software.amazon.awssdk.core.exception.SdkClientException;
 import software.amazon.awssdk.core.interceptor.Context;
 import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
@@ -213,23 +209,6 @@ public void getObject_5xxErrorResponses_shouldNotReuseInitialRequestId() {
         verify(0, getRequestedFor(urlEqualTo(String.format("/%s/%s?partNumber=3", BUCKET, KEY))));
     }
 
-    @Test
-    public void getObject_ioException_shouldRetryAndFail() {
-        String firstRequestId = UUID.randomUUID().toString();
-        String secondRequestId = UUID.randomUUID().toString();
-
-        stubIoError(1);
-        assertThatThrownBy(() -> multipartClient.getObject(b -> b.bucket(BUCKET).key(KEY),
-                                                           AsyncResponseTransformer.toBytes()).join())
-            .isInstanceOf(CompletionException.class)
-            .hasCauseInstanceOf(SdkClientException.class).hasMessageContaining("The connection was closed")
-                .hasStackTraceContaining("Error encountered during GetObjectRequest");
-
-        verify(MAX_ATTEMPTS, getRequestedFor(urlEqualTo(String.format("/%s/%s?partNumber=1", BUCKET, KEY))));
-        verify(0, getRequestedFor(urlEqualTo(String.format("/%s/%s?partNumber=2", BUCKET, KEY))));
-        verify(0, getRequestedFor(urlEqualTo(String.format("/%s/%s?partNumber=3", BUCKET, KEY))));
-    }
-
 
     @Test
     public void multipartDownload_200Response_shouldSucceed() {
@@ -254,16 +233,14 @@ public void multipartDownload_200Response_shouldSucceed() {
     @Test
     public void multipartDownload_secondPartNonRetryableError_shouldFail() {
         stub200SuccessPart1();
-        stubError(2, errorBody(String.valueOf(500), "Internal Error"));
+        stubError(2, internalErrorBody());
 
         CompletableFuture<ResponseBytes<GetObjectResponse>> future =
             multipartClient.getObject(GetObjectRequest.builder().bucket(BUCKET).key(KEY).build(),
                                       AsyncResponseTransformer.toBytes());
 
-        assertThatThrownBy(() -> future.join()).hasCauseInstanceOf(S3Exception.class)
-                                               .hasStackTraceContaining("Error encountered "
-                                                                                                         + "during "
-                                                                                           + "GetObjectRequest");
+        assertThatThrownBy(future::join).hasCauseInstanceOf(S3Exception.class)
+                                        .hasMessageContaining("We encountered an internal error. Please try again. (Service: S3, Status Code: 500");
 
         verify(1, getRequestedFor(urlEqualTo(String.format("/%s/%s?partNumber=1", BUCKET, KEY))));
         verify(MAX_ATTEMPTS, getRequestedFor(urlEqualTo(String.format("/%s/%s?partNumber=2", BUCKET, KEY))));
@@ -369,7 +346,22 @@ public void multipartDownload_503OnFirstPartAndSecondPart_shouldRetrySuccessfull
     }
 
     @Test
-    public void getObject_iOError_shouldRetrySuccessfully() {
+    public void getObject_ioExceptionOnly_shouldExhaustRetriesAndFail() {
+        stubIoError(1);
+        stub200SuccessPart2();
+        stub200SuccessPart3();
+        assertThatThrownBy(() -> multipartClient.getObject(b -> b.bucket(BUCKET).key(KEY),
+                                                           AsyncResponseTransformer.toBytes()).join())
+            .isInstanceOf(CompletionException.class)
+            .hasCauseInstanceOf(SdkClientException.class);
+
+        verify(MAX_ATTEMPTS, getRequestedFor(urlEqualTo(String.format("/%s/%s?partNumber=1", BUCKET, KEY))));
+        verify(0, getRequestedFor(urlEqualTo(String.format("/%s/%s?partNumber=2", BUCKET, KEY))));
+        verify(0, getRequestedFor(urlEqualTo(String.format("/%s/%s?partNumber=3", BUCKET, KEY))));
+    }
+
+    @Test
+    public void getObject_iOErrorThen200Response_shouldRetrySuccessfully() {
         String requestId = UUID.randomUUID().toString();
 
         stubFor(any(anyUrl())
@@ -405,32 +397,16 @@ public void getObject_iOError_shouldRetrySuccessfully() {
         assertEquals(requestId, finalRequestId);
     }
 
-    private String errorBody(String errorCode, String errorMessage) {
-        return "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n"
-               + "<Error>\n"
-               + "  <Code>" + errorCode + "</Code>\n"
-               + "  <Message>" + errorMessage + "</Message>\n"
-               + "</Error>";
-    }
-
-    private String internalErrorBody() {
-        return errorBody("InternalError", "We encountered an internal error. Please try again.");
-    }
-
-    private String slowdownErrorBody() {
-        return errorBody("SlowDown", "Please reduce your request rate.");
-    }
-
     private void stubError(int partNumber, String errorBody) {
-        stubFor(get(urlEqualTo(String.format("/%s/%s?partNumber=" + partNumber, BUCKET, KEY)))
+        stubFor(get(urlEqualTo(String.format("/%s/%s?partNumber=%d", BUCKET, KEY, partNumber)))
                     .willReturn(aResponse()
                                     .withHeader("x-amz-request-id", String.valueOf(UUID.randomUUID()))
                                     .withHeader("x-amz-mp-parts-count", String.valueOf(TOTAL_PARTS))
                                     .withStatus(500).withBody(errorBody)));
     }
 
     private void stubIoError(int partNumber) {
-        stubFor(get(urlEqualTo(String.format("/%s/%s?partNumber=" + partNumber, BUCKET, KEY)))
+        stubFor(get(urlEqualTo(String.format("/%s/%s?partNumber=%d", BUCKET, KEY, partNumber)))
                     .willReturn(aResponse()
                                     .withFault(Fault.CONNECTION_RESET_BY_PEER)));
     }
@@ -520,62 +496,4 @@ public void clear() {
             responses.clear();
         }
     }
-
-    /**
-     * Custom AsyncResponseTransformer that simulates an error occurring after onStream() has been called
-     */
-    private static final class StreamingErrorTransformer
-        implements AsyncResponseTransformer<GetObjectResponse, ResponseBytes<GetObjectResponse>> {
-
-        private final CompletableFuture<ResponseBytes<GetObjectResponse>> future = new CompletableFuture<>();
-        private final AtomicBoolean errorThrown = new AtomicBoolean();
-        private final AtomicBoolean onStreamCalled = new AtomicBoolean();
-
-        @Override
-        public CompletableFuture<ResponseBytes<GetObjectResponse>> prepare() {
-            return future;
-        }
-
-        @Override
-        public void onResponse(GetObjectResponse response) {
-            //
-        }
-
-        @Override
-        public void onStream(SdkPublisher<ByteBuffer> publisher) {
-            onStreamCalled.set(true);
-            publisher.subscribe(new Subscriber<ByteBuffer>() {
-                private Subscription subscription;
-
-                @Override
-                public void onSubscribe(Subscription s) {
-                    this.subscription = s;
-                    s.request(1);
-                }
-
-                @Override
-                public void onNext(ByteBuffer byteBuffer) {
-                    if (errorThrown.compareAndSet(false, true)) {
-                        future.completeExceptionally(new RuntimeException());
-                        subscription.cancel();
-                    }
-                }
-
-                @Override
-                public void onError(Throwable t) {
-                    future.completeExceptionally(t);
-                }
-
-                @Override
-                public void onComplete() {
-                    //
-                }
-            });
-        }
-
-        @Override
-        public void exceptionOccurred(Throwable throwable) {
-            future.completeExceptionally(throwable);
-        }
-    }
 }