diff --git a/src/main/java/graphql/ExecutionInput.java b/src/main/java/graphql/ExecutionInput.java index d65920fcb..a83c741e7 100644 --- a/src/main/java/graphql/ExecutionInput.java +++ b/src/main/java/graphql/ExecutionInput.java @@ -11,7 +11,8 @@ import java.util.Locale; import java.util.Map; -import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.CompletableFuture; + import java.util.function.Consumer; import static graphql.Assert.assertNotNull; @@ -34,7 +35,7 @@ public class ExecutionInput { private final DataLoaderRegistry dataLoaderRegistry; private final ExecutionId executionId; private final Locale locale; - private final AtomicBoolean cancelled; + private final CompletableFuture cancellationFuture; private final boolean profileExecution; /** @@ -60,7 +61,7 @@ private ExecutionInput(Builder builder) { this.locale = builder.locale != null ? builder.locale : Locale.getDefault(); // always have a locale in place this.localContext = builder.localContext; this.extensions = builder.extensions; - this.cancelled = builder.cancelled; + this.cancellationFuture = builder.cancellationFuture; this.profileExecution = builder.profileExecution; } @@ -211,7 +212,7 @@ public Map getExtensions() { * @return true if the execution should be cancelled */ public boolean isCancelled() { - return cancelled.get(); + return cancellationFuture.isDone(); } /** @@ -219,7 +220,18 @@ public boolean isCancelled() { * and the graphql engine needs to be running on a thread to allow is to respect this flag. */ public void cancel() { - cancelled.set(true); + cancellationFuture.complete(null); + } + + /** + * Returns a {@link CompletableFuture} that completes when {@link #cancel()} is called. + * This allows async code to race against cancellation without polling. + * + * @return a future that completes (with null) when this execution is cancelled + */ + @Internal + public CompletableFuture getCancellationFuture() { + return cancellationFuture; } @@ -241,7 +253,7 @@ public ExecutionInput transform(Consumer builderConsumer) { .operationName(this.operationName) .context(this.context) .internalTransferContext(this.graphQLContext) - .internalTransferCancelBoolean(this.cancelled) + .internalTransferCancellationFuture(this.cancellationFuture) .localContext(this.localContext) .root(this.root) .dataLoaderRegistry(this.dataLoaderRegistry) @@ -306,7 +318,7 @@ public static class Builder { private DataLoaderRegistry dataLoaderRegistry = EMPTY_DATALOADER_REGISTRY; private Locale locale = Locale.getDefault(); private ExecutionId executionId; - private AtomicBoolean cancelled = new AtomicBoolean(false); + private CompletableFuture cancellationFuture = new CompletableFuture<>(); private boolean profileExecution; /** @@ -412,9 +424,8 @@ private Builder internalTransferContext(GraphQLContext graphQLContext) { return this; } - // hidden on purpose - private Builder internalTransferCancelBoolean(AtomicBoolean cancelled) { - this.cancelled = cancelled; + private Builder internalTransferCancellationFuture(CompletableFuture cancellationFuture) { + this.cancellationFuture = cancellationFuture; return this; } diff --git a/src/main/java/graphql/GraphQLUnusualConfiguration.java b/src/main/java/graphql/GraphQLUnusualConfiguration.java index e96b0b738..852d8b4bb 100644 --- a/src/main/java/graphql/GraphQLUnusualConfiguration.java +++ b/src/main/java/graphql/GraphQLUnusualConfiguration.java @@ -277,6 +277,14 @@ public ResponseMapFactoryConfig responseMapFactory() { return new ResponseMapFactoryConfig(this); } + /** + * @return an element that allows you to control cancellation behavior + */ + @ExperimentalApi + public CancellationConfig cancellation() { + return new CancellationConfig(this); + } + private void put(String named, Object value) { if (graphQLContext != null) { graphQLContext.put(named, value); @@ -410,4 +418,53 @@ public ResponseMapFactoryConfig setFactory(ResponseMapFactory factory) { return this; } } + + public static class CancellationConfig extends BaseContextConfig { + + /** + * The context key used to enable capturing partial results when an execution is cancelled. + */ + @ExperimentalApi + public static final String CAPTURE_PARTIAL_RESULTS_ON_CANCEL = "graphql.capturePartialResultsOnCancel"; + + /** + * The context key used to store the cancellation {@link java.util.concurrent.CompletableFuture} + * that completes when {@link ExecutionInput#cancel()} is called. + * This is only set when {@link #CAPTURE_PARTIAL_RESULTS_ON_CANCEL} is enabled. + */ + @Internal + public static final String CANCELLATION_FUTURE_KEY = CAPTURE_PARTIAL_RESULTS_ON_CANCEL + ".cancelFuture"; + + private CancellationConfig(GraphQLContextConfiguration contextConfig) { + super(contextConfig); + } + + /** + * Returns true if partial results should be captured when the execution is cancelled via + * {@link ExecutionInput#cancel()}. + * + * @return true if partial results capture on cancel is enabled + */ + @ExperimentalApi + public boolean isCapturePartialResultsOnCancelEnabled() { + return contextConfig.getBoolean(CAPTURE_PARTIAL_RESULTS_ON_CANCEL); + } + + /** + * When enabled, if {@link ExecutionInput#cancel()} is called during execution, the engine will + * return the partial results of any fields that have already completed, along with an error + * indicating the execution was cancelled. + *

+ * By default this is false and cancellation returns only the cancellation error with null data. + * + * @param enable true to enable capturing partial results on cancel + * + * @return this config object for chaining + */ + @ExperimentalApi + public CancellationConfig capturePartialResultsOnCancel(boolean enable) { + contextConfig.put(CAPTURE_PARTIAL_RESULTS_ON_CANCEL, enable); + return this; + } + } } diff --git a/src/main/java/graphql/execution/AbstractAsyncExecutionStrategy.java b/src/main/java/graphql/execution/AbstractAsyncExecutionStrategy.java index 25f2036cb..6dc606426 100644 --- a/src/main/java/graphql/execution/AbstractAsyncExecutionStrategy.java +++ b/src/main/java/graphql/execution/AbstractAsyncExecutionStrategy.java @@ -25,12 +25,25 @@ protected BiConsumer, Throwable> handleResults(ExecutionContext exe exception = executionContext.possibleCancellation(exception); if (exception != null) { + // A cancellation that fired after some fields already completed arrives here as a + // synthesised AbortExecutionException with a non-null results list (a real field + // failure always has null results). When partial capture is enabled we keep those + // results and attach the cancellation error; otherwise we report the error as usual. + if (results != null && capturePartialResults(executionContext)) { + executionContext.addError((AbortExecutionException) exception); + completeResultFuture(overallResult, executionContext, fieldNames, results); + return; + } handleNonNullException(executionContext, overallResult, exception); return; } - Map resolvedValuesByField = executionContext.getResponseMapFactory().createInsertionOrdered(fieldNames, results); - overallResult.complete(new ExecutionResultImpl(resolvedValuesByField, executionContext.getErrors())); + completeResultFuture(overallResult, executionContext, fieldNames, results); }; } + + protected void completeResultFuture(CompletableFuture overallResult, ExecutionContext executionContext, List fieldNames, List results) { + Map resolvedValuesByField = executionContext.getResponseMapFactory().createInsertionOrdered(fieldNames, results); + overallResult.complete(new ExecutionResultImpl(resolvedValuesByField, executionContext.getErrors())); + } } diff --git a/src/main/java/graphql/execution/AsyncExecutionStrategy.java b/src/main/java/graphql/execution/AsyncExecutionStrategy.java index 325a5829a..580176d48 100644 --- a/src/main/java/graphql/execution/AsyncExecutionStrategy.java +++ b/src/main/java/graphql/execution/AsyncExecutionStrategy.java @@ -4,15 +4,17 @@ import graphql.PublicApi; import graphql.execution.incremental.DeferredExecutionSupport; import org.jspecify.annotations.NullMarked; +import org.jspecify.annotations.Nullable; import graphql.execution.instrumentation.ExecutionStrategyInstrumentationContext; import graphql.execution.instrumentation.Instrumentation; import graphql.execution.instrumentation.parameters.InstrumentationExecutionStrategyParameters; import graphql.introspection.Introspection; import java.util.List; +import java.util.Objects; import java.util.Optional; import java.util.concurrent.CompletableFuture; -import java.util.function.BiConsumer; +import java.util.stream.Collectors; /** * The standard graphql execution strategy that runs fields asynchronously non-blocking. @@ -64,24 +66,23 @@ public CompletableFuture execute(ExecutionContext executionCont CompletableFuture overallResult = new CompletableFuture<>(); executionStrategyCtx.onDispatched(); - futures.await().whenComplete((completeValueInfos, throwable) -> { + CompletableFuture cancelCF = getCancellationFuture(executionContext); + futures.await(cancelCF).whenComplete((completeValueInfos, throwable) -> { List fieldsExecutedOnInitialResult = deferredExecutionSupport.getNonDeferredFieldNames(fieldNames); - BiConsumer, Throwable> handleResultsConsumer = handleResults(executionContext, fieldsExecutedOnInitialResult, overallResult); throwable = executionContext.possibleCancellation(throwable); - if (throwable != null) { - handleResultsConsumer.accept(null, throwable.getCause()); + if (throwable != null && !(completeValueInfos != null && capturePartialResults(executionContext))) { + // a genuine field failure, or a cancellation we cannot surface partial results for: + // there is nothing usable to return, so just report the error + handleResults(executionContext, fieldsExecutedOnInitialResult, overallResult).accept(null, throwable); return; } - Async.CombinedBuilder fieldValuesFutures = Async.ofExpectedSize(completeValueInfos.size()); - for (FieldValueInfo completeValueInfo : completeValueInfos) { - fieldValuesFutures.addObject(completeValueInfo.getFieldValueObject()); - } - dataLoaderDispatcherStrategy.executionStrategyOnFieldValuesInfo(completeValueInfos, parameters); - executionStrategyCtx.onFieldValuesInfo(completeValueInfos); - fieldValuesFutures.await().whenComplete(handleResultsConsumer); + // normal completion, or partial-results-on-cancel: completeValueInfos holds the + // FieldValueInfos that completed (with null entries for any cancelled before completing) + completeFieldValues(executionContext, parameters, executionStrategyCtx, dataLoaderDispatcherStrategy, + completeValueInfos, fieldsExecutedOnInitialResult, cancelCF, overallResult); }).exceptionally((ex) -> { // if there are any issues with combining/handling the field results, // complete the future at all costs and bubble up any thrown exception so @@ -96,4 +97,40 @@ public CompletableFuture execute(ExecutionContext executionCont return overallResult; } + /** + * Turns the completed {@link FieldValueInfo}s into field values and completes the {@code overallResult}. + *

+ * When partial-results-on-cancel is in play {@code completeValueInfos} may contain {@code null} + * entries for fields that were cancelled before they completed; those become {@code null} field + * values and are excluded from the instrumentation callbacks. + */ + @SuppressWarnings("FutureReturnValueIgnored") + private void completeFieldValues(ExecutionContext executionContext, + ExecutionStrategyParameters parameters, + ExecutionStrategyInstrumentationContext executionStrategyCtx, + DataLoaderDispatchStrategy dataLoaderDispatcherStrategy, + List completeValueInfos, + List fieldNames, + @Nullable CompletableFuture cancelCF, + CompletableFuture overallResult) { + Async.CombinedBuilder fieldValuesFutures = Async.ofExpectedSize(completeValueInfos.size()); + boolean hasNulls = false; + for (FieldValueInfo completeValueInfo : completeValueInfos) { + if (completeValueInfo != null) { + fieldValuesFutures.addObject(completeValueInfo.getFieldValueObject()); + } else { + hasNulls = true; + fieldValuesFutures.addObject((Object) null); + } + } + // null entries only occur for partial-results-on-cancel; the instrumentation callbacks should + // not see them, so filter only when needed and otherwise pass the list straight through + List valueInfosForInstrumentation = hasNulls + ? completeValueInfos.stream().filter(Objects::nonNull).collect(Collectors.toList()) + : completeValueInfos; + dataLoaderDispatcherStrategy.executionStrategyOnFieldValuesInfo(valueInfosForInstrumentation, parameters); + executionStrategyCtx.onFieldValuesInfo(valueInfosForInstrumentation); + fieldValuesFutures.await(cancelCF).whenComplete(handleResults(executionContext, fieldNames, overallResult)); + } + } diff --git a/src/main/java/graphql/execution/Execution.java b/src/main/java/graphql/execution/Execution.java index 448d1c738..df7a0f6ae 100644 --- a/src/main/java/graphql/execution/Execution.java +++ b/src/main/java/graphql/execution/Execution.java @@ -10,6 +10,7 @@ import graphql.GraphQL; import graphql.GraphQLContext; import graphql.GraphQLError; +import graphql.GraphQLUnusualConfiguration; import graphql.GraphQLException; import graphql.Internal; import graphql.Profiler; @@ -144,6 +145,13 @@ public CompletableFuture execute(Document document, GraphQLSche executionContext.getGraphQLContext().put(ResultNodesInfo.RESULT_NODES_INFO, executionContext.getResultNodesInfo()); + // When partial results on cancel is enabled, store the cancellation future in the context + // so that Async.Many#await(GraphQLContext) can race against it + if (graphQLContext.getBoolean(GraphQLUnusualConfiguration.CancellationConfig.CAPTURE_PARTIAL_RESULTS_ON_CANCEL)) { + graphQLContext.put(GraphQLUnusualConfiguration.CancellationConfig.CANCELLATION_FUTURE_KEY, + executionInput.getCancellationFuture()); + } + InstrumentationExecutionParameters parameters = new InstrumentationExecutionParameters( executionInput, graphQLSchema ); diff --git a/src/main/java/graphql/execution/ExecutionStrategy.java b/src/main/java/graphql/execution/ExecutionStrategy.java index 713da688d..b7bad18b6 100644 --- a/src/main/java/graphql/execution/ExecutionStrategy.java +++ b/src/main/java/graphql/execution/ExecutionStrategy.java @@ -5,6 +5,7 @@ import graphql.EngineRunningState; import graphql.ExecutionResult; import graphql.ExecutionResultImpl; +import graphql.GraphQLContext; import graphql.GraphQLError; import graphql.Internal; import graphql.PublicSpi; @@ -46,6 +47,7 @@ import graphql.schema.LightDataFetcher; import graphql.util.FpKit; import org.jspecify.annotations.NonNull; +import org.jspecify.annotations.Nullable; import java.util.ArrayList; import java.util.Collections; @@ -58,6 +60,8 @@ import java.util.function.Function; import java.util.function.Supplier; +import static graphql.GraphQLUnusualConfiguration.CancellationConfig.CANCELLATION_FUTURE_KEY; +import static graphql.GraphQLUnusualConfiguration.CancellationConfig.CAPTURE_PARTIAL_RESULTS_ON_CANCEL; import static graphql.execution.Async.exceptionallyCompletedFuture; import static graphql.execution.FieldCollectorParameters.newParameters; import static graphql.execution.FieldValueInfo.CompleteValueType.ENUM; @@ -218,6 +222,7 @@ protected Object executeObject(ExecutionContext executionContext, ExecutionStrat resolveObjectCtx.onDispatched(); + CompletableFuture cancelCF = getCancellationFuture(executionContext); Object fieldValueInfosResult = resolvedFieldFutures.awaitPolymorphic(); if (fieldValueInfosResult instanceof CompletableFuture) { CompletableFuture> fieldValueInfos = (CompletableFuture>) fieldValueInfosResult; @@ -232,7 +237,7 @@ protected Object executeObject(ExecutionContext executionContext, ExecutionStrat Async.CombinedBuilder resultFutures = fieldValuesCombinedBuilder(completeValueInfos); dataLoaderDispatcherStrategy.executeObjectOnFieldValuesInfo(completeValueInfos, parameters); resolveObjectCtx.onFieldValuesInfo(completeValueInfos); - resultFutures.await().whenComplete(handleResultsConsumer); + resultFutures.await(cancelCF).whenComplete(handleResultsConsumer); }).exceptionally((ex) -> { // if there are any issues with combining/handling the field results, // complete the future at all costs and bubble up any thrown exception so @@ -273,19 +278,52 @@ protected Object executeObject(ExecutionContext executionContext, ExecutionStrat return resultFutures; } + protected boolean capturePartialResults(ExecutionContext executionContext) { + return executionContext.getGraphQLContext().getBoolean(CAPTURE_PARTIAL_RESULTS_ON_CANCEL); + } + + /** + * Returns the cancellation future from the execution context if partial result capture is enabled, + * or {@code null} otherwise. This is used to pass into {@link Async.CombinedBuilder#await(CompletableFuture)} + * so that the generic utility can race against cancellation without needing to know about GraphQL context. + * + * @param executionContext the execution context + * + * @return the cancellation future, or {@code null} if partial results on cancel is not enabled + */ + protected @Nullable CompletableFuture getCancellationFuture(ExecutionContext executionContext) { + GraphQLContext graphQLContext = executionContext.getGraphQLContext(); + if (!graphQLContext.getBoolean(CAPTURE_PARTIAL_RESULTS_ON_CANCEL)) { + return null; + } + return graphQLContext.get(CANCELLATION_FUTURE_KEY); + } + private BiConsumer, Throwable> buildFieldValueMap(List fieldNames, CompletableFuture> overallResult, ExecutionContext executionContext) { return (List results, Throwable exception) -> { exception = executionContext.possibleCancellation(exception); if (exception != null) { + Throwable cause = exception instanceof CompletionException ? exception.getCause() : exception; + if (cause instanceof AbortExecutionException && results != null + && capturePartialResults(executionContext)) { + // partial results mode: results already has harvested values from completed CFs + executionContext.addError((AbortExecutionException) cause); + completeFieldValueMap(overallResult, executionContext, fieldNames, results); + return; + } handleValueException(overallResult, exception, executionContext); return; } - Map resolvedValuesByField = executionContext.getResponseMapFactory().createInsertionOrdered(fieldNames, results); - overallResult.complete(resolvedValuesByField); + completeFieldValueMap(overallResult, executionContext, fieldNames, results); }; } + protected void completeFieldValueMap(CompletableFuture> overallResult, ExecutionContext executionContext, List fieldNames, List results) { + Map resolvedValuesByField = executionContext.getResponseMapFactory().createInsertionOrdered(fieldNames, results); + overallResult.complete(resolvedValuesByField); + } + DeferredExecutionSupport createDeferredExecutionSupport(ExecutionContext executionContext, ExecutionStrategyParameters parameters) { MergedSelectionSet fields = parameters.getFields(); diff --git a/src/test/groovy/graphql/ExecutionInputTest.groovy b/src/test/groovy/graphql/ExecutionInputTest.groovy index 23c06ea9c..206d3089f 100644 --- a/src/test/groovy/graphql/ExecutionInputTest.groovy +++ b/src/test/groovy/graphql/ExecutionInputTest.groovy @@ -451,6 +451,300 @@ class ExecutionInputTest extends Specification { } + def "capturePartialResultsOnCancel configuration is accessible via unusualConfiguration"() { + // Smoke test: verify the configuration API works correctly + + def sdl = ''' + type Query { + field1 : String + field2 : String + } + ''' + + DataFetcher df = { DataFetchingEnvironment env -> "value" } + + def fetcherMap = ["Query": ["field1": df, "field2": df]] + def schema = TestUtil.schema(sdl, fetcherMap) + def graphQL = GraphQL.newGraphQL(schema).build() + + when: "capturePartialResultsOnCancel is enabled and execution completes normally" + ExecutionInput executionInput = ExecutionInput.newExecutionInput() + .query("{ field1 field2 }") + .graphQLContext({ c -> + GraphQL.unusualConfiguration(c).cancellation().capturePartialResultsOnCancel(true) + }) + .build() + + def er = graphQL.execute(executionInput) + + then: "normal results are returned unchanged" + er.errors.isEmpty() + er.data == [field1: "value", field2: "value"] + + when: "capturePartialResultsOnCancel is enabled but cancel is called before execution" + ExecutionInput cancelledInput = ExecutionInput.newExecutionInput() + .query("{ field1 field2 }") + .graphQLContext({ c -> + GraphQL.unusualConfiguration(c).cancellation().capturePartialResultsOnCancel(true) + }) + .build() + cancelledInput.cancel() + + er = graphQL.execute(cancelledInput) + + then: "cancel error is returned" + !er.errors.isEmpty() + er.errors.any { it["message"].contains("Execution has been asked to be cancelled") } + } + + def "capturePartialResultsOnCancel returns fast field value when slow field is cancelled"() { + // The partial-results harvest happens at the fieldValuesFutures.await(graphQLContext) level. + // FieldValueInfo.getFieldValueObject() for async object fields is a CompletableFuture. + // When the slow object's CF hasn't completed yet and cancel fires, await(graphQLContext) + // harvests the already-done fast CF and returns partial data. + def sdl = ''' + type Query { + fast : Inner + slow : Inner + } + type Inner { + value : String + } + ''' + + CountDownLatch slowStarted = new CountDownLatch(1) + CountDownLatch slowRelease = new CountDownLatch(1) + + DataFetcher fastInnerDf = { DataFetchingEnvironment env -> + return CompletableFuture.completedFuture([value: "fast-value"]) + } + + DataFetcher slowInnerDf = { DataFetchingEnvironment env -> + return CompletableFuture.supplyAsync { + slowStarted.countDown() + slowRelease.await() + return [value: "slow-value"] + } + } + + DataFetcher innerValueDf = { DataFetchingEnvironment env -> + return env.source["value"] + } + + def fetcherMap = [ + "Query": ["fast": fastInnerDf, "slow": slowInnerDf], + "Inner": ["value": innerValueDf] + ] + def schema = TestUtil.schema(sdl, fetcherMap) + def graphQL = GraphQL.newGraphQL(schema).build() + + when: + ExecutionInput executionInput = ExecutionInput.newExecutionInput() + .query("{ fast { value } slow { value } }") + .graphQLContext({ c -> + GraphQL.unusualConfiguration(c).cancellation().capturePartialResultsOnCancel(true) + }) + .build() + + def cf = graphQL.executeAsync(executionInput) + + // wait for slow DF to have started; by this point fast has already completed + slowStarted.await() + + // cancel while slow is still blocked; fast's fieldValue CF is already done + executionInput.cancel() + + // unblock slow so it doesn't hang forever + slowRelease.countDown() + + await().atMost(Duration.ofSeconds(10)).until({ -> cf.isDone() }) + def er = cf.join() + + then: + !cf.isCompletedExceptionally() + !er.errors.isEmpty() + er.errors.any { it["message"].contains("Execution has been asked to be cancelled") } + // fast field completed before cancel — its data should be present + er.data != null + er.data["fast"] != null + er.data["fast"]["value"] == "fast-value" + // slow field was cancelled — should be null + er.data["slow"] == null + } + + def "without capturePartialResultsOnCancel only the cancel error is returned"() { + def sdl = ''' + type Query { + fast : String + slow : String + } + ''' + + CountDownLatch fastDone = new CountDownLatch(1) + CountDownLatch slowLatch = new CountDownLatch(1) + + DataFetcher fastDf = { DataFetchingEnvironment env -> + return CompletableFuture.supplyAsync { + fastDone.countDown() + return "fast-value" + } + } + + DataFetcher slowDf = { DataFetchingEnvironment env -> + return CompletableFuture.supplyAsync { + slowLatch.await() + return "slow-value" + } + } + + def fetcherMap = ["Query": ["fast": fastDf, "slow": slowDf]] + def schema = TestUtil.schema(sdl, fetcherMap) + def graphQL = GraphQL.newGraphQL(schema).build() + + when: + ExecutionInput executionInput = ExecutionInput.newExecutionInput() + .query("{ fast slow }") + .build() + + def cf = graphQL.executeAsync(executionInput) + + fastDone.await() + executionInput.cancel() + slowLatch.countDown() + + await().atMost(Duration.ofSeconds(10)).until({ -> cf.isDone() }) + def er = cf.join() + + then: + !cf.isCompletedExceptionally() + !er.errors.isEmpty() + er.errors[0]["message"] == "Execution has been asked to be cancelled" + // no partial data - old behaviour + er.data == null + } + + def "capturePartialResultsOnCancel returns partial data when nested object sub-field is cancelled"() { + // This exercises the buildFieldValueMap lambda inside ExecutionStrategy.executeObject() + // for a nested object type. The parent DF returns a map synchronously so the nested + // object's executeObject starts immediately. Inside the nested executeObject, 'fast' + // completes while 'slow' blocks. Cancellation fires at the await(cancelCF) inside + // executeObject for the nested Parent type, triggering the partial-results branch in + // buildFieldValueMap. + def sdl = ''' + type Query { + parent : Parent + } + type Parent { + fast : String + slow : String + } + ''' + + CountDownLatch slowStarted = new CountDownLatch(1) + CompletableFuture slowResult = new CompletableFuture<>() + + DataFetcher parentDf = { DataFetchingEnvironment env -> [:] } + DataFetcher fastDf = { DataFetchingEnvironment env -> "fast-value" } + DataFetcher slowDf = { DataFetchingEnvironment env -> + slowStarted.countDown() + return slowResult + } + + def fetcherMap = [ + "Query" : ["parent": parentDf], + "Parent": ["fast": fastDf, "slow": slowDf] + ] + def schema = TestUtil.schema(sdl, fetcherMap) + def graphQL = GraphQL.newGraphQL(schema).build() + + when: + ExecutionInput executionInput = ExecutionInput.newExecutionInput() + .query("{ parent { fast slow } }") + .graphQLContext({ c -> + GraphQL.unusualConfiguration(c).cancellation().capturePartialResultsOnCancel(true) + }) + .build() + + def cf = graphQL.executeAsync(executionInput) + + // wait for the slow DF to have been invoked — fast is already done synchronously + slowStarted.await() + // cancel while slow's CF is still pending; this completes the cancellation future + // which races against the pending slowResult inside await(cancelCF) + executionInput.cancel() + + await().atMost(Duration.ofSeconds(10)).until({ -> cf.isDone() }) + def er = cf.join() + + then: + !cf.isCompletedExceptionally() + !er.errors.isEmpty() + er.errors.any { it["message"] == "Execution has been asked to be cancelled" } + // The top-level partial results capture sees parent's CF as still pending + // (because its nested sub-fields are being resolved), so parent is null. + // The nested buildFieldValueMap lambda also fires on the same cancellation. + er.data != null + er.data.containsKey("parent") + + cleanup: + // release slow so any background threads can finish + slowResult.complete("slow-value") + } + + def "without capturePartialResultsOnCancel a late cancel discards even fully gathered nested data"() { + // capture is OFF, so there is no cancellation race; the top-level 'outer' field value info + // completes quickly while its nested 'value' is still resolving. Cancelling at that point and + // then letting 'value' finish means the field values are fully gathered, but because capture + // is disabled the gathered data must be discarded and only the cancel error returned. + def sdl = ''' + type Query { + outer : Inner + } + type Inner { + value : String + } + ''' + + CountDownLatch valueStarted = new CountDownLatch(1) + CountDownLatch valueRelease = new CountDownLatch(1) + + DataFetcher outerDf = { DataFetchingEnvironment env -> [:] } + DataFetcher valueDf = { DataFetchingEnvironment env -> + return CompletableFuture.supplyAsync { + valueStarted.countDown() + valueRelease.await() + return "value" + } + } + + def fetcherMap = ["Query": ["outer": outerDf], "Inner": ["value": valueDf]] + def schema = TestUtil.schema(sdl, fetcherMap) + def graphQL = GraphQL.newGraphQL(schema).build() + + when: + ExecutionInput executionInput = ExecutionInput.newExecutionInput() + .query("{ outer { value } }") + .build() // capturePartialResultsOnCancel NOT enabled + + def cf = graphQL.executeAsync(executionInput) + + // the nested value DF has started, so 'outer' field value info is already done + valueStarted.await() + // cancel while the nested value is still resolving, then let it finish + executionInput.cancel() + valueRelease.countDown() + + await().atMost(Duration.ofSeconds(10)).until({ -> cf.isDone() }) + def er = cf.join() + + then: + !cf.isCompletedExceptionally() + !er.errors.isEmpty() + er.errors[0]["message"] == "Execution has been asked to be cancelled" + // capture disabled - even though the data was fully gathered, it is discarded + er.data == null + } + private static ExecutionResult awaitAsync(GraphQL graphQL, ExecutionInput executionInput) { def cf = graphQL.executeAsync(executionInput) await().atMost(Duration.ofSeconds(10)).until({ -> cf.isDone() })