From e179924038c3d87ff7fda0b1842e13bda9639e67 Mon Sep 17 00:00:00 2001
From: Julien Ponge <jponge@redhat.com>
Date: Wed, 13 Dec 2023 23:40:02 +0100
Subject: [PATCH] fix: avoid a race condition in the new concatMap operator

Avoids a race-condition between the inner and outer subscribers on the
completion signals.
---
 .../operators/multi/MultiConcatMapOp.java      | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/implementation/src/main/java/io/smallrye/mutiny/operators/multi/MultiConcatMapOp.java b/implementation/src/main/java/io/smallrye/mutiny/operators/multi/MultiConcatMapOp.java
index ef925be53..0b511cb21 100644
--- a/implementation/src/main/java/io/smallrye/mutiny/operators/multi/MultiConcatMapOp.java
+++ b/implementation/src/main/java/io/smallrye/mutiny/operators/multi/MultiConcatMapOp.java
@@ -69,7 +69,7 @@ private enum State {
         private volatile State state = State.INIT;
         private volatile Subscription upstream;
         private Subscription currentUpstream;
-        private boolean upstreamHasCompleted = false;
+        private volatile boolean upstreamHasCompleted = false;
         private Throwable failure;
 
         private static final AtomicReferenceFieldUpdater<ConcatMapSubscriber, Subscription> UPSTREAM_UPDATER = AtomicReferenceFieldUpdater
@@ -148,10 +148,10 @@ private Throwable addFailure(Throwable failure) {
 
         @Override
         public void onCompletion() {
+            upstreamHasCompleted = true;
             if (state == State.CANCELLED) {
                 return;
             }
-            upstreamHasCompleted = true;
             if (STATE_UPDATER.compareAndSet(this, State.WAITING_NEXT_PUBLISHER, State.CANCELLED)
                     || STATE_UPDATER.compareAndSet(this, State.INIT, State.CANCELLED)) {
                 if (failure == null) {
@@ -243,18 +243,18 @@ public void onCompletion() {
                 if (state == State.CANCELLED) {
                     return;
                 }
-                if (!upstreamHasCompleted) {
-                    state = State.WAITING_NEXT_PUBLISHER;
-                    if (demand > 0L) {
-                        upstream.request(1L);
-                    }
-                } else {
-                    state = State.CANCELLED;
+                state = State.WAITING_NEXT_PUBLISHER;
+                if (upstreamHasCompleted && STATE_UPDATER.compareAndSet(ConcatMapSubscriber.this, State.WAITING_NEXT_PUBLISHER,
+                        State.CANCELLED)) {
                     if (failure != null) {
                         downstream.onFailure(failure);
                     } else {
                         downstream.onComplete();
                     }
+                } else {
+                    if (demand > 0L) {
+                        upstream.request(1L);
+                    }
                 }
             }