Skip to content

Commit ce83384

Browse files
bbakermanbbakerman
committed
Added support for async create state in instrumentmentions
1 parent 72645a5 commit ce83384

File tree

7 files changed

+215
-37
lines changed

7 files changed

+215
-37
lines changed

‎src/main/java/graphql/GraphQL.java‎

Lines changed: 27 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package graphql;
22

33
import graphql.execution.AbortExecutionException;
4+
import graphql.execution.Async;
45
import graphql.execution.AsyncExecutionStrategy;
56
import graphql.execution.AsyncSerialExecutionStrategy;
67
import graphql.execution.DataFetcherExceptionHandler;
@@ -421,31 +422,33 @@ public CompletableFuture<ExecutionResult> executeAsync(ExecutionInput executionI
421422
if (logNotSafe.isDebugEnabled()) {
422423
logNotSafe.debug("Executing request. operation name: '{}'. query: '{}'. variables '{}'", executionInput.getOperationName(), executionInput.getQuery(), executionInput.getVariables());
423424
}
424-
executionInput = ensureInputHasId(executionInput);
425+
ExecutionInput executionInputWithId = ensureInputHasId(executionInput);
425426

426-
InstrumentationState instrumentationState = instrumentation.createState(new InstrumentationCreateStateParameters(this.graphQLSchema, executionInput));
427-
try {
428-
InstrumentationExecutionParameters inputInstrumentationParameters = new InstrumentationExecutionParameters(executionInput, this.graphQLSchema, instrumentationState);
429-
executionInput = instrumentation.instrumentExecutionInput(executionInput, inputInstrumentationParameters, instrumentationState);
430-
431-
CompletableFuture<ExecutionResult> beginExecutionCF = new CompletableFuture<>();
432-
InstrumentationExecutionParameters instrumentationParameters = new InstrumentationExecutionParameters(executionInput, this.graphQLSchema, instrumentationState);
433-
InstrumentationContext<ExecutionResult> executionInstrumentation = nonNullCtx(instrumentation.beginExecution(instrumentationParameters, instrumentationState));
434-
executionInstrumentation.onDispatched(beginExecutionCF);
435-
436-
GraphQLSchema graphQLSchema = instrumentation.instrumentSchema(this.graphQLSchema, instrumentationParameters, instrumentationState);
437-
438-
CompletableFuture<ExecutionResult> executionResult = parseValidateAndExecute(executionInput, graphQLSchema, instrumentationState);
439-
//
440-
// finish up instrumentation
441-
executionResult = executionResult.whenComplete(completeInstrumentationCtxCF(executionInstrumentation, beginExecutionCF));
442-
//
443-
// allow instrumentation to tweak the result
444-
executionResult = executionResult.thenCompose(result -> instrumentation.instrumentExecutionResult(result, instrumentationParameters, instrumentationState));
445-
return executionResult;
446-
} catch (AbortExecutionException abortException) {
447-
return handleAbortException(executionInput, instrumentationState, abortException);
448-
}
427+
CompletableFuture<InstrumentationState> instrumentationStateCF = instrumentation.createStateAsync(new InstrumentationCreateStateParameters(this.graphQLSchema, executionInput));
428+
return Async.orNullCompletedFuture(instrumentationStateCF).thenCompose(instrumentationState -> {
429+
try {
430+
InstrumentationExecutionParameters inputInstrumentationParameters = new InstrumentationExecutionParameters(executionInputWithId, this.graphQLSchema, instrumentationState);
431+
ExecutionInput instrumentedExecutionInput = instrumentation.instrumentExecutionInput(executionInputWithId, inputInstrumentationParameters, instrumentationState);
432+
433+
CompletableFuture<ExecutionResult> beginExecutionCF = new CompletableFuture<>();
434+
InstrumentationExecutionParameters instrumentationParameters = new InstrumentationExecutionParameters(instrumentedExecutionInput, this.graphQLSchema, instrumentationState);
435+
InstrumentationContext<ExecutionResult> executionInstrumentation = nonNullCtx(instrumentation.beginExecution(instrumentationParameters, instrumentationState));
436+
executionInstrumentation.onDispatched(beginExecutionCF);
437+
438+
GraphQLSchema graphQLSchema = instrumentation.instrumentSchema(this.graphQLSchema, instrumentationParameters, instrumentationState);
439+
440+
CompletableFuture<ExecutionResult> executionResult = parseValidateAndExecute(instrumentedExecutionInput, graphQLSchema, instrumentationState);
441+
//
442+
// finish up instrumentation
443+
executionResult = executionResult.whenComplete(completeInstrumentationCtxCF(executionInstrumentation, beginExecutionCF));
444+
//
445+
// allow instrumentation to tweak the result
446+
executionResult = executionResult.thenCompose(result -> instrumentation.instrumentExecutionResult(result, instrumentationParameters, instrumentationState));
447+
return executionResult;
448+
} catch (AbortExecutionException abortException) {
449+
return handleAbortException(executionInput, instrumentationState, abortException);
450+
}
451+
});
449452
}
450453

451454
private CompletableFuture<ExecutionResult> handleAbortException(ExecutionInput executionInput, InstrumentationState instrumentationState, AbortExecutionException abortException) {

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22

33
import graphql.Assert;
44
import graphql.Internal;
5+
import org.jetbrains.annotations.NotNull;
6+
import org.jetbrains.annotations.Nullable;
57

68
import java.util.ArrayList;
79
import java.util.Collection;
@@ -207,4 +209,15 @@ public static <T> CompletableFuture<T> exceptionallyCompletedFuture(Throwable ex
207209
return result;
208210
}
209211

212+
/**
213+
* If the passed in CompletableFuture is null then it creates a CompletableFuture that resolves to null
214+
*
215+
* @param completableFuture the CF to use
216+
* @param <T> for two
217+
*
218+
* @return the completableFuture if it's not null or one that always resoles to null
219+
*/
220+
public static <T> @NotNull CompletableFuture<T> orNullCompletedFuture(@Nullable CompletableFuture<T> completableFuture) {
221+
return completableFuture != null ? completableFuture : CompletableFuture.completedFuture(null);
222+
}
210223
}

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

Lines changed: 30 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import graphql.schema.GraphQLSchema;
2323
import graphql.validation.ValidationError;
2424
import org.jetbrains.annotations.NotNull;
25+
import org.jetbrains.annotations.Nullable;
2526

2627
import java.util.Arrays;
2728
import java.util.List;
@@ -80,10 +81,19 @@ private <T> InstrumentationContext<T> chainedCtx(Function<Instrumentation, Instr
8081
return new ChainedInstrumentationContext<>(mapAndDropNulls(instrumentations, mapper));
8182
}
8283

84+
@Override
85+
public InstrumentationState createState() {
86+
return Assert.assertShouldNeverHappen("createStateAsync should only ever be used");
87+
}
88+
89+
@Override
90+
public @Nullable InstrumentationState createState(InstrumentationCreateStateParameters parameters) {
91+
return Assert.assertShouldNeverHappen("createStateAsync should only ever be used");
92+
}
8393

8494
@Override
85-
public InstrumentationState createState(InstrumentationCreateStateParameters parameters) {
86-
return new ChainedInstrumentationState(instrumentations, parameters);
95+
public @NotNull CompletableFuture<InstrumentationState> createStateAsync(InstrumentationCreateStateParameters parameters) {
96+
return ChainedInstrumentationState.combineAll(instrumentations, parameters);
8797
}
8898

8999
@Override
@@ -349,18 +359,31 @@ public CompletableFuture<ExecutionResult> instrumentExecutionResult(ExecutionRes
349359
}
350360

351361
static class ChainedInstrumentationState implements InstrumentationState {
352-
private final Map<Instrumentation, InstrumentationState> instrumentationStates;
362+
private final Map<Instrumentation, InstrumentationState> instrumentationToStates;
353363

354364

355-
private ChainedInstrumentationState(List<Instrumentation> instrumentations, InstrumentationCreateStateParameters parameters) {
356-
instrumentationStates = Maps.newLinkedHashMapWithExpectedSize(instrumentations.size());
357-
instrumentations.forEach(i -> instrumentationStates.put(i, i.createState(parameters)));
365+
private ChainedInstrumentationState(List<Instrumentation> instrumentations, List<InstrumentationState> instrumentationStates) {
366+
instrumentationToStates = Maps.newLinkedHashMapWithExpectedSize(instrumentations.size());
367+
for (int i = 0; i < instrumentations.size(); i++) {
368+
Instrumentation instrumentation = instrumentations.get(i);
369+
InstrumentationState instrumentationState = instrumentationStates.get(i);
370+
instrumentationToStates.put(instrumentation, instrumentationState);
371+
}
358372
}
359373

360374
private InstrumentationState getState(Instrumentation instrumentation) {
361-
return instrumentationStates.get(instrumentation);
375+
return instrumentationToStates.get(instrumentation);
362376
}
363377

378+
private static CompletableFuture<InstrumentationState> combineAll(List<Instrumentation> instrumentations, InstrumentationCreateStateParameters parameters) {
379+
Async.CombinedBuilder<InstrumentationState> builder = Async.ofExpectedSize(instrumentations.size());
380+
for (Instrumentation instrumentation : instrumentations) {
381+
// state can be null including the CF so handle that
382+
CompletableFuture<InstrumentationState> stateCF = Async.orNullCompletedFuture(instrumentation.createStateAsync(parameters));
383+
builder.add(stateCF);
384+
}
385+
return builder.await().thenApply(instrumentationStates -> new ChainedInstrumentationState(instrumentations, instrumentationStates));
386+
}
364387
}
365388

366389
private static class ChainedInstrumentationContext<T> implements InstrumentationContext<T> {

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,11 +63,26 @@ default InstrumentationState createState() {
6363
*
6464
* @return a state object that is passed to each method
6565
*/
66+
@Deprecated
67+
@DeprecatedAt("2023-08-25")
6668
@Nullable
6769
default InstrumentationState createState(InstrumentationCreateStateParameters parameters) {
6870
return createState();
6971
}
7072

73+
/**
74+
* This will be called just before execution to create an object, in an asynchronous manner, that is given back to all instrumentation methods
75+
* to allow them to have per execution request state
76+
*
77+
* @param parameters the parameters to this step
78+
*
79+
* @return a state object that is passed to each method
80+
*/
81+
@Nullable
82+
default CompletableFuture<InstrumentationState> createStateAsync(InstrumentationCreateStateParameters parameters) {
83+
return CompletableFuture.completedFuture(createState(parameters));
84+
}
85+
7186
/**
7287
* This is called right at the start of query execution, and it's the first step in the instrumentation chain.
7388
*

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,12 @@ public InstrumentationState createState() {
5656
return null;
5757
}
5858

59+
@Override
60+
public @Nullable CompletableFuture<InstrumentationState> createStateAsync(InstrumentationCreateStateParameters parameters) {
61+
InstrumentationState state = createState(parameters);
62+
return state == null ? null : CompletableFuture.completedFuture(state);
63+
}
64+
5965
@Override
6066
public @NotNull InstrumentationContext<ExecutionResult> beginExecution(InstrumentationExecutionParameters parameters) {
6167
return assertShouldNeverHappen("The deprecated " + "beginExecution" + " was called");

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

Lines changed: 71 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,13 @@
11
package graphql.execution.instrumentation
22

3+
import graphql.ExecutionInput
34
import graphql.ExecutionResult
45
import graphql.GraphQL
56
import graphql.StarWarsSchema
67
import graphql.execution.AsyncExecutionStrategy
7-
import graphql.execution.instrumentation.parameters.InstrumentationExecuteOperationParameters
8+
import graphql.execution.instrumentation.parameters.InstrumentationCreateStateParameters
89
import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters
9-
import graphql.execution.instrumentation.parameters.InstrumentationExecutionStrategyParameters
10-
import graphql.execution.instrumentation.parameters.InstrumentationFieldFetchParameters
11-
import graphql.execution.instrumentation.parameters.InstrumentationFieldParameters
1210
import graphql.execution.instrumentation.parameters.InstrumentationValidationParameters
13-
import graphql.language.Document
14-
import graphql.schema.DataFetcher
1511
import graphql.validation.ValidationError
1612
import spock.lang.Specification
1713

@@ -279,6 +275,75 @@ class ChainedInstrumentationStateTest extends Specification {
279275

280276
}
281277

278+
279+
class StringInstrumentationState implements InstrumentationState {
280+
StringInstrumentationState(String value) {
281+
this.value = value
282+
}
283+
284+
String value
285+
}
286+
287+
def "can have an multiple async createState() calls in play"() {
288+
289+
290+
given:
291+
292+
def query = '''query Q($var: String!) {
293+
human(id: $var) {
294+
id
295+
name
296+
}
297+
}
298+
'''
299+
300+
301+
def instrumentation1 = new SimplePerformantInstrumentation() {
302+
@Override
303+
CompletableFuture<InstrumentationState> createStateAsync(InstrumentationCreateStateParameters parameters) {
304+
return CompletableFuture.supplyAsync {
305+
return new StringInstrumentationState("I1")
306+
} as CompletableFuture<InstrumentationState>
307+
}
308+
309+
@Override
310+
CompletableFuture<ExecutionResult> instrumentExecutionResult(ExecutionResult executionResult, InstrumentationExecutionParameters parameters, InstrumentationState state) {
311+
return CompletableFuture.completedFuture(
312+
executionResult.transform { it.addExtension("i1", ((StringInstrumentationState) state).value) }
313+
)
314+
}
315+
}
316+
def instrumentation2 = new SimplePerformantInstrumentation() {
317+
@Override
318+
CompletableFuture<InstrumentationState> createStateAsync(InstrumentationCreateStateParameters parameters) {
319+
return CompletableFuture.supplyAsync {
320+
return new StringInstrumentationState("I2")
321+
} as CompletableFuture<InstrumentationState>
322+
}
323+
324+
@Override
325+
CompletableFuture<ExecutionResult> instrumentExecutionResult(ExecutionResult executionResult, InstrumentationExecutionParameters parameters, InstrumentationState state) {
326+
return CompletableFuture.completedFuture(
327+
executionResult.transform { it.addExtension("i2", ((StringInstrumentationState) state).value) }
328+
)
329+
}
330+
331+
}
332+
333+
def graphQL = GraphQL
334+
.newGraphQL(StarWarsSchema.starWarsSchema)
335+
.instrumentation(new ChainedInstrumentation([instrumentation1, instrumentation2]))
336+
.doNotAddDefaultInstrumentations() // important, otherwise a chained one wil be used
337+
.build()
338+
339+
when:
340+
def variables = [var: "1001"]
341+
def er = graphQL.execute(ExecutionInput.newExecutionInput().query(query).variables(variables)) // Luke
342+
343+
then:
344+
er.extensions == [i1: "I1", i2: "I2"]
345+
}
346+
282347
private void assertCalls(NamedInstrumentation instrumentation) {
283348
assert instrumentation.dfInvocations[0].getFieldDefinition().name == 'hero'
284349
assert instrumentation.dfInvocations[0].getExecutionStepInfo().getPath().toList() == ['hero']

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

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import graphql.ExecutionResult
55
import graphql.GraphQL
66
import graphql.StarWarsSchema
77
import graphql.execution.AsyncExecutionStrategy
8+
import graphql.execution.instrumentation.parameters.InstrumentationCreateStateParameters
89
import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters
910
import graphql.execution.instrumentation.parameters.InstrumentationExecutionStrategyParameters
1011
import graphql.execution.instrumentation.parameters.InstrumentationFieldFetchParameters
@@ -404,4 +405,56 @@ class InstrumentationTest extends Specification {
404405

405406
instrumentation.executionList == expected
406407
}
408+
409+
class StringInstrumentationState implements InstrumentationState {
410+
StringInstrumentationState(String value) {
411+
this.value = value
412+
}
413+
414+
String value
415+
}
416+
417+
def "can have an single async createState() in play"() {
418+
419+
420+
given:
421+
422+
def query = '''query Q($var: String!) {
423+
human(id: $var) {
424+
id
425+
name
426+
}
427+
}
428+
'''
429+
430+
431+
def instrumentation1 = new SimplePerformantInstrumentation() {
432+
@Override
433+
CompletableFuture<InstrumentationState> createStateAsync(InstrumentationCreateStateParameters parameters) {
434+
return CompletableFuture.supplyAsync {
435+
return new StringInstrumentationState("I1")
436+
} as CompletableFuture<InstrumentationState>
437+
}
438+
439+
@Override
440+
CompletableFuture<ExecutionResult> instrumentExecutionResult(ExecutionResult executionResult, InstrumentationExecutionParameters parameters, InstrumentationState state) {
441+
return CompletableFuture.completedFuture(
442+
executionResult.transform { it.addExtension("i1", ((StringInstrumentationState) state).value) }
443+
)
444+
}
445+
}
446+
447+
def graphQL = GraphQL
448+
.newGraphQL(StarWarsSchema.starWarsSchema)
449+
.instrumentation(instrumentation1)
450+
.doNotAddDefaultInstrumentations() // important, otherwise a chained one wil be used
451+
.build()
452+
453+
when:
454+
def variables = [var: "1001"]
455+
def er = graphQL.execute(ExecutionInput.newExecutionInput().query(query).variables(variables)) // Luke
456+
457+
then:
458+
er.extensions == [i1: "I1"]
459+
}
407460
}

0 commit comments

Comments
 (0)