graphql-java · bbakerman · May 29, 2026 · May 29, 2026 · May 29, 2026 · May 29, 2026
diff --git a/src/main/java/graphql/execution/Async.java b/src/main/java/graphql/execution/Async.java
@@ -57,6 +57,20 @@ public interface CombinedBuilder<T> {
          */
         CompletableFuture<List<T>> await();
 
+        /**
+         * Like {@link #await()} but races against the given cancellation future. If the cancellation future
+         * completes before all the tracked futures complete, the already-completed futures will have their
+         * values harvested and returned as partial results (with {@code null} for incomplete entries)
+         * rather than completing exceptionally.
+         *
+         * <p>If {@code cancellationFuture} is {@code null}, this behaves identically to {@link #await()}.
+         *
+         * @param cancellationFuture a future that, when completed, signals cancellation; may be {@code null}
+         *
+         * @return a CompletableFuture to a List of values (possibly partial on cancellation)
+         */
+        CompletableFuture<List<T>> await(@Nullable CompletableFuture<Void> cancellationFuture);
+
         /**
          * This will return a {@code CompletableFuture<List<T>>} if ANY of the input values are async
          * otherwise it just return a materialised {@code List<T>}
@@ -104,6 +118,11 @@ public CompletableFuture<List<T>> await() {
             return typedEmpty();
         }
 
+        @Override
+        public CompletableFuture<List<T>> await(@Nullable CompletableFuture<Void> cancellationFuture) {
+            return await();
+        }
+
         @Override
         public Object awaitPolymorphic() {
             Assert.assertTrue(ix == 0, () -> "expected size was " + 0 + " got " + ix);
@@ -145,6 +164,33 @@ public CompletableFuture<List<T>> await() {
                 CompletableFuture<T> cf = (CompletableFuture<T>) value;
                 return cf.thenApply(Collections::singletonList);
             }
+            return materialisedValue();
+        }
+
+        @Override
+        public CompletableFuture<List<T>> await(@Nullable CompletableFuture<Void> cancellationFuture) {
+            commonSizeAssert();
+            if (cancellationFuture == null) {
+                return await();
+            }
+
+            if (value instanceof CompletableFuture) {
+                CompletableFuture<List<T>> overallResult = new CompletableFuture<>();
+                //noinspection unchecked
+                CompletableFuture<T> valueCF = (CompletableFuture<T>) value;
+                CompletableFuture.anyOf(valueCF, cancellationFuture).whenComplete((ignored, exception) -> {
+                    if (exception != null) {
+                        overallResult.completeExceptionally(exception);
+                        return;
+                    }
+                    overallResult.complete(Collections.singletonList(doneOrNull(valueCF)));
+                });
+                return overallResult;
+            }
+            return materialisedValue();
+        }
+
+        private @NonNull CompletableFuture<List<T>> materialisedValue() {
             //noinspection unchecked
             return CompletableFuture.completedFuture(Collections.singletonList((T) value));
         }
@@ -232,6 +278,51 @@ public CompletableFuture<List<T>> await() {
             return overallResult;
         }
 
+        @Override
+        public CompletableFuture<List<T>> await(@Nullable CompletableFuture<Void> cancellationFuture) {
+            commonSizeAssert();
+            if (cfCount == 0) {
+                return CompletableFuture.completedFuture(materialisedList(array));
+            }
+            if (cancellationFuture == null) {
+                return await();
+            }
+
+            CompletableFuture<List<T>> overallResult = new CompletableFuture<>();
+            CompletableFuture<Void> allOf = CompletableFuture.allOf(copyOnlyCFsToArray());
+
+            // Race "all field futures complete" against cancellation. The cancellation future always
+            // completes normally (see ExecutionInput#cancel), so anyOf can only complete exceptionally
+            // when a field future fails - in which case we propagate that failure.
+            CompletableFuture.anyOf(allOf, cancellationFuture).whenComplete((ignored, exception) -> {
+                if (exception != null) {
+                    overallResult.completeExceptionally(exception);
+                    return;
+                }
+                // Either every field future is done (allOf won) or cancellation won the race. In both
+                // cases we harvest whatever has completed; field futures that are not yet done become
+                // null. join() is safe here: if allOf is not done then no field future has failed (a
+                // failure would have completed allOf exceptionally and taken the branch above).
+                overallResult.complete(harvestResults(array));
+            });
+
+            return overallResult;
+        }
+
+        @SuppressWarnings("unchecked")
+        private List<T> harvestResults(Object[] array) {
+            List<T> results = new ArrayList<>(array.length);
+            for (Object object : array) {
+                if (object instanceof CompletableFuture) {
+                    CompletableFuture<T> cf = (CompletableFuture<T>) object;
+                    results.add(doneOrNull(cf));
+                } else {
+                    results.add((T) object);
+                }
+            }
+            return results;
+        }
+
         @SuppressWarnings("unchecked")
         @NonNull
         private CompletableFuture<T>[] copyOnlyCFsToArray() {
@@ -273,6 +364,10 @@ private void commonSizeAssert() {
 
     }
 
+    private static <T> @Nullable T doneOrNull(CompletableFuture<T> valueCF) {
+        return valueCF.isDone() ? valueCF.join() : null;
+    }
+
     @SuppressWarnings("unchecked")
     public static <T, U> CompletableFuture<List<U>> each(Collection<T> list, Function<T, Object> cfOrMaterialisedValueFactory) {
         Object l = eachPolymorphic(list, cfOrMaterialisedValueFactory);

diff --git a/src/test/groovy/graphql/execution/AsyncTest.groovy b/src/test/groovy/graphql/execution/AsyncTest.groovy
@@ -9,6 +9,7 @@ import java.util.function.BiFunction
 import java.util.function.Function
 
 import static java.util.concurrent.CompletableFuture.completedFuture
+import static java.util.concurrent.CompletableFuture.runAsync
 
 class AsyncTest extends Specification {
 
@@ -421,4 +422,244 @@ class AsyncTest extends Specification {
             return awaited
         }
     }
+
+    def "await with null cancelCF behaves like plain await"() {
+        when:
+        def asyncBuilder = Async.ofExpectedSize(3)
+        asyncBuilder.add(completedFuture("A"))
+        asyncBuilder.add(completedFuture("B"))
+        asyncBuilder.add(completedFuture("C"))
+        def list = asyncBuilder.await((CompletableFuture) null).join()
+
+        then:
+        list == ["A", "B", "C"]
+    }
+
+    def "await with cancelCF returns all results when all CFs complete before cancellation"() {
+        when:
+        def cancelCF = new CompletableFuture<Void>()
+        def asyncBuilder = Async.ofExpectedSize(3)
+        asyncBuilder.add(completedFuture("A"))
+        asyncBuilder.add(completedFuture("B"))
+        asyncBuilder.add(completedFuture("C"))
+        def list = asyncBuilder.await(cancelCF).join()
+
+        then:
+        list == ["A", "B", "C"]
+    }
+
+    def "await with cancelCF returns partial results when cancellation fires before all CFs complete"() {
+        when:
+        def cancelCF = new CompletableFuture<Void>()
+        def pending1 = new CompletableFuture<String>()
+        def pending2 = new CompletableFuture<String>()
+
+        def asyncBuilder = Async.ofExpectedSize(4)
+        asyncBuilder.add(completedFuture("A"))
+        asyncBuilder.add(pending1)
+        asyncBuilder.add(completedFuture("C"))
+        asyncBuilder.add(pending2)
+
+        def resultCF = asyncBuilder.await(cancelCF)
+
+        // cancel before pending CFs complete
+        cancelCF.complete(null)
+
+        def list = resultCF.join()
+
+        then:
+        list == ["A", null, "C", null]
+    }
+
+    def "await with cancelCF returns partial results with mixed objects and CFs"() {
+        when:
+        def cancelCF = new CompletableFuture<Void>()
+        def pending = new CompletableFuture<String>()
+
+        def asyncBuilder = Async.ofExpectedSize(4)
+        asyncBuilder.addObject("A")
+        asyncBuilder.add(completedFuture("B"))
+        asyncBuilder.add(pending)
+        asyncBuilder.addObject("D")
+
+        def resultCF = asyncBuilder.await(cancelCF)
+
+        cancelCF.complete(null)
+        def list = resultCF.join()
+
+        then:
+        list == ["A", "B", null, "D"]
+    }
+
+    def "await with cancelCF returns full results when all CFs complete even if cancelCF completes later"() {
+        when:
+        def cancelCF = new CompletableFuture<Void>()
+        def cf1 = new CompletableFuture<String>()
+        def cf2 = new CompletableFuture<String>()
+
+        def asyncBuilder = Async.ofExpectedSize(2)
+        asyncBuilder.add(cf1)
+        asyncBuilder.add(cf2)
+
+        def resultCF = asyncBuilder.await(cancelCF)
+
+        // complete all CFs before cancellation
+        cf1.complete("X")
+        cf2.complete("Y")
+
+        def list = resultCF.join()
+
+        then: "full results returned despite cancel firing after"
+        list == ["X", "Y"]
+
+        when: "cancelCF completes after all CFs - should not affect result"
+        cancelCF.complete(null)
+
+        then: "result is unchanged"
+        resultCF.join() == ["X", "Y"]
+    }
+
+    def "await with cancelCF propagates exception if a CF fails before cancellation"() {
+        when:
+        def cancelCF = new CompletableFuture<Void>()
+        def failing = new CompletableFuture<String>()
+        failing.completeExceptionally(new RuntimeException("boom"))
+
+        def asyncBuilder = Async.ofExpectedSize(2)
+        asyncBuilder.add(completedFuture("A"))
+        asyncBuilder.add(failing)
+
+        def resultCF = asyncBuilder.await(cancelCF)
+        resultCF.join()
+
+        then:
+        thrown(CompletionException)
+    }
+
+    def "await with cancelCF works with all materialised values"() {
+        when:
+        def cancelCF = new CompletableFuture<Void>()
+        def asyncBuilder = Async.ofExpectedSize(3)
+        asyncBuilder.addObject("A")
+        asyncBuilder.addObject("B")
+        asyncBuilder.addObject("C")
+        // make cancel happen soon but off thread
+        runAsync({ -> cancelCF.complete(null) })
+        def list = asyncBuilder.await(cancelCF).join()
+
+        then:
+        list == ["A", "B", "C"]
+    }
+
+    def "await with null cancelCF delegates to plain await"() {
+        when: "a many builder is awaited with a null cancellation future"
+        def asyncBuilder = Async.ofExpectedSize(2)
+        asyncBuilder.add(completedFuture("A"))
+        asyncBuilder.add(completedFuture("B"))
+        def list = asyncBuilder.await((CompletableFuture<Void>) null).join()
+
+        then: "it behaves identically to await() and returns all results"
+        list == ["A", "B"]
+    }
+
+    def "await with cancelCF on empty builder returns empty list"() {
+        when:
+        def cancelCF = new CompletableFuture<Void>()
+        def asyncBuilder = Async.ofExpectedSize(0)
+        def list = asyncBuilder.await(cancelCF).join()
+
+        then:
+        list == []
+    }
+
+    def "await with cancelCF on single builder can return completed values"() {
+        when: "single builder with a completed CF"
+        def cancelCF = new CompletableFuture<Void>()
+        def asyncBuilder = Async.ofExpectedSize(1)
+        asyncBuilder.add(completedFuture("A"))
+        // make cancel happen soon but off thread
+        runAsync({ -> cancelCF.complete(null) })
+        def list = asyncBuilder.await(cancelCF).join()
+
+        then: "result is returned normally"
+        list == ["A"]
+    }
+
+    def "await with cancelCF on single builder can return exceptions"() {
+        when: "single builder with a completed CF"
+        def cancelCF = new CompletableFuture<Void>()
+        def failing = new CompletableFuture<String>()
+        failing.completeExceptionally(new RuntimeException("boom"))
+
+        def asyncBuilder = Async.ofExpectedSize(1)
+        asyncBuilder.add(failing)
+
+        // make cancel happen soon but off thread
+        runAsync({ -> cancelCF.complete(null) })
+        def list = asyncBuilder.await(cancelCF).join()
+
+        then: "result is exceptional"
+        thrown(CompletionException)
+    }
+
+    def "await with null cancelCF on single builder will return completed values"() {
+        when: "single builder with a completed CF"
+        def asyncBuilder = Async.ofExpectedSize(1)
+        asyncBuilder.add(completedFuture("A"))
+        def list = asyncBuilder.await(null).join()
+
+        then: "result is returned normally"
+        list == ["A"]
+    }
+
+    def "await with null cancelCF on single builder will return materialised value"() {
+        when: "single builder with a completed CF"
+        def asyncBuilder = Async.ofExpectedSize(1)
+        asyncBuilder.addObject("A")
+        def list = asyncBuilder.await(null).join()
+
+        then: "result is returned normally"
+        list == ["A"]
+    }
+
+    def "await with cancelCF on single builder can be cancelled"() {
+        when: "single builder with a completed CF"
+        def cancelCF = new CompletableFuture<Void>()
+        def asyncBuilder = Async.ofExpectedSize(1)
+        asyncBuilder.add(new CompletableFuture<Object>())
+
+        // make cancel happen soon but off thread
+        runAsync({ -> cancelCF.complete(null) })
+
+        def list = asyncBuilder.await(cancelCF).join()
+
+        then: "the single value is null since it never completed"
+        list == [null]
+    }
+
+    def "await with cancelCF on single builder with materialised value returns it"() {
+        when:
+        def cancelCF = new CompletableFuture<Void>()
+        def asyncBuilder = Async.ofExpectedSize(1)
+        asyncBuilder.addObject("A")
+
+        // make cancel happen soon but off thread
+        runAsync({ -> cancelCF.complete(null) })
+
+        def list = asyncBuilder.await(cancelCF).join()
+
+        then:
+        list == ["A"]
+    }
+
+    def "await with null cancelCF on single builder with materialised value returns it"() {
+        when:
+        def asyncBuilder = Async.ofExpectedSize(1)
+        asyncBuilder.addObject("A")
+
+        def list = asyncBuilder.await(null).join()
+
+        then:
+        list == ["A"]
+    }
 }