ConcatMap flaky test :

Change how the upstream request is deferred until the next downstream request

(cherry picked from commit a81c239)
Signed-off-by: Julien Ponge <[email protected]>

implementation/src/main/java/io/smallrye/mutiny/operators/multi/MultiConcatMapOp.java

@@ -1,5 +1,6 @@
 package io.smallrye.mutiny.operators.multi;
 
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
@@ -79,6 +80,8 @@ public static final class ConcatMapMainSubscriber<I, O> implements MultiSubscrib
 
         final ConcatMapInner<O> inner;
 
+        private final AtomicBoolean deferredUpstreamRequest = new AtomicBoolean(false);
+
         ConcatMapMainSubscriber(
                 MultiSubscriber<? super O> downstream,
                 Function<? super I, ? extends Publisher<? extends O>> mapper,
@@ -94,13 +97,16 @@ public void request(long n) {
             if (n > 0) {
                 if (state.compareAndSet(STATE_NEW, STATE_READY)) {
                     upstream.request(1);
-                    // No outstanding requests from inner, forward the request to upstream
-                } else if (state.get() == STATE_READY && inner.requested() == 0) {
+                }
+                if (deferredUpstreamRequest.compareAndSet(true, false)) {
                     upstream.request(1);
                 }
                 inner.request(n);
+                if (inner.requested() != 0L && deferredUpstreamRequest.compareAndSet(true, false)) {
+                    upstream.request(1);
+                }
             } else {
-                downstream.onFailure(new IllegalArgumentException("Invalid requests, must be greater than 0"));
+                downstream.onFailure(Subscriptions.getInvalidRequestException());
             }
         }
 
@@ -188,26 +194,17 @@ public synchronized void tryEmit(O value) {
         }
 
         public void innerComplete(long emitted) {
-            while (true) {
-                int state = this.state.get();
-                if (state == STATE_EMITTING) {
-                    if (this.state.compareAndSet(state, STATE_READY)) {
-                        // Inner completed but there are outstanding requests from inner,
-                        // Or the inner completed without producing any items
-                        // Request new item from upstream
-                        if (inner.requested() != 0L || emitted == 0) {
-                            upstream.request(1);
-                        }
-                        return;
-                    }
-                } else if (state == STATE_OUTER_TERMINATED) {
-                    if (this.state.compareAndSet(state, STATE_TERMINATED)) {
-                        terminateDownstream();
-                        return;
-                    }
+            if (this.state.compareAndSet(STATE_EMITTING, STATE_READY)) {
+                // Inner completed but there are outstanding requests from inner,
+                // Or the inner completed without producing any items
+                // Request new item from upstream
+                if (inner.requested() != 0L || emitted == 0) {
+                    upstream.request(1);
                 } else {
-                    return;
+                    deferredUpstreamRequest.set(true);
                 }
+            } else if (this.state.compareAndSet(STATE_OUTER_TERMINATED, STATE_TERMINATED)) {
+                terminateDownstream();
             }
         }
 
@@ -272,7 +269,7 @@ static final class ConcatMapInner<O> extends SwitchableSubscriptionSubscriber<O>
         /**
          * Downstream passed as {@code null} to {@link SwitchableSubscriptionSubscriber} as accessors are not reachable.
          * Effective downstream is {@code parent}.
-         * 
+         *
          * @param parent parent as downstream
          */
         ConcatMapInner(ConcatMapMainSubscriber<?, O> parent) {

implementation/src/test/java/io/smallrye/mutiny/operators/multi/MultiConcatMapNoPrefetchTest.java

@@ -69,9 +69,9 @@ void testTransformToMulti(boolean prefetch, int[] upstreamRequests) {
         Multi<Integer> result = upstream.onItem()
                 .transformToMulti(i -> Multi.createFrom().items(i, i))
                 .concatenate(prefetch);
-        AssertSubscriber<Integer> ts = new AssertSubscriber<>(5);
+        AssertSubscriber<Integer> ts = new AssertSubscriber<>();
         result.runSubscriptionOn(Infrastructure.getDefaultExecutor()).subscribe(ts);
-        ts.request(5);
+        ts.request(10);
         ts.awaitItems(10);
         await().untilAsserted(() -> assertThat(upstreamRequestCount).hasValue(upstreamRequests[0]));
         ts.request(1);
@@ -135,6 +135,7 @@ void testNoPrefetchWithConcatMapContainingEmpty() {
                 .concatenate(false);
         AssertSubscriber<Integer> ts = new AssertSubscriber<>(5);
         result.runSubscriptionOn(Infrastructure.getDefaultExecutor()).subscribe(ts);
+        ts.awaitSubscription();
         ts.request(5);
         ts.awaitItems(10);
         await().untilAsserted(() -> assertThat(upstreamRequestCount).hasValueGreaterThan(10));