From 46a265c126f951b28524d330b9021558c3ca4961 Mon Sep 17 00:00:00 2001 From: Andreas Marek Date: Tue, 19 May 2026 08:17:02 +1000 Subject: [PATCH] Fix subscription alternative call context cleanup --- .../execution/DataLoaderDispatchStrategy.java | 4 +++ .../graphql/execution/ExecutionStrategy.java | 15 ++++---- .../ExecutionStrategyParameters.java | 34 +++++++++---------- .../execution/NonNullableFieldValidator.java | 4 +-- .../SubscriptionExecutionStrategy.java | 16 ++++++--- .../incremental/DeferredExecutionSupport.java | 2 +- .../ExhaustedDataLoaderDispatchStrategy.java | 12 ++++--- .../PerLevelDataLoaderDispatchStrategy.java | 18 ++++++---- .../schema/DataFetchingEnvironmentImpl.java | 12 +++++-- .../graphql/schema/DataLoaderWithContext.java | 4 +-- .../SubscriptionExecutionStrategyTest.groovy | 22 +++++++++++- ...ustedDataLoaderDispatchStrategyTest.groovy | 26 +++++++++++++- ...LevelDataLoaderDispatchStrategyTest.groovy | 22 ++++++++++++ 13 files changed, 143 insertions(+), 48 deletions(-) diff --git a/src/main/java/graphql/execution/DataLoaderDispatchStrategy.java b/src/main/java/graphql/execution/DataLoaderDispatchStrategy.java index ae73dc2fe2..1231a803a9 100644 --- a/src/main/java/graphql/execution/DataLoaderDispatchStrategy.java +++ b/src/main/java/graphql/execution/DataLoaderDispatchStrategy.java @@ -65,6 +65,10 @@ default void subscriptionEventCompletionDone(AlternativeCallContext alternativeC } + default void subscriptionEventExecutionDone(AlternativeCallContext alternativeCallContext) { + + } + default void finishedFetching(ExecutionContext executionContext, ExecutionStrategyParameters newParameters) { } diff --git a/src/main/java/graphql/execution/ExecutionStrategy.java b/src/main/java/graphql/execution/ExecutionStrategy.java index 9f402a24d6..713da688d0 100644 --- a/src/main/java/graphql/execution/ExecutionStrategy.java +++ b/src/main/java/graphql/execution/ExecutionStrategy.java @@ -14,6 +14,7 @@ import graphql.UnresolvedTypeError; import graphql.execution.directives.QueryDirectives; import graphql.execution.directives.QueryDirectivesImpl; +import graphql.execution.incremental.AlternativeCallContext; import graphql.execution.incremental.DeferredExecutionSupport; import graphql.execution.incremental.IncrementalExecutionContextKeys; import graphql.execution.instrumentation.ExecuteObjectInstrumentationContext; @@ -457,7 +458,7 @@ private Object fetchField(GraphQLFieldDefinition fieldDef, ExecutionContext exec .parentType(parentType) .selectionSet(fieldCollector) .queryDirectives(queryDirectives) - .deferredCallContext(parameters.getDeferredCallContext()) + .alternativeCallContext(parameters.getAlternativeCallContext()) .level(parameters.getPath().getLevel()) .build(); }); @@ -1122,18 +1123,20 @@ private Supplier createExecutionStepInfo(ExecutionContext exe return FpKit.intraThreadMemoize(() -> createExecutionStepInfo(executionContext, parameters, fieldDef, null)); } - // Errors that result from the execution of deferred fields are kept in the deferred context only. + // Errors in alternative execution paths are kept in the alternative call context. private static void addErrorToRightContext(GraphQLError error, ExecutionStrategyParameters parameters, ExecutionContext executionContext) { - if (parameters.getDeferredCallContext() != null) { - parameters.getDeferredCallContext().addError(error); + AlternativeCallContext alternativeCallContext = parameters.getAlternativeCallContext(); + if (alternativeCallContext != null) { + alternativeCallContext.addError(error); } else { executionContext.addError(error); } } private static void addErrorsToRightContext(List errors, ExecutionStrategyParameters parameters, ExecutionContext executionContext) { - if (parameters.getDeferredCallContext() != null) { - parameters.getDeferredCallContext().addErrors(errors); + AlternativeCallContext alternativeCallContext = parameters.getAlternativeCallContext(); + if (alternativeCallContext != null) { + alternativeCallContext.addErrors(errors); } else { executionContext.addErrors(errors); } diff --git a/src/main/java/graphql/execution/ExecutionStrategyParameters.java b/src/main/java/graphql/execution/ExecutionStrategyParameters.java index 21b828b7d5..502ae220af 100644 --- a/src/main/java/graphql/execution/ExecutionStrategyParameters.java +++ b/src/main/java/graphql/execution/ExecutionStrategyParameters.java @@ -76,30 +76,26 @@ public ResultPath getPath() { return parent; } + /** + * Returns the alternative call context if this execution is scoped to an alternative execution path. + * This is used for deferred fragment execution and subscription event execution. + * @return the alternative call context or null if execution is not scoped to an alternative execution path + */ + @Nullable + @Internal + public AlternativeCallContext getAlternativeCallContext() { + return alternativeCallContext; + } + /** * Returns the deferred call context if we're in the scope of a deferred call. - * A new DeferredCallContext is created for each @defer block, and is passed down to all fields within the deferred call. - * - * 
-     *     query {
-     *        ... @defer {
-     *            field1 {        # new DeferredCallContext created here
-     *                field1a     # DeferredCallContext passed down to this field
-     *            }
-     *        }
-     *
-     *        ... @defer {
-     *            field2          # new DeferredCallContext created here
-     *        }
-     *     }
-     *
* * @return the deferred call context or null if we're not in the scope of a deferred call */ @Nullable @Internal public AlternativeCallContext getDeferredCallContext() { - return alternativeCallContext; + return getAlternativeCallContext(); } /** @@ -293,11 +289,15 @@ public Builder parent(ExecutionStrategyParameters parent) { return this; } - public Builder deferredCallContext(AlternativeCallContext alternativeCallContext) { + public Builder alternativeCallContext(AlternativeCallContext alternativeCallContext) { this.alternativeCallContext = alternativeCallContext; return this; } + public Builder deferredCallContext(AlternativeCallContext alternativeCallContext) { + return alternativeCallContext(alternativeCallContext); + } + public ExecutionStrategyParameters build() { return new ExecutionStrategyParameters(executionStepInfo, source, localContext, fields, nonNullableFieldValidator, path, currentField, parent, alternativeCallContext); } diff --git a/src/main/java/graphql/execution/NonNullableFieldValidator.java b/src/main/java/graphql/execution/NonNullableFieldValidator.java index b59f633bac..4680d9a7b1 100644 --- a/src/main/java/graphql/execution/NonNullableFieldValidator.java +++ b/src/main/java/graphql/execution/NonNullableFieldValidator.java @@ -50,8 +50,8 @@ public T validate(ExecutionStrategyParameters parameters, T result) throws N NonNullableFieldWasNullException nonNullException = new NonNullableFieldWasNullException(executionStepInfo, path); final GraphQLError error = new NonNullableFieldWasNullError(nonNullException); - if(parameters.getDeferredCallContext() != null) { - parameters.getDeferredCallContext().addError(error); + if(parameters.getAlternativeCallContext() != null) { + parameters.getAlternativeCallContext().addError(error); } else { executionContext.addError(error, path); } diff --git a/src/main/java/graphql/execution/SubscriptionExecutionStrategy.java b/src/main/java/graphql/execution/SubscriptionExecutionStrategy.java index 89c77e967a..bb57bdf736 100644 --- a/src/main/java/graphql/execution/SubscriptionExecutionStrategy.java +++ b/src/main/java/graphql/execution/SubscriptionExecutionStrategy.java @@ -180,14 +180,20 @@ private CompletableFuture executeSubscriptionEvent(ExecutionCon )); - executionContext.getDataLoaderDispatcherStrategy().newSubscriptionExecution(newParameters.getDeferredCallContext()); + AlternativeCallContext alternativeCallContext = assertNotNull( + newParameters.getAlternativeCallContext(), + "alternativeCallContext must not be null"); + executionContext.getDataLoaderDispatcherStrategy().newSubscriptionExecution(alternativeCallContext); Object fetchedValue = unboxPossibleDataFetcherResult(newExecutionContext, newParameters, eventPayload); FieldValueInfo fieldValueInfo = completeField(newExecutionContext, newParameters, fetchedValue); - executionContext.getDataLoaderDispatcherStrategy().subscriptionEventCompletionDone(newParameters.getDeferredCallContext()); + executionContext.getDataLoaderDispatcherStrategy().subscriptionEventCompletionDone(alternativeCallContext); CompletableFuture overallResult = fieldValueInfo .getFieldValueFuture() - .thenApply(val -> new ExecutionResultImpl(val, assertNotNull(newParameters.getDeferredCallContext(), "deferredCallContext must not be null").getErrors())) - .thenApply(executionResult -> wrapWithRootFieldName(newParameters, executionResult)); + .thenApply(val -> new ExecutionResultImpl(val, alternativeCallContext.getErrors())) + .thenApply(executionResult -> wrapWithRootFieldName(newParameters, executionResult)) + .whenComplete((executionResult, throwable) -> { + executionContext.getDataLoaderDispatcherStrategy().subscriptionEventExecutionDone(alternativeCallContext); + }); // dispatch instrumentation so they can know about each subscription event subscribedFieldCtx.onDispatched(); @@ -230,7 +236,7 @@ private ExecutionStrategyParameters firstFieldOfSubscriptionSelection(ExecutionC .path(fieldPath) .nonNullFieldValidator(nonNullableFieldValidator); if (newCallContext) { - builder.deferredCallContext(new AlternativeCallContext(1, 1)); + builder.alternativeCallContext(new AlternativeCallContext(1, 1)); } }); diff --git a/src/main/java/graphql/execution/incremental/DeferredExecutionSupport.java b/src/main/java/graphql/execution/incremental/DeferredExecutionSupport.java index f7ddbfe6f7..1ca1c81314 100644 --- a/src/main/java/graphql/execution/incremental/DeferredExecutionSupport.java +++ b/src/main/java/graphql/execution/incremental/DeferredExecutionSupport.java @@ -155,7 +155,7 @@ private Supplier dataFetchingEnvironment) { CallStack callStack = getCallStack(executionStrategyParameters); int level = executionStrategyParameters.getPath().getLevel(); - AlternativeCallContext deferredCallContext = executionStrategyParameters.getDeferredCallContext(); - if (level == 1 || (deferredCallContext != null && level == deferredCallContext.getStartLevel())) { + AlternativeCallContext alternativeCallContext = executionStrategyParameters.getAlternativeCallContext(); + if (level == 1 || (alternativeCallContext != null && level == alternativeCallContext.getStartLevel())) { int happenedFirstLevelFetchCount = callStack.happenedFirstLevelFetchCount.incrementAndGet(); if (happenedFirstLevelFetchCount == callStack.expectedFirstLevelFetchCount) { callStack.dispatchedLevels.add(level); @@ -395,20 +395,25 @@ public void subscriptionEventCompletionDone(AlternativeCallContext alternativeCa onCompletionFinished(0, callStack); } + @Override + public void subscriptionEventExecutionDone(AlternativeCallContext alternativeCallContext) { + alternativeCallContextMap.remove(alternativeCallContext); + } + @Override public void deferredOnFieldValue(String resultKey, FieldValueInfo fieldValueInfo, Throwable throwable, ExecutionStrategyParameters parameters) { CallStack callStack = getCallStack(parameters); int deferredFragmentRootFieldsCompleted = callStack.deferredFragmentRootFieldsCompleted.incrementAndGet(); - Assert.assertNotNull(parameters.getDeferredCallContext()); - if (deferredFragmentRootFieldsCompleted == parameters.getDeferredCallContext().getFields()) { - onCompletionFinished(parameters.getDeferredCallContext().getStartLevel() - 1, callStack); + Assert.assertNotNull(parameters.getAlternativeCallContext()); + if (deferredFragmentRootFieldsCompleted == parameters.getAlternativeCallContext().getFields()) { + onCompletionFinished(parameters.getAlternativeCallContext().getStartLevel() - 1, callStack); } } private CallStack getCallStack(ExecutionStrategyParameters parameters) { - return getCallStack(parameters.getDeferredCallContext()); + return getCallStack(parameters.getAlternativeCallContext()); } private CallStack getCallStack(@Nullable AlternativeCallContext alternativeCallContext) { @@ -520,4 +525,3 @@ public void newDataLoaderInvocation(int level, } - diff --git a/src/main/java/graphql/schema/DataFetchingEnvironmentImpl.java b/src/main/java/graphql/schema/DataFetchingEnvironmentImpl.java index b9cfce9485..eb13820559 100644 --- a/src/main/java/graphql/schema/DataFetchingEnvironmentImpl.java +++ b/src/main/java/graphql/schema/DataFetchingEnvironmentImpl.java @@ -458,11 +458,15 @@ public Builder queryDirectives(QueryDirectives queryDirectives) { return this; } - public Builder deferredCallContext(AlternativeCallContext alternativeCallContext) { + public Builder alternativeCallContext(AlternativeCallContext alternativeCallContext) { this.alternativeCallContext = alternativeCallContext; return this; } + public Builder deferredCallContext(AlternativeCallContext alternativeCallContext) { + return alternativeCallContext(alternativeCallContext); + } + public DataFetchingEnvironment build() { return new DataFetchingEnvironmentImpl(this); } @@ -499,10 +503,14 @@ public DataLoaderDispatchStrategy getDataLoaderDispatchStrategy() { return dataLoaderDispatchStrategy; } - public AlternativeCallContext getDeferredCallContext() { + public AlternativeCallContext getAlternativeCallContext() { return alternativeCallContext; } + public AlternativeCallContext getDeferredCallContext() { + return getAlternativeCallContext(); + } + public Profiler getProfiler() { return profiler; } diff --git a/src/main/java/graphql/schema/DataLoaderWithContext.java b/src/main/java/graphql/schema/DataLoaderWithContext.java index 3d4224b364..af4f74c693 100644 --- a/src/main/java/graphql/schema/DataLoaderWithContext.java +++ b/src/main/java/graphql/schema/DataLoaderWithContext.java @@ -68,11 +68,11 @@ private void newDataLoaderInvocation() { DataFetchingEnvironmentImpl dfeImpl = (DataFetchingEnvironmentImpl) dfe; DataFetchingEnvironmentImpl.DFEInternalState dfeInternalState = (DataFetchingEnvironmentImpl.DFEInternalState) dfeImpl.toInternal(); if (dfeInternalState.getDataLoaderDispatchStrategy() instanceof PerLevelDataLoaderDispatchStrategy) { - AlternativeCallContext alternativeCallContext = dfeInternalState.getDeferredCallContext(); + AlternativeCallContext alternativeCallContext = dfeInternalState.getAlternativeCallContext(); int level = dfeImpl.getLevel(); ((PerLevelDataLoaderDispatchStrategy) dfeInternalState.dataLoaderDispatchStrategy).newDataLoaderInvocation(level, delegate, alternativeCallContext); } else if (dfeInternalState.getDataLoaderDispatchStrategy() instanceof ExhaustedDataLoaderDispatchStrategy) { - AlternativeCallContext alternativeCallContext = dfeInternalState.getDeferredCallContext(); + AlternativeCallContext alternativeCallContext = dfeInternalState.getAlternativeCallContext(); ((ExhaustedDataLoaderDispatchStrategy) dfeInternalState.dataLoaderDispatchStrategy).newDataLoaderInvocation(alternativeCallContext); } } diff --git a/src/test/groovy/graphql/execution/SubscriptionExecutionStrategyTest.groovy b/src/test/groovy/graphql/execution/SubscriptionExecutionStrategyTest.groovy index af3dcaf5ce..0a206779b7 100644 --- a/src/test/groovy/graphql/execution/SubscriptionExecutionStrategyTest.groovy +++ b/src/test/groovy/graphql/execution/SubscriptionExecutionStrategyTest.groovy @@ -11,8 +11,9 @@ import graphql.TestUtil import graphql.TypeMismatchError import graphql.execution.instrumentation.InstrumentationState import graphql.execution.instrumentation.LegacyTestingInstrumentation -import graphql.execution.instrumentation.dataloader.DataLoaderDispatchingContextKeys import graphql.execution.instrumentation.ModernTestingInstrumentation +import graphql.execution.instrumentation.SimplePerformantInstrumentation +import graphql.execution.instrumentation.dataloader.DataLoaderDispatchingContextKeys import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters import graphql.execution.pubsub.CapturingSubscriber import graphql.execution.pubsub.FlowMessagePublisher @@ -36,6 +37,7 @@ import spock.lang.Unroll import java.util.concurrent.CompletableFuture import java.util.concurrent.CopyOnWriteArrayList import java.util.concurrent.atomic.AtomicInteger +import java.util.concurrent.atomic.AtomicReference import static graphql.schema.idl.TypeRuntimeWiring.newTypeWiring @@ -830,6 +832,16 @@ class SubscriptionExecutionStrategyTest extends Specification { def dataLoader = DataLoaderFactory.newDataLoader("dogsNameLoader", batchLoader) DataLoaderRegistry dataLoaderRegistry = new DataLoaderRegistry() dataLoaderRegistry.register("dogsNameLoader", dataLoader) + AtomicReference capturedExecutionContext = new AtomicReference<>() + def instrumentation = Spy(SimplePerformantInstrumentation) { + instrumentExecutionContext(_, _, _) >> { + ExecutionContext executionContext, + InstrumentationExecutionParameters parameters, + InstrumentationState state -> + capturedExecutionContext.set(executionContext) + executionContext + } + } DataFetcher dogsNameDF = { env -> println "dogsNameDF called" @@ -857,6 +869,7 @@ class SubscriptionExecutionStrategyTest extends Specification { .dataLoaderRegistry(dataLoaderRegistry) .build() def graphQL = GraphQL.newGraphQL(schema) + .instrumentation(instrumentation) .build() if (exhaustedStrategy) { @@ -878,11 +891,18 @@ class SubscriptionExecutionStrategyTest extends Specification { events[0].data == ["newDogs": [[name: "Luna"], [name: "Skipper"]]] events[1].data == ["newDogs": [[name: "Luna"], [name: "Skipper"]]] events[2].data == ["newDogs": [[name: "Luna"], [name: "Skipper"]]] + alternativeCallContextMap(capturedExecutionContext.get().dataLoaderDispatcherStrategy).size() == 0 where: exhaustedStrategy << [false, true] } + private Map alternativeCallContextMap(DataLoaderDispatchStrategy dataLoaderDispatchStrategy) { + def field = dataLoaderDispatchStrategy.class.getDeclaredField("alternativeCallContextMap") + field.accessible = true + field.get(dataLoaderDispatchStrategy) as Map + } + def "can instrument subscription reactive ending"() { diff --git a/src/test/groovy/graphql/execution/instrumentation/dataloader/ExhaustedDataLoaderDispatchStrategyTest.groovy b/src/test/groovy/graphql/execution/instrumentation/dataloader/ExhaustedDataLoaderDispatchStrategyTest.groovy index dd61b39071..6fed8f4d2b 100644 --- a/src/test/groovy/graphql/execution/instrumentation/dataloader/ExhaustedDataLoaderDispatchStrategyTest.groovy +++ b/src/test/groovy/graphql/execution/instrumentation/dataloader/ExhaustedDataLoaderDispatchStrategyTest.groovy @@ -18,6 +18,7 @@ import graphql.schema.GraphQLSchema import org.dataloader.BatchLoader import org.dataloader.DataLoaderFactory import org.dataloader.DataLoaderRegistry +import spock.lang.Issue import spock.lang.Specification import java.util.concurrent.CompletableFuture @@ -242,6 +243,23 @@ class ExhaustedDataLoaderDispatchStrategyTest extends Specification { batchLoaderInvocations.get() == 1 } + @Issue("https://github.com/graphql-java/graphql-java/issues/4314") + def "subscription event call stacks are removed after execution is done"() { + given: + setupStrategy(simpleBatchLoader()) + + when: + 3.times { + def alternativeCallContext = new AlternativeCallContext(1, 1) + strategy.newSubscriptionExecution(alternativeCallContext) + strategy.subscriptionEventCompletionDone(alternativeCallContext) + strategy.subscriptionEventExecutionDone(alternativeCallContext) + } + + then: + alternativeCallContextMap().size() == 0 + } + def "startComplete and stopComplete affect dispatch"() { given: setupStrategy(simpleBatchLoader()) @@ -279,7 +297,7 @@ class ExhaustedDataLoaderDispatchStrategyTest extends Specification { .source(new Object()) .fields(graphql.execution.MergedSelectionSet.newMergedSelectionSet().build()) .nonNullFieldValidator(new NonNullableFieldValidator(executionContext)) - .deferredCallContext(deferCtx) + .alternativeCallContext(deferCtx) .build() when: @@ -430,4 +448,10 @@ class ExhaustedDataLoaderDispatchStrategyTest extends Specification { completed roundCount.get() == 2 } + + private Map alternativeCallContextMap() { + def field = ExhaustedDataLoaderDispatchStrategy.getDeclaredField("alternativeCallContextMap") + field.accessible = true + field.get(strategy) as Map + } } diff --git a/src/test/groovy/graphql/execution/instrumentation/dataloader/PerLevelDataLoaderDispatchStrategyTest.groovy b/src/test/groovy/graphql/execution/instrumentation/dataloader/PerLevelDataLoaderDispatchStrategyTest.groovy index b67dc7e37b..5491945237 100644 --- a/src/test/groovy/graphql/execution/instrumentation/dataloader/PerLevelDataLoaderDispatchStrategyTest.groovy +++ b/src/test/groovy/graphql/execution/instrumentation/dataloader/PerLevelDataLoaderDispatchStrategyTest.groovy @@ -16,9 +16,11 @@ import graphql.execution.NonNullableFieldValidator import graphql.execution.ResultPath import graphql.execution.ValueUnboxer import graphql.execution.instrumentation.SimplePerformantInstrumentation +import graphql.execution.incremental.AlternativeCallContext import graphql.schema.DataFetcher import graphql.schema.DataFetchingEnvironment import org.dataloader.DataLoaderRegistry +import spock.lang.Issue import spock.lang.Specification import java.util.concurrent.CountDownLatch @@ -60,6 +62,20 @@ class PerLevelDataLoaderDispatchStrategyTest extends Specification { strategy = new PerLevelDataLoaderDispatchStrategy(executionContext) } + @Issue("https://github.com/graphql-java/graphql-java/issues/4314") + def "subscription event call stacks are removed after execution is done"() { + when: + 3.times { + def alternativeCallContext = new AlternativeCallContext(1, 1) + strategy.newSubscriptionExecution(alternativeCallContext) + strategy.subscriptionEventCompletionDone(alternativeCallContext) + strategy.subscriptionEventExecutionDone(alternativeCallContext) + } + + then: + alternativeCallContextMap().size() == 0 + } + private ExecutionStrategyParameters paramsAtLevel(int level) { def path = ResultPath.rootPath() for (int i = 0; i < level; i++) { @@ -179,4 +195,10 @@ class PerLevelDataLoaderDispatchStrategyTest extends Specification { then: strategy.initialCallStack.get(0).happenedCompletionFinishedCount > 0 } + + private Map alternativeCallContextMap() { + def field = PerLevelDataLoaderDispatchStrategy.getDeclaredField("alternativeCallContextMap") + field.accessible = true + field.get(strategy) as Map + } }