This has a fix to buffer published events to keep them in order - tweaks on tests and how we handle things

bbakerman · bbakerman · commit 916232b5e73e · 2024-04-23T16:45:54.000+10:00
diff --git a/src/main/java/graphql/execution/reactive/CompletionStageMappingPublisher.java b/src/main/java/graphql/execution/reactive/CompletionStageMappingPublisher.java
@@ -66,14 +66,14 @@ public class CompletionStageSubscriber implements Subscriber<U> {
         final LockKit.ReentrantLock lock = new LockKit.ReentrantLock();
         final AtomicReference<Runnable> onCompleteOrErrorRun;
         final AtomicBoolean onCompleteOrErrorRunCalled;
-        final AtomicBoolean upstreamCancelled;
+        final AtomicBoolean subscriptionCancelled;
 
         public CompletionStageSubscriber(Subscriber<? super D> downstreamSubscriber) {
             this.downstreamSubscriber = downstreamSubscriber;
             inFlightDataQ = new ArrayDeque<>();
             onCompleteOrErrorRun = new AtomicReference<>();
             onCompleteOrErrorRunCalled = new AtomicBoolean(false);
-            upstreamCancelled = new AtomicBoolean(false);
+            subscriptionCancelled = new AtomicBoolean(false);
         }
 
 
@@ -123,6 +123,9 @@ private BiConsumer<D, Throwable> whenNextFinished() {
         }
 
         private void emptyInFlightQueueIfWeCan() {
+            if (subscriptionCancelled.get()) {
+                return;
+            }
             // done inside a memory lock, so we cant offer new CFs to the queue
             // until we have processed any completed ones from the start of
             // the queue.
@@ -162,7 +165,7 @@ private void emptyInFlightQueueIfWeCan() {
 
         private void handleThrowable(Throwable throwable) {
             // only do this once
-            if (upstreamCancelled.compareAndSet(false,true)) {
+            if (subscriptionCancelled.compareAndSet(false,true)) {
                 downstreamSubscriber.onError(throwable);
                 //
                 // Reactive semantics say that IF an exception happens on a publisher,
diff --git a/src/test/groovy/graphql/execution/reactive/CompletionStageMappingPublisherTest.groovy b/src/test/groovy/graphql/execution/reactive/CompletionStageMappingPublisherTest.groovy
@@ -10,8 +10,6 @@ import spock.lang.Specification
 import java.util.concurrent.CompletableFuture
 import java.util.concurrent.CompletionException
 import java.util.concurrent.CompletionStage
-import java.util.concurrent.CountDownLatch
-import java.util.concurrent.TimeUnit
 import java.util.function.Function
 
 class CompletionStageMappingPublisherTest extends Specification {
@@ -132,30 +130,29 @@ class CompletionStageMappingPublisherTest extends Specification {
         when:
         Publisher<Integer> rxIntegers = Flowable.range(0, 10)
 
-        CountDownLatch latch = new CountDownLatch(5)
-        CountDownLatch exceptionLatch = new CountDownLatch(1)
+        List<Runnable> completions = []
+
         def mapper = new Function<Integer, CompletionStage<String>>() {
             @Override
             CompletionStage<String> apply(Integer integer) {
 
-                if (integer == 5) {
-                    return CompletableFuture.supplyAsync {
-                        exceptionLatch.await(10_000, TimeUnit.SECONDS)
-                        sleep(100)
-                        throw new RuntimeException("Bang")
-                    }
+                if (integer < 5) {
+                    def cf = new CompletableFuture<String>()
+                    completions.add({ cf.complete(String.valueOf(integer)) })
+                    return cf
+                } else if (integer == 5) {
+                    def cf = new CompletableFuture<String>()
+                    completions.add({ cf.completeExceptionally(new RuntimeException("Bang")) })
+                    return cf
                 } else if (integer == 6) {
-                    return CompletableFuture.supplyAsync {
-                        latch.countDown()
-                        exceptionLatch.countDown() // number 5 is now alive
-                        return String.valueOf(integer)
+                    // complete 5,4,3,2,1 so we have to queue
+                    def reverse = completions.reverse()
+                    for (Runnable r : (reverse)) {
+                        r.run()
                     }
+                    return CompletableFuture.completedFuture(String.valueOf(integer))
                 } else {
-                    return CompletableFuture.supplyAsync {
-                        latch.countDown()
-                        latch.await(10_000, TimeUnit.SECONDS)
-                        return String.valueOf(integer)
-                    }
+                    return CompletableFuture.completedFuture(String.valueOf(integer))
                 }
             }
         }
