Skip to content

Commit a590cfa

Browse files
andimarekandimarek
committed
refactor engine tracking and simplify it
1 parent a25f93b commit a590cfa

17 files changed

+617
-503
lines changed

‎src/main/java/graphql/EngineRunningState.java‎

Lines changed: 203 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,203 @@
1+
package graphql;
2+
3+
import graphql.execution.EngineRunningObserver;
4+
import graphql.execution.ExecutionId;
5+
import org.jspecify.annotations.Nullable;
6+
7+
import java.util.concurrent.CompletableFuture;
8+
import java.util.concurrent.CompletionStage;
9+
import java.util.concurrent.atomic.AtomicInteger;
10+
import java.util.function.BiConsumer;
11+
import java.util.function.BiFunction;
12+
import java.util.function.Function;
13+
import java.util.function.Supplier;
14+
15+
import static graphql.Assert.assertTrue;
16+
import static graphql.execution.EngineRunningObserver.RunningState.NOT_RUNNING;
17+
import static graphql.execution.EngineRunningObserver.RunningState.RUNNING;
18+
19+
@Internal
20+
public class EngineRunningState {
21+
22+
@Nullable
23+
private final EngineRunningObserver engineRunningObserver;
24+
@Nullable
25+
private final GraphQLContext graphQLContext;
26+
@Nullable
27+
private volatile ExecutionId executionId;
28+
29+
private final AtomicInteger isRunning = new AtomicInteger(0);
30+
31+
@VisibleForTesting
32+
public EngineRunningState() {
33+
this.engineRunningObserver = null;
34+
this.graphQLContext = null;
35+
this.executionId = null;
36+
}
37+
38+
public EngineRunningState(ExecutionInput executionInput) {
39+
EngineRunningObserver engineRunningObserver = executionInput.getGraphQLContext().get(EngineRunningObserver.ENGINE_RUNNING_OBSERVER_KEY);
40+
if (engineRunningObserver != null) {
41+
this.engineRunningObserver = engineRunningObserver;
42+
this.graphQLContext = executionInput.getGraphQLContext();
43+
this.executionId = executionInput.getExecutionId();
44+
} else {
45+
this.engineRunningObserver = null;
46+
this.graphQLContext = null;
47+
this.executionId = null;
48+
}
49+
}
50+
51+
public <U, T> CompletableFuture<U> handle(CompletableFuture<T> src, BiFunction<? super T, Throwable, ? extends U> fn) {
52+
if (engineRunningObserver == null) {
53+
return src.handle(fn);
54+
}
55+
src = observeCompletableFutureStart(src);
56+
CompletableFuture<U> result = src.handle((t, throwable) -> {
57+
// because we added an artificial dependent CF on src (in observeCompletableFutureStart) , a throwable is a CompletionException
58+
// that needs to be unwrapped
59+
if (throwable != null) {
60+
throwable = throwable.getCause();
61+
}
62+
return fn.apply(t, throwable);
63+
});
64+
observerCompletableFutureEnd(src);
65+
return result;
66+
}
67+
68+
public <T> CompletableFuture<T> whenComplete(CompletableFuture<T> src, BiConsumer<? super T, ? super Throwable> fn) {
69+
if (engineRunningObserver == null) {
70+
return src.whenComplete(fn);
71+
}
72+
src = observeCompletableFutureStart(src);
73+
CompletableFuture<T> result = src.whenComplete((t, throwable) -> {
74+
// because we added an artificial dependent CF on src (in observeCompletableFutureStart) , a throwable is a CompletionException
75+
// that needs to be unwrapped
76+
if (throwable != null) {
77+
throwable = throwable.getCause();
78+
}
79+
fn.accept(t, throwable);
80+
});
81+
observerCompletableFutureEnd(src);
82+
return result;
83+
}
84+
85+
public <U, T> CompletableFuture<U> compose(CompletableFuture<T> src, Function<? super T, ? extends CompletionStage<U>> fn) {
86+
if (engineRunningObserver == null) {
87+
return src.thenCompose(fn);
88+
}
89+
CompletableFuture<U> result = new CompletableFuture<>();
90+
src = observeCompletableFutureStart(src);
91+
src.whenComplete((u, t) -> run(() -> {
92+
CompletionStage<U> innerCF = fn.apply(u).toCompletableFuture();
93+
innerCF.whenComplete((u1, t1) -> run(() -> {
94+
if (t1 != null) {
95+
result.completeExceptionally(t1);
96+
} else {
97+
result.complete(u1);
98+
}
99+
}));
100+
}));
101+
observerCompletableFutureEnd(src);
102+
return result;
103+
}
104+
105+
public <T, U> CompletableFuture<U> chain(CompletableFuture<T> src, Function<CompletableFuture<T>, CompletableFuture<U>> fn) {
106+
if (engineRunningObserver == null) {
107+
return fn.apply(src);
108+
}
109+
src = observeCompletableFutureStart(src);
110+
CompletableFuture<U> result = fn.apply(src);
111+
observerCompletableFutureEnd(result);
112+
return result;
113+
}
114+
115+
116+
public <T> CompletableFuture<T> observeCompletableFutureStart(CompletableFuture<T> future) {
117+
if (engineRunningObserver == null) {
118+
return future;
119+
}
120+
future = future.thenApply(Function.identity());
121+
incrementRunningWhenCompleted(future);
122+
return future;
123+
}
124+
125+
public void observerCompletableFutureEnd(CompletableFuture<?> future) {
126+
if (engineRunningObserver == null) {
127+
return;
128+
}
129+
decrementRunningWhenCompleted(future);
130+
}
131+
132+
133+
private void incrementRunningWhenCompleted(CompletableFuture<?> cf) {
134+
cf.whenComplete((result, throwable) -> {
135+
incrementRunning();
136+
});
137+
}
138+
139+
private void decrementRunningWhenCompleted(CompletableFuture<?> cf) {
140+
cf.whenComplete((result, throwable) -> {
141+
decrementRunning();
142+
});
143+
144+
}
145+
146+
private void decrementRunning() {
147+
assertTrue(isRunning.get() > 0);
148+
if (isRunning.decrementAndGet() == 0) {
149+
changeOfState(NOT_RUNNING);
150+
}
151+
}
152+
153+
154+
private void incrementRunning() {
155+
if (engineRunningObserver == null) {
156+
return;
157+
}
158+
assertTrue(isRunning.get() >= 0);
159+
if (isRunning.incrementAndGet() == 1) {
160+
changeOfState(RUNNING);
161+
}
162+
163+
}
164+
165+
166+
public void updateExecutionId(ExecutionId executionId) {
167+
if (engineRunningObserver == null) {
168+
return;
169+
}
170+
this.executionId = executionId;
171+
}
172+
173+
private void changeOfState(EngineRunningObserver.RunningState runningState) {
174+
engineRunningObserver.runningStateChanged(executionId, graphQLContext, runningState);
175+
}
176+
177+
public void run(Runnable runnable) {
178+
if (engineRunningObserver == null) {
179+
runnable.run();
180+
return;
181+
}
182+
incrementRunning();
183+
try {
184+
runnable.run();
185+
} finally {
186+
decrementRunning();
187+
}
188+
}
189+
190+
public <T> T call(Supplier<T> supplier) {
191+
if (engineRunningObserver == null) {
192+
return supplier.get();
193+
}
194+
incrementRunning();
195+
try {
196+
return supplier.get();
197+
} finally {
198+
decrementRunning();
199+
}
200+
}
201+
202+
203+
}

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

Lines changed: 40 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -412,41 +412,44 @@ public CompletableFuture<ExecutionResult> executeAsync(UnaryOperator<ExecutionIn
412412
* @return a promise to an {@link ExecutionResult} which can include errors
413413
*/
414414
public CompletableFuture<ExecutionResult> executeAsync(ExecutionInput executionInput) {
415-
ExecutionInput executionInputWithId = ensureInputHasId(executionInput);
416-
417-
CompletableFuture<InstrumentationState> instrumentationStateCF = instrumentation.createStateAsync(new InstrumentationCreateStateParameters(this.graphQLSchema, executionInputWithId));
418-
return Async.orNullCompletedFuture(instrumentationStateCF).thenCompose(instrumentationState -> {
419-
try {
420-
InstrumentationExecutionParameters inputInstrumentationParameters = new InstrumentationExecutionParameters(executionInputWithId, this.graphQLSchema);
421-
ExecutionInput instrumentedExecutionInput = instrumentation.instrumentExecutionInput(executionInputWithId, inputInstrumentationParameters, instrumentationState);
422-
423-
InstrumentationExecutionParameters instrumentationParameters = new InstrumentationExecutionParameters(instrumentedExecutionInput, this.graphQLSchema);
424-
InstrumentationContext<ExecutionResult> executionInstrumentation = nonNullCtx(instrumentation.beginExecution(instrumentationParameters, instrumentationState));
425-
executionInstrumentation.onDispatched();
426-
427-
GraphQLSchema graphQLSchema = instrumentation.instrumentSchema(this.graphQLSchema, instrumentationParameters, instrumentationState);
428-
429-
CompletableFuture<ExecutionResult> executionResult = parseValidateAndExecute(instrumentedExecutionInput, graphQLSchema, instrumentationState);
430-
//
431-
// finish up instrumentation
432-
executionResult = executionResult.whenComplete(completeInstrumentationCtxCF(executionInstrumentation));
433-
//
434-
// allow instrumentation to tweak the result
435-
executionResult = executionResult.thenCompose(result -> instrumentation.instrumentExecutionResult(result, instrumentationParameters, instrumentationState));
436-
return executionResult;
437-
} catch (AbortExecutionException abortException) {
438-
return handleAbortException(executionInput, instrumentationState, abortException);
439-
}
415+
EngineRunningState engineRunningState = new EngineRunningState(executionInput);
416+
return engineRunningState.call(() -> {
417+
ExecutionInput executionInputWithId = ensureInputHasId(executionInput);
418+
engineRunningState.updateExecutionId(executionInputWithId.getExecutionId());
419+
420+
CompletableFuture<InstrumentationState> instrumentationStateCF = instrumentation.createStateAsync(new InstrumentationCreateStateParameters(this.graphQLSchema, executionInputWithId));
421+
instrumentationStateCF = Async.orNullCompletedFuture(instrumentationStateCF);
422+
423+
return engineRunningState.compose(instrumentationStateCF, (instrumentationState -> {
424+
try {
425+
InstrumentationExecutionParameters inputInstrumentationParameters = new InstrumentationExecutionParameters(executionInputWithId, this.graphQLSchema);
426+
ExecutionInput instrumentedExecutionInput = instrumentation.instrumentExecutionInput(executionInputWithId, inputInstrumentationParameters, instrumentationState);
427+
428+
InstrumentationExecutionParameters instrumentationParameters = new InstrumentationExecutionParameters(instrumentedExecutionInput, this.graphQLSchema);
429+
InstrumentationContext<ExecutionResult> executionInstrumentation = nonNullCtx(instrumentation.beginExecution(instrumentationParameters, instrumentationState));
430+
executionInstrumentation.onDispatched();
431+
432+
GraphQLSchema graphQLSchema = instrumentation.instrumentSchema(this.graphQLSchema, instrumentationParameters, instrumentationState);
433+
434+
CompletableFuture<ExecutionResult> executionResult = parseValidateAndExecute(instrumentedExecutionInput, graphQLSchema, instrumentationState, engineRunningState);
435+
//
436+
// finish up instrumentation
437+
executionResult = executionResult.whenComplete(completeInstrumentationCtxCF(executionInstrumentation));
438+
//
439+
// allow instrumentation to tweak the result
440+
executionResult = engineRunningState.compose(executionResult, (result -> instrumentation.instrumentExecutionResult(result, instrumentationParameters, instrumentationState)));
441+
return executionResult;
442+
} catch (AbortExecutionException abortException) {
443+
return handleAbortException(executionInput, instrumentationState, abortException);
444+
}
445+
}));
440446
});
441447
}
442448

449+
443450
private CompletableFuture<ExecutionResult> handleAbortException(ExecutionInput executionInput, InstrumentationState instrumentationState, AbortExecutionException abortException) {
444-
CompletableFuture<ExecutionResult> executionResult = CompletableFuture.completedFuture(abortException.toExecutionResult());
445451
InstrumentationExecutionParameters instrumentationParameters = new InstrumentationExecutionParameters(executionInput, this.graphQLSchema);
446-
//
447-
// allow instrumentation to tweak the result
448-
executionResult = executionResult.thenCompose(result -> instrumentation.instrumentExecutionResult(result, instrumentationParameters, instrumentationState));
449-
return executionResult;
452+
return instrumentation.instrumentExecutionResult(abortException.toExecutionResult(), instrumentationParameters, instrumentationState);
450453
}
451454

452455
private ExecutionInput ensureInputHasId(ExecutionInput executionInput) {
@@ -460,24 +463,24 @@ private ExecutionInput ensureInputHasId(ExecutionInput executionInput) {
460463
}
461464

462465

463-
private CompletableFuture<ExecutionResult> parseValidateAndExecute(ExecutionInput executionInput, GraphQLSchema graphQLSchema, InstrumentationState instrumentationState) {
466+
private CompletableFuture<ExecutionResult> parseValidateAndExecute(ExecutionInput executionInput, GraphQLSchema graphQLSchema, InstrumentationState instrumentationState, EngineRunningState engineRunningState) {
464467
AtomicReference<ExecutionInput> executionInputRef = new AtomicReference<>(executionInput);
465468
Function<ExecutionInput, PreparsedDocumentEntry> computeFunction = transformedInput -> {
466469
// if they change the original query in the pre-parser, then we want to see it downstream from then on
467470
executionInputRef.set(transformedInput);
468471
return parseAndValidate(executionInputRef, graphQLSchema, instrumentationState);
469472
};
470473
CompletableFuture<PreparsedDocumentEntry> preparsedDoc = preparsedDocumentProvider.getDocumentAsync(executionInput, computeFunction);
471-
return preparsedDoc.thenCompose(preparsedDocumentEntry -> {
474+
return engineRunningState.compose(preparsedDoc, (preparsedDocumentEntry -> {
472475
if (preparsedDocumentEntry.hasErrors()) {
473476
return CompletableFuture.completedFuture(new ExecutionResultImpl(preparsedDocumentEntry.getErrors()));
474477
}
475478
try {
476-
return execute(executionInputRef.get(), preparsedDocumentEntry.getDocument(), graphQLSchema, instrumentationState);
479+
return execute(executionInputRef.get(), preparsedDocumentEntry.getDocument(), graphQLSchema, instrumentationState, engineRunningState);
477480
} catch (AbortExecutionException e) {
478481
return CompletableFuture.completedFuture(e.toExecutionResult());
479482
}
480-
});
483+
}));
481484
}
482485

483486
private PreparsedDocumentEntry parseAndValidate(AtomicReference<ExecutionInput> executionInputRef, GraphQLSchema graphQLSchema, InstrumentationState instrumentationState) {
@@ -536,13 +539,14 @@ private List<ValidationError> validate(ExecutionInput executionInput, Document d
536539
private CompletableFuture<ExecutionResult> execute(ExecutionInput executionInput,
537540
Document document,
538541
GraphQLSchema graphQLSchema,
539-
InstrumentationState instrumentationState
542+
InstrumentationState instrumentationState,
543+
EngineRunningState engineRunningState
540544
) {
541545

542546
Execution execution = new Execution(queryStrategy, mutationStrategy, subscriptionStrategy, instrumentation, valueUnboxer, doNotAutomaticallyDispatchDataLoader);
543547
ExecutionId executionId = executionInput.getExecutionId();
544548

545-
return execution.execute(document, graphQLSchema, executionId, executionInput, instrumentationState);
549+
return execution.execute(document, graphQLSchema, executionId, executionInput, instrumentationState, engineRunningState);
546550
}
547551

548552
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ public AbstractAsyncExecutionStrategy(DataFetcherExceptionHandler dataFetcherExc
2222
}
2323

2424
protected BiConsumer<List<Object>, Throwable> handleResults(ExecutionContext executionContext, List<String> fieldNames, CompletableFuture<ExecutionResult> overallResult) {
25-
return (List<Object> results, Throwable exception) -> executionContext.run(exception, () -> {
25+
return (List<Object> results, Throwable exception) -> {
2626
if (exception != null) {
2727
handleNonNullException(executionContext, overallResult, exception);
2828
return;
@@ -35,6 +35,6 @@ protected BiConsumer<List<Object>, Throwable> handleResults(ExecutionContext exe
3535
resolvedValuesByField.put(fieldName, result);
3636
}
3737
overallResult.complete(new ExecutionResultImpl(resolvedValuesByField, executionContext.getErrors()));
38-
});
38+
};
3939
}
4040
}

0 commit comments

Comments
 (0)