Fix race in FlatteningSubscriber where onClose called before items pa… (#3735)

* Fix race in FlatteningSubscriber where onClose called before items passed to onNext

* Add stochastic unit test for race condition in FlatteningSubscriber

---------

Co-authored-by: Dongie Agnir <[email protected]>

Author: stevenshan

.changes/next-release/bugfix-AWSSDKforJavav2-56b73c0.json

@@ -0,0 +1,6 @@
+{
+    "type": "bugfix",
+    "category": "AWS SDK for Java v2",
+    "contributor": "stevenshan",
+    "description": "Fix race in FlatteningSubscriber where onClose called before items passed to onNext. See [#3734](https://github.com/aws/aws-sdk-java-v2/issues/3734)"
+}

utils/src/main/java/software/amazon/awssdk/utils/async/FlatteningSubscriber.java

@@ -235,7 +235,7 @@ private void handleOnCompleteState() {
      * result is subject to change.
      */
     private boolean onCompleteNeeded() {
-        return allItems.isEmpty() && onCompleteCalledByUpstream && !terminalCallMadeDownstream;
+        return onCompleteCalledByUpstream && allItems.isEmpty() && !terminalCallMadeDownstream;
     }
 
     /**

utils/src/test/java/software/amazon/awssdk/utils/async/FlatteningSubscriberTest.java

@@ -18,10 +18,17 @@
 import static org.mockito.Mockito.times;
 
 import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mockito;
+import org.reactivestreams.Publisher;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
 
@@ -195,10 +202,67 @@ public void requestsFromDownstreamDoNothingAfterOnError() {
         Mockito.verifyNoMoreInteractions(mockUpstream, mockDelegate);
     }
 
+    @Test
+    public void stochastic_dataFlushedBeforeOnComplete() {
+        ExecutorService exec = Executors.newSingleThreadExecutor();
+        try {
+            for (int i = 0; i < 30_000_000; ++i) {
+                Publisher<List<String>> iterablePublisher = subscriber -> subscriber.onSubscribe(new Subscription() {
+                    @Override
+                    public void request(long l) {
+                        exec.submit(() -> {
+                            subscriber.onNext(Collections.singletonList("data"));
+                            subscriber.onComplete();
+                        });
+                    }
+
+                    @Override
+                    public void cancel() {
+                    }
+                });
+
+                AtomicInteger seen = new AtomicInteger(0);
+                CompletableFuture<Void> finished = new CompletableFuture<>();
+                FlatteningSubscriber<String> elementSubscriber = new FlatteningSubscriber<>(new Subscriber<String>() {
+                    @Override
+                    public void onSubscribe(Subscription subscription) {
+                        subscription.request(1);
+                    }
+
+                    @Override
+                    public void onNext(String s) {
+                        seen.incrementAndGet();
+                    }
+
+                    @Override
+                    public void onError(Throwable e) {
+                        finished.completeExceptionally(e);
+                    }
+
+                    @Override
+                    public void onComplete() {
+                        if (seen.get() != 1) {
+                            finished.completeExceptionally(
+                                new RuntimeException("Should have gotten 1 element before onComplete"));
+                        } else {
+                            finished.complete(null);
+                        }
+                    }
+                });
+
+                iterablePublisher.subscribe(elementSubscriber);
+
+                finished.join();
+            }
+        } finally {
+            exec.shutdown();
+        }
+    }
+
     private Subscription getDownstreamFromDelegate() {
         ArgumentCaptor<Subscription> subscriptionCaptor = ArgumentCaptor.forClass(Subscription.class);
         Mockito.verify(mockDelegate).onSubscribe(subscriptionCaptor.capture());
         return subscriptionCaptor.getValue();
     }
 
-}
+}