Fix intermittent NPE race in EmittingSubscription when cancel() nulls subscriber during emit

.changes/next-release/bugfix-AWSSDKforJavav2-c89f928.json

@@ -0,0 +1,6 @@
+{
+    "type": "bugfix",
+    "category": "AWS SDK for Java v2",
+    "contributor": "cthiebault",
+    "description": "Fixed an intermittent NullPointerException when downloading a single-part object with a multipart-enabled async S3 client. A race between the emit loop and subscription cancellation could dereference a cleared subscriber reference; the reference is now stable and reads are null-safe."
+}

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

@@ -35,7 +35,7 @@
 public final class EmittingSubscription<T> implements Subscription {
     private static final Logger log = Logger.loggerFor(EmittingSubscription.class);
 
-    private Subscriber<? super T> downstreamSubscriber;
+    private final Subscriber<? super T> downstreamSubscriber;
     private final AtomicBoolean emitting;
     private final AtomicLong outstandingDemand;
     private final Runnable onCancel;
@@ -74,7 +74,6 @@ public void request(long n) {
     @Override
     public void cancel() {
         isCancelled.set(true);
-        downstreamSubscriber = null;
         onCancel.run();
     }
 

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

@@ -41,7 +41,7 @@ public class FileAsyncResponseTransformerPublisher<T extends SdkResponse>
 
     private final Path path;
     private final FileTransformerConfiguration initialConfig;
-    private Subscriber<?> subscriber;
+    private volatile Subscriber<?> subscriber;
     private final AtomicLong transformerCount;
 
 
@@ -70,9 +70,22 @@ private AsyncResponseTransformer<T, T> createTransformer() {
     }
 
     private void onCancel() {
+        // Drop the subscriber reference on cancel, per Reactive Streams rule 3.13. Reads snapshot the volatile field and
+        // null-check it, so this clear cannot race into a NullPointerException.
         subscriber = null;
     }
 
+    /**
+     * Signals an error to the downstream subscriber if still present. Reads the volatile field once so a concurrent
+     * {@link #onCancel()} cannot cause a {@link NullPointerException}.
+     */
+    private void signalError(Throwable e) {
+        Subscriber<?> currentSubscriber = subscriber;
+        if (currentSubscriber != null) {
+            currentSubscriber.onError(e);
+        }
+    }
+
     /**
      * This is the AsyncResponseTransformer that will be used for each individual requests.
      * <p>
@@ -107,13 +120,13 @@ public void onResponse(T response) {
                     Optional<String> contentLength = response.sdkHttpResponse().firstMatchingHeader("content-length");
                     long transformerCount = FileAsyncResponseTransformerPublisher.this.transformerCount.get();
                     // Error out if content range header is missing and this is not the initial request
-                    if (subscriber != null && transformerCount > 0) {
-                        subscriber.onError(new IllegalStateException("Content range header is missing"));
+                    if (transformerCount > 0) {
+                        signalError(new IllegalStateException("Content range header is missing"));
                         return;
                     }
 
                     if (!contentLength.isPresent()) {
-                        subscriber.onError(new IllegalStateException("Content length header is missing"));
+                        signalError(new IllegalStateException("Content length header is missing"));
                         return;
                     }
                     String totalLength = contentLength.get();
@@ -125,10 +138,7 @@ public void onResponse(T response) {
             String contentRange = contentRangeOpt.get();
             Optional<Pair<Long, Long>> contentRangePair = ContentRangeParser.range(contentRange);
             if (!contentRangePair.isPresent()) {
-                if (subscriber != null) {
-                    IllegalStateException e = new IllegalStateException("Could not parse content range header " + contentRange);
-                    handleError(e);
-                }
+                handleError(new IllegalStateException("Could not parse content range header " + contentRange));
                 return;
             }
 
@@ -141,7 +151,7 @@ public void onResponse(T response) {
         }
 
         private void handleError(Throwable e) {
-            subscriber.onError(e);
+            signalError(e);
             future.completeExceptionally(e);
         }
 

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

@@ -0,0 +1,148 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.core.internal.async;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
+
+class EmittingSubscriptionTest {
+
+    @Test
+    void request_withPositiveDemand_emitsOneItemPerRequested() {
+        RecordingSubscriber subscriber = new RecordingSubscriber();
+        EmittingSubscription<Object> subscription = subscriptionFor(subscriber);
+
+        subscription.request(3);
+
+        assertThat(subscriber.onNextCount.get()).isEqualTo(3);
+        assertThat(subscriber.onError.get()).isNull();
+    }
+
+    @ParameterizedTest
+    @ValueSource(longs = {0L, -1L, Long.MIN_VALUE})
+    void request_withNonPositiveDemand_signalsIllegalArgumentException(long demand) {
+        RecordingSubscriber subscriber = new RecordingSubscriber();
+        EmittingSubscription<Object> subscription = subscriptionFor(subscriber);
+
+        subscription.request(demand);
+
+        assertThat(subscriber.onError.get()).isInstanceOf(IllegalArgumentException.class);
+        assertThat(subscriber.onNextCount.get()).isZero();
+    }
+
+    @Test
+    void cancel_thenRequest_emitsNothing() {
+        RecordingSubscriber subscriber = new RecordingSubscriber();
+        EmittingSubscription<Object> subscription = subscriptionFor(subscriber);
+
+        subscription.cancel();
+        subscription.request(5);
+
+        assertThat(subscriber.onNextCount.get()).isZero();
+        assertThat(subscriber.onError.get()).isNull();
+    }
+
+    /**
+     * A {@code cancel()} concurrent with an in-flight emit loop must not throw.
+     */
+    @Test
+    @Timeout(value = 30, unit = TimeUnit.SECONDS)
+    void cancel_concurrentWithEmit_neverThrowsNullPointerException() throws InterruptedException {
+        int rounds = 1_000;
+        // A large demand keeps the emit loop running long enough to overlap the cancel.
+        long demandPerRound = 10_000;
+
+        for (int round = 0; round < rounds; round++) {
+            RecordingSubscriber subscriber = new RecordingSubscriber();
+            EmittingSubscription<Object> subscription = subscriptionFor(subscriber);
+
+            AtomicReference<Throwable> emitFailure = new AtomicReference<>();
+            CountDownLatch startSignal = new CountDownLatch(1);
+
+            Thread emitter = new Thread(() -> {
+                awaitQuietly(startSignal);
+                try {
+                    subscription.request(demandPerRound);
+                } catch (Throwable t) {
+                    emitFailure.set(t);
+                }
+            });
+            Thread canceller = new Thread(() -> {
+                awaitQuietly(startSignal);
+                subscription.cancel();
+            });
+
+            emitter.start();
+            canceller.start();
+            startSignal.countDown();
+            emitter.join();
+            canceller.join();
+
+            assertThat(emitFailure.get())
+                .withFailMessage("cancel() racing the emit loop threw: %s", emitFailure.get())
+                .isNull();
+        }
+    }
+
+    private static EmittingSubscription<Object> subscriptionFor(Subscriber<Object> subscriber) {
+        return EmittingSubscription.<Object>builder()
+                                   .downstreamSubscriber(subscriber)
+                                   .onCancel(() -> { })
+                                   .supplier(Object::new)
+                                   .build();
+    }
+
+    private static void awaitQuietly(CountDownLatch latch) {
+        try {
+            latch.await();
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+        }
+    }
+
+    private static final class RecordingSubscriber implements Subscriber<Object> {
+        private final AtomicInteger onNextCount = new AtomicInteger();
+        private final AtomicReference<Throwable> onError = new AtomicReference<>();
+
+        @Override
+        public void onSubscribe(Subscription s) {
+        }
+
+        @Override
+        public void onNext(Object item) {
+            onNextCount.incrementAndGet();
+        }
+
+        @Override
+        public void onError(Throwable t) {
+            onError.compareAndSet(null, t);
+        }
+
+        @Override
+        public void onComplete() {
+        }
+    }
+}

services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/ParallelMultipartDownloaderSubscriber.java

@@ -198,7 +198,9 @@ public void onSubscribe(Subscription s) {
             return;
         }
         this.subscription = s;
-        subscription.request(maxInFlightParts);
+        synchronized (subscriptionLock) {
+            subscription.request(maxInFlightParts);
+        }
     }
 
     @Override
@@ -352,7 +354,9 @@ private boolean isMultipartObject(GetObjectResponse response) {
         if (response.partsCount() == null || response.partsCount() == 1) {
             // Single part object detected, skip multipart and complete everything now
             log.debug(() -> "Single Part object detected, skipping multipart download");
-            subscription.cancel();
+            synchronized (subscriptionLock) {
+                subscription.cancel();
+            }
             resultFuture.complete(response);
             return false;
         }