Skip to content

Fixed deferred support to have proper Instrumentation #4083

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
bbakerman merged 3 commits into from
Aug 15, 2025
Merged

Fixed deferred support to have proper Instrumentation #4083

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
890ddb3
Fixed deferred support to have proper Instrumentation
bbakerman Aug 12, 2025
b831e21
Fixed deferred support to have proper Instrumentation - fixed bugs
bbakerman Aug 12, 2025
215be0b
Added NoContextChainedInstrumentation.java support
bbakerman Aug 12, 2025
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
8 changes: 7 additions & 1 deletion src/main/java/graphql/execution/ExecutionStrategy.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -292,7 +292,8 @@ DeferredExecutionSupport createDeferredExecutionSupport(ExecutionContext executi
fields,
parameters,
executionContext,
(ec, esp) -> Async.toCompletableFuture(resolveFieldWithInfo(ec, esp))
(ec, esp) -> Async.toCompletableFuture(resolveFieldWithInfo(ec, esp)),
this::createExecutionStepInfo
) : DeferredExecutionSupport.NOOP;

}
Expand Down Expand Up @@ -1096,6 +1097,11 @@ protected ExecutionStepInfo createExecutionStepInfo(ExecutionContext executionCo
fieldContainer);
}

private Supplier<ExecutionStepInfo> createExecutionStepInfo(ExecutionContext executionContext, ExecutionStrategyParameters parameters) {
GraphQLFieldDefinition fieldDef = getFieldDef(executionContext, parameters, parameters.getField().getSingleField());
return FpKit.intraThreadMemoize(() -> createExecutionStepInfo(executionContext, parameters, fieldDef, null));
Copy link
Member Author

@bbakerman bbakerman Aug 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I need to be able to create a ExecutionStepInfo for a field but later

}

// Errors that result from the execution of deferred fields are kept in the deferred context only.
private static void addErrorToRightContext(GraphQLError error, ExecutionStrategyParameters parameters, ExecutionContext executionContext) {
if (parameters.getDeferredCallContext() != null) {
Expand Down
64 changes: 41 additions & 23 deletions 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 @@ -7,14 +7,18 @@
import graphql.ExecutionResultImpl;
import graphql.Internal;
import graphql.execution.ExecutionContext;
import graphql.execution.ExecutionStepInfo;
import graphql.execution.ExecutionStrategyParameters;
import graphql.execution.FieldValueInfo;
import graphql.execution.MergedField;
import graphql.execution.MergedSelectionSet;
import graphql.execution.ResultPath;
import graphql.execution.instrumentation.Instrumentation;
import graphql.execution.instrumentation.InstrumentationContext;
import graphql.execution.instrumentation.parameters.InstrumentationFieldParameters;
import graphql.incremental.IncrementalPayload;
import graphql.util.FpKit;
import org.jspecify.annotations.NonNull;

import java.util.Collections;
import java.util.HashMap;
Expand All @@ -27,6 +31,8 @@
import java.util.function.BiFunction;
import java.util.function.Supplier;

import static graphql.execution.instrumentation.SimpleInstrumentationContext.nonNullCtx;

/**
* The purpose of this class hierarchy is to encapsulate most of the logic for deferring field execution, thus
* keeping the main execution strategy code clean and focused on the main execution logic.
Expand Down Expand Up @@ -59,16 +65,19 @@ class DeferredExecutionSupportImpl implements DeferredExecutionSupport {
private final ExecutionStrategyParameters parameters;
private final ExecutionContext executionContext;
private final BiFunction<ExecutionContext, ExecutionStrategyParameters, CompletableFuture<FieldValueInfo>> resolveFieldWithInfoFn;
private final BiFunction<ExecutionContext, ExecutionStrategyParameters, Supplier<ExecutionStepInfo>> executionStepInfoFn;
private final Map<String, Supplier<CompletableFuture<DeferredFragmentCall.FieldWithExecutionResult>>> dfCache = new HashMap<>();

public DeferredExecutionSupportImpl(
MergedSelectionSet mergedSelectionSet,
ExecutionStrategyParameters parameters,
ExecutionContext executionContext,
BiFunction<ExecutionContext, ExecutionStrategyParameters, CompletableFuture<FieldValueInfo>> resolveFieldWithInfoFn
BiFunction<ExecutionContext, ExecutionStrategyParameters, CompletableFuture<FieldValueInfo>> resolveFieldWithInfoFn,
BiFunction<ExecutionContext, ExecutionStrategyParameters, Supplier<ExecutionStepInfo>> executionStepInfoFn
Copy link
Member Author

@bbakerman bbakerman Aug 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I need to be able to build the ExecutionStepInfo later and this allows me to do that back in the execution strategy

) {
this.executionContext = executionContext;
this.resolveFieldWithInfoFn = resolveFieldWithInfoFn;
this.executionStepInfoFn = executionStepInfoFn;
ImmutableListMultimap.Builder<DeferredExecution, MergedField> deferredExecutionToFieldsBuilder = ImmutableListMultimap.builder();
ImmutableSet.Builder<MergedField> deferredFieldsBuilder = ImmutableSet.builder();
ImmutableList.Builder<String> nonDeferredFieldNamesBuilder = ImmutableList.builder();
Expand Down Expand Up @@ -153,37 +162,46 @@ private Supplier<CompletableFuture<DeferredFragmentCall.FieldWithExecutionResult
}
);


Instrumentation instrumentation = executionContext.getInstrumentation();

instrumentation.beginDeferredField(executionContext.getInstrumentationState());

// todo: handle cached computations
return dfCache.computeIfAbsent(
currentField.getResultKey(),
// The same field can be associated with multiple defer executions, so
// we memoize the field resolution to avoid multiple calls to the same data fetcher
key -> FpKit.interThreadMemoize(() -> {
CompletableFuture<FieldValueInfo> fieldValueResult = resolveFieldWithInfoFn.apply(executionContext, executionStrategyParameters);
key -> FpKit.interThreadMemoize(resolveDeferredFieldValue(currentField, executionContext, executionStrategyParameters)
)
Copy link
Member Author

@bbakerman bbakerman Aug 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

moved most of the code into a method resolveDeferredFieldValue - just because it was a lambda in a lambda

);
}

fieldValueResult.whenComplete((fieldValueInfo, throwable) -> {
executionContext.getDataLoaderDispatcherStrategy().deferredOnFieldValue(currentField.getResultKey(), fieldValueInfo, throwable, executionStrategyParameters);
});
@NonNull
private Supplier<CompletableFuture<DeferredFragmentCall.FieldWithExecutionResult>> resolveDeferredFieldValue(MergedField currentField, ExecutionContext executionContext, ExecutionStrategyParameters executionStrategyParameters) {
return () -> {

Instrumentation instrumentation = executionContext.getInstrumentation();
Supplier<ExecutionStepInfo> executionStepInfo = executionStepInfoFn.apply(executionContext, executionStrategyParameters);
InstrumentationFieldParameters fieldParameters = new InstrumentationFieldParameters(executionContext, executionStepInfo);
InstrumentationContext<Object> deferredFieldCtx = nonNullCtx(instrumentation.beginDeferredField(fieldParameters, executionContext.getInstrumentationState()));

CompletableFuture<ExecutionResult> executionResultCF = fieldValueResult
.thenCompose(fvi -> fvi
.getFieldValueFuture()
.thenApply(fv -> ExecutionResultImpl.newExecutionResult().data(fv).build())
);
CompletableFuture<FieldValueInfo> fieldValueResult = resolveFieldWithInfoFn.apply(this.executionContext, executionStrategyParameters);

return executionResultCF
.thenApply(executionResult ->
new DeferredFragmentCall.FieldWithExecutionResult(currentField.getResultKey(), executionResult)
);
}
)
);
deferredFieldCtx.onDispatched();

fieldValueResult.whenComplete((fieldValueInfo, throwable) -> {
this.executionContext.getDataLoaderDispatcherStrategy().deferredOnFieldValue(currentField.getResultKey(), fieldValueInfo, throwable, executionStrategyParameters);
deferredFieldCtx.onCompleted(fieldValueInfo, throwable);
});


CompletableFuture<ExecutionResult> executionResultCF = fieldValueResult
.thenCompose(fvi -> fvi
.getFieldValueFuture()
.thenApply(fv -> ExecutionResultImpl.newExecutionResult().data(fv).build())
);

return executionResultCF
.thenApply(executionResult ->
new DeferredFragmentCall.FieldWithExecutionResult(currentField.getResultKey(), executionResult)
);
};
}
}

Expand Down
5 changes: 2 additions & 3 deletions src/main/java/graphql/execution/instrumentation/ChainedInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -166,11 +166,10 @@ public ExecutionStrategyInstrumentationContext beginExecutionStrategy(Instrument

@ExperimentalApi
@Override
public InstrumentationContext<Object> beginDeferredField(InstrumentationState instrumentationState) {
return new ChainedDeferredExecutionStrategyInstrumentationContext(chainedMapAndDropNulls(instrumentationState, Instrumentation::beginDeferredField));
public InstrumentationContext<Object> beginDeferredField(InstrumentationFieldParameters parameters, InstrumentationState state) {
return chainedCtx(state, (instrumentation, specificState) -> instrumentation.beginDeferredField(parameters, specificState));
Copy link
Member Author

@bbakerman bbakerman Aug 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was just wrong and it had not tests

}


@Override
public InstrumentationContext<ExecutionResult> beginSubscribedFieldEvent(InstrumentationFieldParameters parameters, InstrumentationState state) {
return chainedCtx(state, (instrumentation, specificState) -> instrumentation.beginSubscribedFieldEvent(parameters, specificState));
Expand Down
7 changes: 4 additions & 3 deletions src/main/java/graphql/execution/instrumentation/Instrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -153,12 +153,13 @@ default ExecuteObjectInstrumentationContext beginExecuteObject(InstrumentationEx
* <p>
* This is an EXPERIMENTAL instrumentation callback. The method signature will definitely change.
*
* @param state the state created during the call to {@link #createStateAsync(InstrumentationCreateStateParameters)}
* @param parameters the parameters to this step
* @param state the state created during the call to {@link #createStateAsync(InstrumentationCreateStateParameters)}
*
* @return a nullable {@link ExecutionStrategyInstrumentationContext} object that will be called back when the step ends (assuming it's not null)
* @return a nullable {@link InstrumentationContext} object that will be called back when the step ends (assuming it's not null)
*/
@ExperimentalApi
default InstrumentationContext<Object> beginDeferredField(InstrumentationState state) {
default InstrumentationContext<Object> beginDeferredField(InstrumentationFieldParameters parameters, InstrumentationState state) {
return noOp();
Copy link
Member Author

@bbakerman bbakerman Aug 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we now have field parameters

}

Expand Down
5 changes: 5 additions & 0 deletions src/main/java/graphql/execution/instrumentation/NoContextChainedInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,11 @@ public ExecutionStrategyInstrumentationContext beginExecutionStrategy(Instrument
return runAll(state, (instrumentation, specificState) -> instrumentation.beginExecuteObject(parameters, specificState));
}

@Override
public InstrumentationContext<Object> beginDeferredField(InstrumentationFieldParameters parameters, InstrumentationState state) {
return runAll(state, (instrumentation, specificState) -> instrumentation.beginDeferredField(parameters, specificState));
}

@Override
public InstrumentationContext<ExecutionResult> beginSubscribedFieldEvent(InstrumentationFieldParameters parameters, InstrumentationState state) {
return runAll(state, (instrumentation, specificState) -> instrumentation.beginSubscribedFieldEvent(parameters, specificState));
Expand Down
21 changes: 21 additions & 0 deletions src/test/groovy/graphql/TestUtil.groovy
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,9 @@ package graphql

import graphql.execution.MergedField
import graphql.execution.MergedSelectionSet
import graphql.execution.pubsub.CapturingSubscriber
import graphql.incremental.DelayedIncrementalPartialResult
import graphql.incremental.IncrementalExecutionResult
import graphql.introspection.Introspection.DirectiveLocation
import graphql.language.Document
import graphql.language.Field
Expand Down Expand Up @@ -31,6 +34,8 @@ import graphql.schema.idl.TypeRuntimeWiring
import graphql.schema.idl.WiringFactory
import graphql.schema.idl.errors.SchemaProblem
import groovy.json.JsonOutput
import org.awaitility.Awaitility
import org.reactivestreams.Publisher

import java.util.function.Supplier
import java.util.stream.Collectors
Expand Down Expand Up @@ -323,4 +328,20 @@ class TestUtil {
return rn.nextInt(max - min + 1) + min
}


static List<Map<String, Object>> getIncrementalResults(IncrementalExecutionResult initialResult) {
Publisher<DelayedIncrementalPartialResult> deferredResultStream = initialResult.incrementalItemPublisher

def subscriber = new CapturingSubscriber<DelayedIncrementalPartialResult>()

deferredResultStream.subscribe(subscriber)

Awaitility.await().untilTrue(subscriber.isDone())
if (subscriber.throwable != null) {
throw new RuntimeException(subscriber.throwable)
}
return subscriber.getEvents()
.collect { it.toSpecification() }
}

}
78 changes: 78 additions & 0 deletions src/test/groovy/graphql/execution/instrumentation/InstrumentationTest.groovy
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,13 @@ import graphql.ExecutionInput
import graphql.ExecutionResult
import graphql.GraphQL
import graphql.StarWarsSchema
import graphql.TestUtil
import graphql.execution.AsyncExecutionStrategy
import graphql.execution.instrumentation.parameters.InstrumentationCreateStateParameters
import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters
import graphql.execution.instrumentation.parameters.InstrumentationExecutionStrategyParameters
import graphql.execution.instrumentation.parameters.InstrumentationFieldFetchParameters
import graphql.incremental.IncrementalExecutionResult
import graphql.language.AstPrinter
import graphql.parser.Parser
import graphql.schema.DataFetcher
Expand Down Expand Up @@ -496,4 +498,80 @@ class InstrumentationTest extends Specification {
then:
er.extensions == [i1: "I1"]
}

def "can instrumented deferred fields"() {

given:

def query = """
{
hero {
id
... @defer(label: "name") {
name
}
}
}
"""


when:

def instrumentation = new ModernTestingInstrumentation()

def graphQL = GraphQL
.newGraphQL(StarWarsSchema.starWarsSchema)
.queryExecutionStrategy(new AsyncExecutionStrategy())
.instrumentation(instrumentation)
.build()

def ei = ExecutionInput.newExecutionInput(query).graphQLContext { it ->
GraphQL.unusualConfiguration(it).incrementalSupport().enableIncrementalSupport(true)
}.build()

IncrementalExecutionResult incrementalER = graphQL.execute(ei) as IncrementalExecutionResult
//
// cause the defer Publish to be finished
def results = TestUtil.getIncrementalResults(incrementalER)


then:

instrumentation.executionList == ["start:execution",
"start:parse",
"end:parse",
"start:validation",
"end:validation",
"start:execute-operation",
"start:execution-strategy",
"start:field-hero",
"start:fetch-hero",
"end:fetch-hero",
"start:complete-hero",
"start:execute-object",
"start:field-id",
"start:fetch-id",
"end:fetch-id",
"start:complete-id",
"end:complete-id",
"end:field-id",

"end:execute-object",
"end:complete-hero",
"end:field-hero",
"end:execution-strategy",
"end:execute-operation",
"end:execution",
//
// the deferred field resolving now happens after the operation has come back
"start:deferred-field-name",
"start:field-name",
"start:fetch-name",
"end:fetch-name",
"start:complete-name",
"end:complete-name",
"end:field-name",
"end:deferred-field-name",
]
}
}
6 changes: 6 additions & 0 deletions src/test/groovy/graphql/execution/instrumentation/ModernTestingInstrumentation.groovy
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,12 @@ class ModernTestingInstrumentation implements Instrumentation {
return new TestingInstrumentContext("complete-list-$parameters.field.name", executionList, throwableList, useOnDispatch)
}

@Override
InstrumentationContext<Object> beginDeferredField(InstrumentationFieldParameters parameters, InstrumentationState state) {
assert state == instrumentationState
return new TestingInstrumentContext("deferred-field-$parameters.field.name", executionList, throwableList, useOnDispatch)
}

@Override
ExecutionInput instrumentExecutionInput(ExecutionInput executionInput, InstrumentationExecutionParameters parameters, InstrumentationState state) {
assert state == instrumentationState
Expand Down