Propagate exception to inputstream subscriber

Within InputStreamResponseTransformer, if exceptionOccurred is invoked
after we already have an InputStream subscriber, propate the exception
to that subscriber.

Author: dagnir

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

@@ -17,6 +17,8 @@
 
 import java.nio.ByteBuffer;
 import java.util.concurrent.CompletableFuture;
+import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.core.ResponseInputStream;
 import software.amazon.awssdk.core.SdkResponse;
@@ -35,6 +37,7 @@ public class InputStreamResponseTransformer<ResponseT extends SdkResponse>
 
     private volatile CompletableFuture<ResponseInputStream<ResponseT>> future;
     private volatile ResponseT response;
+    private volatile WaitForSubscribeOnErrorWrapper subscriber;
 
     @Override
     public CompletableFuture<ResponseInputStream<ResponseT>> prepare() {
@@ -51,17 +54,74 @@ public void onResponse(ResponseT response) {
     @Override
     public void onStream(SdkPublisher<ByteBuffer> publisher) {
         AbortableInputStreamSubscriber inputStreamSubscriber = AbortableInputStreamSubscriber.builder().build();
-        publisher.subscribe(inputStreamSubscriber);
+        WaitForSubscribeOnErrorWrapper waitForSubscribeSubscriber = new WaitForSubscribeOnErrorWrapper(inputStreamSubscriber);
+
+        this.subscriber = waitForSubscribeSubscriber;
+
+        publisher.subscribe(waitForSubscribeSubscriber);
         future.complete(new ResponseInputStream<>(response, inputStreamSubscriber));
     }
 
     @Override
     public void exceptionOccurred(Throwable error) {
         future.completeExceptionally(error);
+        if (subscriber != null) {
+            this.subscriber.onError(error);
+        }
     }
 
     @Override
     public String name() {
         return TransformerType.STREAM.getName();
     }
+
+    // Simple wrapper subscriber that ensures we don't forward the `onError` to the delegate until onSubscribe is called, to be
+    // compliant with the reactive streams spec. We use onError for forwarding the exception given to exceptionOccurred.
+    private static final class WaitForSubscribeOnErrorWrapper implements Subscriber<ByteBuffer> {
+        private final Object lock = new Object();
+        private final AbortableInputStreamSubscriber delegate;
+
+        private boolean subscribed = false;
+        private Throwable transformerException;
+
+
+        private WaitForSubscribeOnErrorWrapper(AbortableInputStreamSubscriber delegate) {
+            this.delegate = delegate;
+        }
+
+        @Override
+        public void onSubscribe(Subscription s) {
+            synchronized (lock) {
+                subscribed = true;
+                delegate.onSubscribe(s);
+
+                if (transformerException != null) {
+                    delegate.onError(transformerException);
+                    transformerException = null;
+                }
+            }
+        }
+
+        @Override
+        public void onNext(ByteBuffer byteBuffer) {
+            this.delegate.onNext(byteBuffer);
+        }
+
+        @Override
+        public void onError(Throwable t) {
+            synchronized (lock) {
+                if (subscribed) {
+                    delegate.onError(t);
+                } else {
+                    // We're not subscribed yet, save the exception for until we are.
+                    transformerException = t;
+                }
+            }
+        }
+
+        @Override
+        public void onComplete() {
+            this.delegate.onComplete();
+        }
+    }
 }

core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/async/InputStreamResponseTransformerTest.java

@@ -16,13 +16,20 @@
 package software.amazon.awssdk.core.internal.async;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatNoException;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.ByteBuffer;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Phaser;
+import java.util.concurrent.atomic.AtomicBoolean;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.reactivestreams.Subscription;
 import software.amazon.awssdk.core.ResponseInputStream;
 import software.amazon.awssdk.core.SdkResponse;
 import software.amazon.awssdk.core.async.SdkPublisher;
@@ -80,4 +87,59 @@ public void inputStreamArrayReadsAreFromPublisher() throws IOException {
         assertThat(data[2]).isEqualTo((byte) 2);
         assertThat(stream.read(data)).isEqualTo(-1);
     }
+
+    @Test
+    void exceptionOccurred_onStreamCalled_exceptionPropagatedToRead() {
+        ResponseInputStream<SdkResponse> inputStream = resultFuture.join();
+        transformer.exceptionOccurred(new IOException("I/O issue"));
+        assertThatThrownBy(inputStream::read).isInstanceOf(IOException.class);
+    }
+
+    @Test
+    void exceptionOccurred_onStreamNotCalled_doesNotFail() {
+        InputStreamResponseTransformer<SdkResponse> responseTransformer = new InputStreamResponseTransformer<>();
+        responseTransformer.prepare();
+        responseTransformer.onResponse(response);
+        assertThatNoException().isThrownBy(() -> responseTransformer.exceptionOccurred(new IOException("I/O issue")));
+    }
+
+    @Test
+    public void exceptionOccurred_onStreamCalled_onSubscribeNotCalled_doesNotThrow() throws Exception {
+        ExecutorService exec = Executors.newSingleThreadExecutor();
+
+        for (int i = 0; i < 1024; ++i) {
+            InputStreamResponseTransformer<SdkResponse> responseTransformer = new InputStreamResponseTransformer<>();
+            CompletableFuture<ResponseInputStream<SdkResponse>> cf = responseTransformer.prepare();
+            responseTransformer.onResponse(response);
+
+            AtomicBoolean onSubscribeCalled = new AtomicBoolean();
+
+            Phaser phaser = new Phaser(2);
+
+            responseTransformer.exceptionOccurred(new IOException("I/O issue"));
+            phaser.arrive();
+
+            exec.submit(() -> {
+                phaser.arriveAndAwaitAdvance();
+
+                responseTransformer.onStream(SdkPublisher.adapt(s -> {
+                    // the exception should not be forwarded to the inputstream subscriber if onSubscribe is not yet called
+
+                    s.onSubscribe(new Subscription() {
+                        @Override
+                        public void request(long n) {
+                        }
+
+                        @Override
+                        public void cancel() {
+                        }
+                    });
+                    onSubscribeCalled.set(true);
+                }));
+            }).get();
+
+            assertThatThrownBy(() -> cf.join().read()).hasMessageContaining("I/O issue");
+            assertThat(onSubscribeCalled.get()).isTrue();
+        }
+    }
 }