Merge pull request #1923 from smallrye/fix/uni-memoize-forwarding-race

fix: race condition in UniMemoizeOp on awaiters forwarding

Author: jponge

context-propagation/src/test/java/io/smallrye/mutiny/context/UniContextPropagationTest.java

@@ -5,13 +5,24 @@
 import static org.junit.jupiter.api.Assertions.assertNull;
 
 import java.time.Duration;
-import java.util.concurrent.*;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ForkJoinPool;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 
 import org.eclipse.microprofile.context.ThreadContext;
-import org.junit.jupiter.api.*;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
 
 import io.smallrye.context.CleanAutoCloseable;
 import io.smallrye.context.SmallRyeThreadContext;

implementation/src/main/java/io/smallrye/mutiny/operators/uni/UniMemoizeOp.java

@@ -2,6 +2,8 @@
 
 import static io.smallrye.mutiny.helpers.ParameterValidation.nonNull;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.BooleanSupplier;
@@ -47,27 +49,39 @@ public UniMemoizeOp(Uni<? extends I> upstream, BooleanSupplier invalidationReque
     @Override
     public void subscribe(UniSubscriber<? super I> subscriber) {
         nonNull(subscriber, "subscriber");
+
+        boolean shouldSubscribeUpstream = false;
+        Object cached = null;
+        boolean wasInCachingState = false;
+
         internalLock.lock();
-        checkForInvalidation();
-        switch (state) {
-            case INIT:
-                state = State.WAITING_FOR_UPSTREAM;
-                awaiters.add(subscriber);
-                currentContext = subscriber.context();
-                internalLock.unlock();
-                subscriber.onSubscribe(new MemoizedSubscription(subscriber));
-                upstream().subscribe().withSubscriber(this);
-                break;
-            case WAITING_FOR_UPSTREAM:
-                awaiters.add(subscriber);
-                internalLock.unlock();
-                subscriber.onSubscribe(new MemoizedSubscription(subscriber));
-                break;
-            case CACHING:
-                internalLock.unlock();
-                subscriber.onSubscribe(new MemoizedSubscription(subscriber));
-                forwardTo(subscriber);
-                break;
+        try {
+            checkForInvalidation(); // May throw an exception
+            switch (state) {
+                case INIT:
+                    state = State.WAITING_FOR_UPSTREAM;
+                    awaiters.add(subscriber);
+                    currentContext = subscriber.context();
+                    shouldSubscribeUpstream = true;
+                    break;
+                case WAITING_FOR_UPSTREAM:
+                    awaiters.add(subscriber);
+                    break;
+                case CACHING:
+                    cached = cachedResult;
+                    wasInCachingState = true;
+                    break;
+            }
+        } finally {
+            internalLock.unlock();
+        }
+
+        subscriber.onSubscribe(new MemoizedSubscription(subscriber));
+
+        if (shouldSubscribeUpstream) {
+            upstream().subscribe().withSubscriber(this);
+        } else if (wasInCachingState) {
+            forwardTo(subscriber, cached);
         }
     }
 
@@ -91,35 +105,49 @@ public void onSubscribe(UniSubscription subscription) {
     @Override
     public void onItem(I item) {
         internalLock.lock();
+        List<UniSubscriber<? super I>> toNotify = null;
         if (state == State.WAITING_FOR_UPSTREAM) {
             state = State.CACHING;
             cachedResult = item;
-            internalLock.unlock();
-            notifyAwaiters();
-        } else {
-            internalLock.unlock();
+            toNotify = gatherAwaiters();
+        }
+        internalLock.unlock();
+        if (toNotify != null) {
+            notifyAwaiters(toNotify, item);
+        }
+    }
+
+    private List<UniSubscriber<? super I>> gatherAwaiters() {
+        return new ArrayList<>(awaiters);
+    }
+
+    private void notifyAwaiters(List<UniSubscriber<? super I>> toNotify, Object result) {
+        for (UniSubscriber<? super I> awaiter : toNotify) {
+            forwardTo(awaiter, result);
         }
     }
 
     @Override
     public void onFailure(Throwable failure) {
         internalLock.lock();
+        List<UniSubscriber<? super I>> toNotify = null;
         if (state == State.WAITING_FOR_UPSTREAM) {
             state = State.CACHING;
             cachedResult = failure;
-            internalLock.unlock();
-            notifyAwaiters();
-        } else {
-            internalLock.unlock();
+            toNotify = gatherAwaiters();
+        }
+        internalLock.unlock();
+        if (toNotify != null) {
+            notifyAwaiters(toNotify, failure);
         }
     }
 
     @SuppressWarnings("unchecked")
-    private void forwardTo(UniSubscriber<? super I> subscriber) {
-        if (cachedResult instanceof Throwable) {
-            subscriber.onFailure((Throwable) cachedResult);
+    private void forwardTo(UniSubscriber<? super I> subscriber, Object result) {
+        if (result instanceof Throwable) {
+            subscriber.onFailure((Throwable) result);
         } else {
-            subscriber.onItem((I) cachedResult);
+            subscriber.onItem((I) result);
         }
     }
 
@@ -128,13 +156,6 @@ public Context context() {
         return currentContext;
     }
 
-    private void notifyAwaiters() {
-        UniSubscriber<? super I> awaiter;
-        while ((awaiter = awaiters.poll()) != null) {
-            forwardTo(awaiter);
-        }
-    }
-
     private class MemoizedSubscription implements UniSubscription {
 
         private final UniSubscriber<? super I> subscriber;

implementation/src/test/java/io/smallrye/mutiny/groups/UniMemoizeTest.java

@@ -590,4 +590,22 @@ void reproducer_1910() {
             resultFuture.cancel(true);
         }
     }
+
+    @RepeatedTest(1000)
+    public void testCachingRaceInNotification() {
+        AtomicInteger sub = new AtomicInteger();
+        AtomicInteger count = new AtomicInteger();
+
+        ExecutorService pool = ForkJoinPool.commonPool();
+        Uni<Integer> uni = Uni.createFrom().item(count::incrementAndGet)
+                .emitOn(pool)
+                .runSubscriptionOn(pool)
+                .memoize().until(() -> sub.incrementAndGet() > 2);
+
+        assertThat(uni.await().atMost(Duration.ofMillis(100))).isEqualTo(1);
+        assertThat(uni.await().atMost(Duration.ofMillis(100))).isEqualTo(1);
+        assertThat(uni.await().atMost(Duration.ofMillis(100))).isEqualTo(2);
+        assertThat(uni.await().atMost(Duration.ofMillis(100))).isEqualTo(3);
+        assertThat(uni.await().atMost(Duration.ofMillis(100))).isEqualTo(4);
+    }
 }