Fix race condition in CompletionStageSubscriber.onComplete()

andimarek · claude · andimarek · commit cdc892a72ad5 · 2026-03-06T14:10:58.000+10:00
The spec317_mustSupportACumulativePendingElementCountUpToLongMaxValue TCK
test intermittently hangs on CI (seen on Java 21). Root cause: a race
between onComplete() and finallyAfterEachPromiseFinishes() where the
queue-emptiness check and the onCompleteRun read/write are not atomic —
the last CF can complete between the check and the set, leaving the
runnable stranded forever.

Fix by extending the existing ReentrantLock scope to also cover
onCompleteRun, making the emptiness check and runnable set/claim a
single atomic operation on both sides. Runnables are executed outside
the lock to avoid deadlocks with user code.

Co-Authored-By: Claude Opus 4.6 &lt;noreply@anthropic.com&gt;
diff --git a/src/main/java/graphql/execution/reactive/CompletionStageOrderedSubscriber.java b/src/main/java/graphql/execution/reactive/CompletionStageOrderedSubscriber.java
@@ -32,8 +32,7 @@ protected void whenNextFinished(CompletionStage<D> completionStage, D d, Throwab
                 emptyInFlightQueueIfWeCan();
             }
         } finally {
-            boolean empty = inFlightQIsEmpty();
-            finallyAfterEachPromiseFinishes(empty);
+            finallyAfterEachPromiseFinishes();
         }
     }
 
diff --git a/src/main/java/graphql/execution/reactive/CompletionStageSubscriber.java b/src/main/java/graphql/execution/reactive/CompletionStageSubscriber.java
@@ -95,20 +95,19 @@ protected void whenNextFinished(CompletionStage<D> completionStage, D d, Throwab
                 downstreamSubscriber.onNext(d);
             }
         } finally {
-            boolean empty = removeFromInFlightQAndCheckIfEmpty(completionStage);
-            finallyAfterEachPromiseFinishes(empty);
+            removeFromInFlightQ(completionStage);
+            finallyAfterEachPromiseFinishes();
         }
     }
 
-    protected void finallyAfterEachPromiseFinishes(boolean isInFlightEmpty) {
-        //
-        // if the runOnCompleteOrErrorRun runnable is set, the upstream has
-        // called onComplete() already, but the CFs have not all completed
-        // yet, so we have to check whenever a CF completes
-        //
-        Runnable runOnCompleteOrErrorRun = onCompleteRun.get();
-        if (isInFlightEmpty && runOnCompleteOrErrorRun != null) {
-            onCompleteRun.set(null);
+    protected void finallyAfterEachPromiseFinishes() {
+        Runnable runOnCompleteOrErrorRun = lock.callLocked(() -> {
+            if (inFlightDataQ.isEmpty()) {
+                return onCompleteRun.getAndSet(null);
+            }
+            return null;
+        });
+        if (runOnCompleteOrErrorRun != null) {
             runOnCompleteOrErrorRun.run();
         }
     }
@@ -149,11 +148,15 @@ public void onComplete() {
     }
 
     private void onComplete(Runnable doneCodeToRun) {
-        if (inFlightQIsEmpty()) {
-            // run right now
-            doneCodeToRun.run();
-        } else {
+        boolean runNow = lock.callLocked(() -> {
+            if (inFlightDataQ.isEmpty()) {
+                return true;
+            }
             onCompleteRun.set(doneCodeToRun);
+            return false;
+        });
+        if (runNow) {
+            doneCodeToRun.run();
         }
     }
 
@@ -163,12 +166,8 @@ protected void offerToInFlightQ(CompletionStage<?> completionStage) {
         );
     }
 
-    private boolean removeFromInFlightQAndCheckIfEmpty(CompletionStage<?> completionStage) {
-        // uncontested locks in java are cheap - we don't expect much contention here
-        return lock.callLocked(() -> {
-            inFlightDataQ.remove(completionStage);
-            return inFlightDataQ.isEmpty();
-        });
+    private void removeFromInFlightQ(CompletionStage<?> completionStage) {
+        lock.runLocked(() -> inFlightDataQ.remove(completionStage));
     }
 
     /**
@@ -186,10 +185,6 @@ private void cancelInFlightFutures() {
         });
     }
 
-    protected boolean inFlightQIsEmpty() {
-        return lock.callLocked(inFlightDataQ::isEmpty);
-    }
-
     /**
      * The two terminal states are onComplete or onError
      *

Original file line number	Diff line number	Diff line change
`@@ -32,8 +32,7 @@ protected void whenNextFinished(CompletionStage<D> completionStage, D d, Throwab`
`32`	`32`	`emptyInFlightQueueIfWeCan();`
`33`	`33`	`}`
`34`	`34`	`} finally {`
`35`		`- boolean empty = inFlightQIsEmpty();`
`36`		`- finallyAfterEachPromiseFinishes(empty);`
	`35`	`+ finallyAfterEachPromiseFinishes();`
`37`	`36`	`}`
`38`	`37`	`}`
`39`	`38`