Skip to content

Fix subscription alternative call context cleanup #4384

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
andimarek merged 1 commit into from
May 18, 2026
Merged

Fix subscription alternative call context cleanup #4384

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions src/main/java/graphql/execution/DataLoaderDispatchStrategy.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,10 @@ default void subscriptionEventCompletionDone(AlternativeCallContext alternativeC

}

default void subscriptionEventExecutionDone(AlternativeCallContext alternativeCallContext) {

}

default void finishedFetching(ExecutionContext executionContext, ExecutionStrategyParameters newParameters) {

}
Expand Down
15 changes: 9 additions & 6 deletions src/main/java/graphql/execution/ExecutionStrategy.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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();
});
Expand Down Expand Up @@ -1122,18 +1123,20 @@ private Supplier<ExecutionStepInfo> 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<GraphQLError> 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);
}
Expand Down
25 changes: 5 additions & 20 deletions src/main/java/graphql/execution/ExecutionStrategyParameters.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,28 +77,13 @@ public ResultPath getPath() {
}

/**
* 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.
*
* <pre>
* query {
* ... @defer {
* field1 { # new DeferredCallContext created here
* field1a # DeferredCallContext passed down to this field
* }
* }
*
* ... @defer {
* field2 # new DeferredCallContext created here
* }
* }
* </pre>
*
* @return the deferred call context or null if we're not in the scope of a deferred call
* 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 getDeferredCallContext() {
public AlternativeCallContext getAlternativeCallContext() {
return alternativeCallContext;
}

Expand Down Expand Up @@ -293,7 +278,7 @@ public Builder parent(ExecutionStrategyParameters parent) {
return this;
}

public Builder deferredCallContext(AlternativeCallContext alternativeCallContext) {
public Builder alternativeCallContext(AlternativeCallContext alternativeCallContext) {
this.alternativeCallContext = alternativeCallContext;
return this;
}
Expand Down
4 changes: 2 additions & 2 deletions src/main/java/graphql/execution/NonNullableFieldValidator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,8 +50,8 @@ public <T> 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);
}
Expand Down
16 changes: 11 additions & 5 deletions src/main/java/graphql/execution/SubscriptionExecutionStrategy.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,14 +180,20 @@ private CompletableFuture<ExecutionResult> 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<ExecutionResult> 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();
Expand Down Expand Up @@ -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));
}
});

Expand Down
2 changes: 1 addition & 1 deletion src/main/java/graphql/execution/incremental/DeferredExecutionSupport.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -155,7 +155,7 @@ private Supplier<CompletableFuture<DeferredFragmentCall.FieldWithExecutionResult
{
MergedSelectionSet mergedSelectionSet = MergedSelectionSet.newMergedSelectionSet().subFields(fields).build();
ResultPath path = parameters.getPath().segment(currentField.getResultKey());
builder.deferredCallContext(alternativeCallContext)
builder.alternativeCallContext(alternativeCallContext)
.field(currentField)
.fields(mergedSelectionSet)
.path(path)
Expand Down
12 changes: 8 additions & 4 deletions ...ava/graphql/execution/instrumentation/dataloader/ExhaustedDataLoaderDispatchStrategy.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -173,12 +173,17 @@ public void subscriptionEventCompletionDone(AlternativeCallContext alternativeCa
decrementObjectRunningAndMaybeDispatch(callStack);
}

@Override
public void subscriptionEventExecutionDone(AlternativeCallContext alternativeCallContext) {
alternativeCallContextMap.remove(alternativeCallContext);
}

@Override
public void deferFieldFetched(ExecutionStrategyParameters parameters) {
CallStack callStack = getCallStack(parameters);
int deferredFragmentRootFieldsCompleted = callStack.deferredFragmentRootFieldsCompleted.incrementAndGet();
Assert.assertNotNull(parameters.getDeferredCallContext());
if (deferredFragmentRootFieldsCompleted == parameters.getDeferredCallContext().getFields()) {
Assert.assertNotNull(parameters.getAlternativeCallContext());
if (deferredFragmentRootFieldsCompleted == parameters.getAlternativeCallContext().getFields()) {
decrementObjectRunningAndMaybeDispatch(callStack);
}
}
Expand All @@ -195,7 +200,7 @@ public void stopComplete(ExecutionStrategyParameters parameters) {
}

private CallStack getCallStack(ExecutionStrategyParameters parameters) {
return getCallStack(parameters.getDeferredCallContext());
return getCallStack(parameters.getAlternativeCallContext());
}

private CallStack getCallStack(@Nullable AlternativeCallContext alternativeCallContext) {
Expand Down Expand Up @@ -281,4 +286,3 @@ public void newDataLoaderInvocation(@Nullable AlternativeCallContext alternative


}

18 changes: 11 additions & 7 deletions ...java/graphql/execution/instrumentation/dataloader/PerLevelDataLoaderDispatchStrategy.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -362,8 +362,8 @@ public void fieldFetched(ExecutionContext executionContext,
Supplier<DataFetchingEnvironment> 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);
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -520,4 +525,3 @@ public void newDataLoaderInvocation(int level,


}

4 changes: 2 additions & 2 deletions src/main/java/graphql/schema/DataFetchingEnvironmentImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -458,7 +458,7 @@ public Builder queryDirectives(QueryDirectives queryDirectives) {
return this;
}

public Builder deferredCallContext(AlternativeCallContext alternativeCallContext) {
public Builder alternativeCallContext(AlternativeCallContext alternativeCallContext) {
this.alternativeCallContext = alternativeCallContext;
return this;
}
Expand Down Expand Up @@ -499,7 +499,7 @@ public DataLoaderDispatchStrategy getDataLoaderDispatchStrategy() {
return dataLoaderDispatchStrategy;
}

public AlternativeCallContext getDeferredCallContext() {
public AlternativeCallContext getAlternativeCallContext() {
return alternativeCallContext;
}

Expand Down
4 changes: 2 additions & 2 deletions src/main/java/graphql/schema/DataLoaderWithContext.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}
Expand Down
22 changes: 21 additions & 1 deletion src/test/groovy/graphql/execution/SubscriptionExecutionStrategyTest.groovy
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand Down Expand Up @@ -830,6 +832,16 @@ class SubscriptionExecutionStrategyTest extends Specification {
def dataLoader = DataLoaderFactory.newDataLoader("dogsNameLoader", batchLoader)
DataLoaderRegistry dataLoaderRegistry = new DataLoaderRegistry()
dataLoaderRegistry.register("dogsNameLoader", dataLoader)
AtomicReference<ExecutionContext> capturedExecutionContext = new AtomicReference<>()
def instrumentation = Spy(SimplePerformantInstrumentation) {
instrumentExecutionContext(_, _, _) >> {
ExecutionContext executionContext,
InstrumentationExecutionParameters parameters,
InstrumentationState state ->
capturedExecutionContext.set(executionContext)
executionContext
}
}

DataFetcher dogsNameDF = { env ->
println "dogsNameDF called"
Expand Down Expand Up @@ -857,6 +869,7 @@ class SubscriptionExecutionStrategyTest extends Specification {
.dataLoaderRegistry(dataLoaderRegistry)
.build()
def graphQL = GraphQL.newGraphQL(schema)
.instrumentation(instrumentation)
.build()

if (exhaustedStrategy) {
Expand All @@ -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"() {

Expand Down
26 changes: 25 additions & 1 deletion ...aphql/execution/instrumentation/dataloader/ExhaustedDataLoaderDispatchStrategyTest.groovy
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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())
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
}
}
Loading