Skip to content

Commit 9c92725

Browse files
bbakermanbbakerman
authored
Merge pull request #4083 from graphql-java/fix-deferred-field-instrumentation
Fixed deferred support to have proper Instrumentation
2 parents 3ba9750 + 215be0b commit 9c92725

File tree

8 files changed

+164
-30
lines changed

8 files changed

+164
-30
lines changed

‎src/main/java/graphql/execution/ExecutionStrategy.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -292,7 +292,8 @@ DeferredExecutionSupport createDeferredExecutionSupport(ExecutionContext executi
292292
fields,
293293
parameters,
294294
executionContext,
295-
(ec, esp) -> Async.toCompletableFuture(resolveFieldWithInfo(ec, esp))
295+
(ec, esp) -> Async.toCompletableFuture(resolveFieldWithInfo(ec, esp)),
296+
this::createExecutionStepInfo
296297
) : DeferredExecutionSupport.NOOP;
297298

298299
}
@@ -1096,6 +1097,11 @@ protected ExecutionStepInfo createExecutionStepInfo(ExecutionContext executionCo
10961097
fieldContainer);
10971098
}
10981099

1100+
private Supplier<ExecutionStepInfo> createExecutionStepInfo(ExecutionContext executionContext, ExecutionStrategyParameters parameters) {
1101+
GraphQLFieldDefinition fieldDef = getFieldDef(executionContext, parameters, parameters.getField().getSingleField());
1102+
return FpKit.intraThreadMemoize(() -> createExecutionStepInfo(executionContext, parameters, fieldDef, null));
1103+
}
1104+
10991105
// Errors that result from the execution of deferred fields are kept in the deferred context only.
11001106
private static void addErrorToRightContext(GraphQLError error, ExecutionStrategyParameters parameters, ExecutionContext executionContext) {
11011107
if (parameters.getDeferredCallContext() != null) {

‎src/main/java/graphql/execution/incremental/DeferredExecutionSupport.java‎

Lines changed: 41 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -7,14 +7,18 @@
77
import graphql.ExecutionResultImpl;
88
import graphql.Internal;
99
import graphql.execution.ExecutionContext;
10+
import graphql.execution.ExecutionStepInfo;
1011
import graphql.execution.ExecutionStrategyParameters;
1112
import graphql.execution.FieldValueInfo;
1213
import graphql.execution.MergedField;
1314
import graphql.execution.MergedSelectionSet;
1415
import graphql.execution.ResultPath;
1516
import graphql.execution.instrumentation.Instrumentation;
17+
import graphql.execution.instrumentation.InstrumentationContext;
18+
import graphql.execution.instrumentation.parameters.InstrumentationFieldParameters;
1619
import graphql.incremental.IncrementalPayload;
1720
import graphql.util.FpKit;
21+
import org.jspecify.annotations.NonNull;
1822

1923
import java.util.Collections;
2024
import java.util.HashMap;
@@ -27,6 +31,8 @@
2731
import java.util.function.BiFunction;
2832
import java.util.function.Supplier;
2933

34+
import static graphql.execution.instrumentation.SimpleInstrumentationContext.nonNullCtx;
35+
3036
/**
3137
* The purpose of this class hierarchy is to encapsulate most of the logic for deferring field execution, thus
3238
* keeping the main execution strategy code clean and focused on the main execution logic.
@@ -59,16 +65,19 @@ class DeferredExecutionSupportImpl implements DeferredExecutionSupport {
5965
private final ExecutionStrategyParameters parameters;
6066
private final ExecutionContext executionContext;
6167
private final BiFunction<ExecutionContext, ExecutionStrategyParameters, CompletableFuture<FieldValueInfo>> resolveFieldWithInfoFn;
68+
private final BiFunction<ExecutionContext, ExecutionStrategyParameters, Supplier<ExecutionStepInfo>> executionStepInfoFn;
6269
private final Map<String, Supplier<CompletableFuture<DeferredFragmentCall.FieldWithExecutionResult>>> dfCache = new HashMap<>();
6370

6471
public DeferredExecutionSupportImpl(
6572
MergedSelectionSet mergedSelectionSet,
6673
ExecutionStrategyParameters parameters,
6774
ExecutionContext executionContext,
68-
BiFunction<ExecutionContext, ExecutionStrategyParameters, CompletableFuture<FieldValueInfo>> resolveFieldWithInfoFn
75+
BiFunction<ExecutionContext, ExecutionStrategyParameters, CompletableFuture<FieldValueInfo>> resolveFieldWithInfoFn,
76+
BiFunction<ExecutionContext, ExecutionStrategyParameters, Supplier<ExecutionStepInfo>> executionStepInfoFn
6977
) {
7078
this.executionContext = executionContext;
7179
this.resolveFieldWithInfoFn = resolveFieldWithInfoFn;
80+
this.executionStepInfoFn = executionStepInfoFn;
7281
ImmutableListMultimap.Builder<DeferredExecution, MergedField> deferredExecutionToFieldsBuilder = ImmutableListMultimap.builder();
7382
ImmutableSet.Builder<MergedField> deferredFieldsBuilder = ImmutableSet.builder();
7483
ImmutableList.Builder<String> nonDeferredFieldNamesBuilder = ImmutableList.builder();
@@ -153,37 +162,46 @@ private Supplier<CompletableFuture<DeferredFragmentCall.FieldWithExecutionResult
153162
}
154163
);
155164

156-
157-
Instrumentation instrumentation = executionContext.getInstrumentation();
158-
159-
instrumentation.beginDeferredField(executionContext.getInstrumentationState());
160-
161165
// todo: handle cached computations
162166
return dfCache.computeIfAbsent(
163167
currentField.getResultKey(),
164168
// The same field can be associated with multiple defer executions, so
165169
// we memoize the field resolution to avoid multiple calls to the same data fetcher
166-
key -> FpKit.interThreadMemoize(() -> {
167-
CompletableFuture<FieldValueInfo> fieldValueResult = resolveFieldWithInfoFn.apply(executionContext, executionStrategyParameters);
170+
key -> FpKit.interThreadMemoize(resolveDeferredFieldValue(currentField, executionContext, executionStrategyParameters)
171+
)
172+
);
173+
}
168174

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

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

174-
CompletableFuture<ExecutionResult> executionResultCF = fieldValueResult
175-
.thenCompose(fvi -> fvi
176-
.getFieldValueFuture()
177-
.thenApply(fv -> ExecutionResultImpl.newExecutionResult().data(fv).build())
178-
);
184+
CompletableFuture<FieldValueInfo> fieldValueResult = resolveFieldWithInfoFn.apply(this.executionContext, executionStrategyParameters);
179185

180-
return executionResultCF
181-
.thenApply(executionResult ->
182-
new DeferredFragmentCall.FieldWithExecutionResult(currentField.getResultKey(), executionResult)
183-
);
184-
}
185-
)
186-
);
186+
deferredFieldCtx.onDispatched();
187+
188+
fieldValueResult.whenComplete((fieldValueInfo, throwable) -> {
189+
this.executionContext.getDataLoaderDispatcherStrategy().deferredOnFieldValue(currentField.getResultKey(), fieldValueInfo, throwable, executionStrategyParameters);
190+
deferredFieldCtx.onCompleted(fieldValueInfo, throwable);
191+
});
192+
193+
194+
CompletableFuture<ExecutionResult> executionResultCF = fieldValueResult
195+
.thenCompose(fvi -> fvi
196+
.getFieldValueFuture()
197+
.thenApply(fv -> ExecutionResultImpl.newExecutionResult().data(fv).build())
198+
);
199+
200+
return executionResultCF
201+
.thenApply(executionResult ->
202+
new DeferredFragmentCall.FieldWithExecutionResult(currentField.getResultKey(), executionResult)
203+
);
204+
};
187205
}
188206
}
189207

‎src/main/java/graphql/execution/instrumentation/ChainedInstrumentation.java‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -166,11 +166,10 @@ public ExecutionStrategyInstrumentationContext beginExecutionStrategy(Instrument
166166

167167
@ExperimentalApi
168168
@Override
169-
public InstrumentationContext<Object> beginDeferredField(InstrumentationState instrumentationState) {
170-
return new ChainedDeferredExecutionStrategyInstrumentationContext(chainedMapAndDropNulls(instrumentationState, Instrumentation::beginDeferredField));
169+
public InstrumentationContext<Object> beginDeferredField(InstrumentationFieldParameters parameters, InstrumentationState state) {
170+
return chainedCtx(state, (instrumentation, specificState) -> instrumentation.beginDeferredField(parameters, specificState));
171171
}
172172

173-
174173
@Override
175174
public InstrumentationContext<ExecutionResult> beginSubscribedFieldEvent(InstrumentationFieldParameters parameters, InstrumentationState state) {
176175
return chainedCtx(state, (instrumentation, specificState) -> instrumentation.beginSubscribedFieldEvent(parameters, specificState));

‎src/main/java/graphql/execution/instrumentation/Instrumentation.java‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -153,12 +153,13 @@ default ExecuteObjectInstrumentationContext beginExecuteObject(InstrumentationEx
153153
* <p>
154154
* This is an EXPERIMENTAL instrumentation callback. The method signature will definitely change.
155155
*
156-
* @param state the state created during the call to {@link #createStateAsync(InstrumentationCreateStateParameters)}
156+
* @param parameters the parameters to this step
157+
* @param state the state created during the call to {@link #createStateAsync(InstrumentationCreateStateParameters)}
157158
*
158-
* @return a nullable {@link ExecutionStrategyInstrumentationContext} object that will be called back when the step ends (assuming it's not null)
159+
* @return a nullable {@link InstrumentationContext} object that will be called back when the step ends (assuming it's not null)
159160
*/
160161
@ExperimentalApi
161-
default InstrumentationContext<Object> beginDeferredField(InstrumentationState state) {
162+
default InstrumentationContext<Object> beginDeferredField(InstrumentationFieldParameters parameters, InstrumentationState state) {
162163
return noOp();
163164
}
164165

‎src/main/java/graphql/execution/instrumentation/NoContextChainedInstrumentation.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,11 @@ public ExecutionStrategyInstrumentationContext beginExecutionStrategy(Instrument
8383
return runAll(state, (instrumentation, specificState) -> instrumentation.beginExecuteObject(parameters, specificState));
8484
}
8585

86+
@Override
87+
public InstrumentationContext<Object> beginDeferredField(InstrumentationFieldParameters parameters, InstrumentationState state) {
88+
return runAll(state, (instrumentation, specificState) -> instrumentation.beginDeferredField(parameters, specificState));
89+
}
90+
8691
@Override
8792
public InstrumentationContext<ExecutionResult> beginSubscribedFieldEvent(InstrumentationFieldParameters parameters, InstrumentationState state) {
8893
return runAll(state, (instrumentation, specificState) -> instrumentation.beginSubscribedFieldEvent(parameters, specificState));

‎src/test/groovy/graphql/TestUtil.groovy‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@ package graphql
22

33
import graphql.execution.MergedField
44
import graphql.execution.MergedSelectionSet
5+
import graphql.execution.pubsub.CapturingSubscriber
6+
import graphql.incremental.DelayedIncrementalPartialResult
7+
import graphql.incremental.IncrementalExecutionResult
58
import graphql.introspection.Introspection.DirectiveLocation
69
import graphql.language.Document
710
import graphql.language.Field
@@ -31,6 +34,8 @@ import graphql.schema.idl.TypeRuntimeWiring
3134
import graphql.schema.idl.WiringFactory
3235
import graphql.schema.idl.errors.SchemaProblem
3336
import groovy.json.JsonOutput
37+
import org.awaitility.Awaitility
38+
import org.reactivestreams.Publisher
3439

3540
import java.util.function.Supplier
3641
import java.util.stream.Collectors
@@ -323,4 +328,20 @@ class TestUtil {
323328
return rn.nextInt(max - min + 1) + min
324329
}
325330

331+
332+
static List<Map<String, Object>> getIncrementalResults(IncrementalExecutionResult initialResult) {
333+
Publisher<DelayedIncrementalPartialResult> deferredResultStream = initialResult.incrementalItemPublisher
334+
335+
def subscriber = new CapturingSubscriber<DelayedIncrementalPartialResult>()
336+
337+
deferredResultStream.subscribe(subscriber)
338+
339+
Awaitility.await().untilTrue(subscriber.isDone())
340+
if (subscriber.throwable != null) {
341+
throw new RuntimeException(subscriber.throwable)
342+
}
343+
return subscriber.getEvents()
344+
.collect { it.toSpecification() }
345+
}
346+
326347
}

‎src/test/groovy/graphql/execution/instrumentation/InstrumentationTest.groovy‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,13 @@ import graphql.ExecutionInput
44
import graphql.ExecutionResult
55
import graphql.GraphQL
66
import graphql.StarWarsSchema
7+
import graphql.TestUtil
78
import graphql.execution.AsyncExecutionStrategy
89
import graphql.execution.instrumentation.parameters.InstrumentationCreateStateParameters
910
import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters
1011
import graphql.execution.instrumentation.parameters.InstrumentationExecutionStrategyParameters
1112
import graphql.execution.instrumentation.parameters.InstrumentationFieldFetchParameters
13+
import graphql.incremental.IncrementalExecutionResult
1214
import graphql.language.AstPrinter
1315
import graphql.parser.Parser
1416
import graphql.schema.DataFetcher
@@ -496,4 +498,80 @@ class InstrumentationTest extends Specification {
496498
then:
497499
er.extensions == [i1: "I1"]
498500
}
501+
502+
def "can instrumented deferred fields"() {
503+
504+
given:
505+
506+
def query = """
507+
{
508+
hero {
509+
id
510+
... @defer(label: "name") {
511+
name
512+
}
513+
}
514+
}
515+
"""
516+
517+
518+
when:
519+
520+
def instrumentation = new ModernTestingInstrumentation()
521+
522+
def graphQL = GraphQL
523+
.newGraphQL(StarWarsSchema.starWarsSchema)
524+
.queryExecutionStrategy(new AsyncExecutionStrategy())
525+
.instrumentation(instrumentation)
526+
.build()
527+
528+
def ei = ExecutionInput.newExecutionInput(query).graphQLContext { it ->
529+
GraphQL.unusualConfiguration(it).incrementalSupport().enableIncrementalSupport(true)
530+
}.build()
531+
532+
IncrementalExecutionResult incrementalER = graphQL.execute(ei) as IncrementalExecutionResult
533+
//
534+
// cause the defer Publish to be finished
535+
def results = TestUtil.getIncrementalResults(incrementalER)
536+
537+
538+
then:
539+
540+
instrumentation.executionList == ["start:execution",
541+
"start:parse",
542+
"end:parse",
543+
"start:validation",
544+
"end:validation",
545+
"start:execute-operation",
546+
"start:execution-strategy",
547+
"start:field-hero",
548+
"start:fetch-hero",
549+
"end:fetch-hero",
550+
"start:complete-hero",
551+
"start:execute-object",
552+
"start:field-id",
553+
"start:fetch-id",
554+
"end:fetch-id",
555+
"start:complete-id",
556+
"end:complete-id",
557+
"end:field-id",
558+
559+
"end:execute-object",
560+
"end:complete-hero",
561+
"end:field-hero",
562+
"end:execution-strategy",
563+
"end:execute-operation",
564+
"end:execution",
565+
//
566+
// the deferred field resolving now happens after the operation has come back
567+
"start:deferred-field-name",
568+
"start:field-name",
569+
"start:fetch-name",
570+
"end:fetch-name",
571+
"start:complete-name",
572+
"end:complete-name",
573+
"end:field-name",
574+
"end:deferred-field-name",
575+
]
576+
}
499577
}

‎src/test/groovy/graphql/execution/instrumentation/ModernTestingInstrumentation.groovy‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,12 @@ class ModernTestingInstrumentation implements Instrumentation {
103103
return new TestingInstrumentContext("complete-list-$parameters.field.name", executionList, throwableList, useOnDispatch)
104104
}
105105

106+
@Override
107+
InstrumentationContext<Object> beginDeferredField(InstrumentationFieldParameters parameters, InstrumentationState state) {
108+
assert state == instrumentationState
109+
return new TestingInstrumentContext("deferred-field-$parameters.field.name", executionList, throwableList, useOnDispatch)
110+
}
111+
106112
@Override
107113
ExecutionInput instrumentExecutionInput(ExecutionInput executionInput, InstrumentationExecutionParameters parameters, InstrumentationState state) {
108114
assert state == instrumentationState

0 commit comments

Comments
 (0)