graphql-java · andimarek · Mar 30, 2025 · Mar 30, 2025 · Mar 30, 2025
diff --git a/src/main/java/graphql/execution/AbstractAsyncExecutionStrategy.java b/src/main/java/graphql/execution/AbstractAsyncExecutionStrategy.java
@@ -5,7 +5,6 @@
 import graphql.ExecutionResultImpl;
 import graphql.PublicSpi;
 
-import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
@@ -24,8 +23,10 @@ public AbstractAsyncExecutionStrategy(DataFetcherExceptionHandler dataFetcherExc
 
     protected BiConsumer<List<Object>, Throwable> handleResults(ExecutionContext executionContext, List<String> fieldNames, CompletableFuture<ExecutionResult> overallResult) {
         return (List<Object> results, Throwable exception) -> {
+            executionContext.running();
             if (exception != null) {
                 handleNonNullException(executionContext, overallResult, exception);
+                executionContext.finished();
                 return;
             }
             Map<String, Object> resolvedValuesByField = Maps.newLinkedHashMapWithExpectedSize(fieldNames.size());
@@ -35,6 +36,7 @@ protected BiConsumer<List<Object>, Throwable> handleResults(ExecutionContext exe
                 resolvedValuesByField.put(fieldName, result);
             }
             overallResult.complete(new ExecutionResultImpl(resolvedValuesByField, executionContext.getErrors()));
+            executionContext.finished();
         };
     }
 }
diff --git a/src/main/java/graphql/execution/AsyncExecutionStrategy.java b/src/main/java/graphql/execution/AsyncExecutionStrategy.java
@@ -38,6 +38,7 @@ public AsyncExecutionStrategy(DataFetcherExceptionHandler exceptionHandler) {
     @Override
     @SuppressWarnings("FutureReturnValueIgnored")
     public CompletableFuture<ExecutionResult> execute(ExecutionContext executionContext, ExecutionStrategyParameters parameters) throws NonNullableFieldWasNullException {
+        executionContext.running();
         DataLoaderDispatchStrategy dataLoaderDispatcherStrategy = executionContext.getDataLoaderDispatcherStrategy();
         dataLoaderDispatcherStrategy.executionStrategy(executionContext, parameters);
         Instrumentation instrumentation = executionContext.getInstrumentation();
@@ -50,6 +51,7 @@ public CompletableFuture<ExecutionResult> execute(ExecutionContext executionCont
 
         Optional<ExecutionResult> isNotSensible = Introspection.isIntrospectionSensible(fields, executionContext);
         if (isNotSensible.isPresent()) {
+            executionContext.finished();
             return CompletableFuture.completedFuture(isNotSensible.get());
         }
 
@@ -60,11 +62,13 @@ public CompletableFuture<ExecutionResult> execute(ExecutionContext executionCont
         executionStrategyCtx.onDispatched();
 
         futures.await().whenComplete((completeValueInfos, throwable) -> {
+            executionContext.running();
             List<String> fieldsExecutedOnInitialResult = deferredExecutionSupport.getNonDeferredFieldNames(fieldNames);
 
             BiConsumer<List<Object>, Throwable> handleResultsConsumer = handleResults(executionContext, fieldsExecutedOnInitialResult, overallResult);
             if (throwable != null) {
                 handleResultsConsumer.accept(null, throwable.getCause());
+                executionContext.finished();
                 return;
             }
 
@@ -75,17 +79,21 @@ public CompletableFuture<ExecutionResult> execute(ExecutionContext executionCont
             dataLoaderDispatcherStrategy.executionStrategyOnFieldValuesInfo(completeValueInfos);
             executionStrategyCtx.onFieldValuesInfo(completeValueInfos);
             fieldValuesFutures.await().whenComplete(handleResultsConsumer);
+            executionContext.finished();
         }).exceptionally((ex) -> {
+            executionContext.running();
             // if there are any issues with combining/handling the field results,
             // complete the future at all costs and bubble up any thrown exception so
             // the execution does not hang.
             dataLoaderDispatcherStrategy.executionStrategyOnFieldValuesException(ex);
             executionStrategyCtx.onFieldValuesException();
             overallResult.completeExceptionally(ex);
+            executionContext.finished();
             return null;
         });
 
         overallResult.whenComplete(executionStrategyCtx::onCompleted);
+        executionContext.finished();
         return overallResult;
     }
 }
diff --git a/src/main/java/graphql/execution/DataLoaderDispatchStrategy.java b/src/main/java/graphql/execution/DataLoaderDispatchStrategy.java
@@ -2,8 +2,10 @@
 
 import graphql.Internal;
 import graphql.schema.DataFetcher;
+import graphql.schema.DataFetchingEnvironment;
 
 import java.util.List;
+import java.util.function.Supplier;
 
 @Internal
 public interface DataLoaderDispatchStrategy {
@@ -44,7 +46,7 @@ default void executeObjectOnFieldValuesException(Throwable t, ExecutionStrategyP
     default void fieldFetched(ExecutionContext executionContext,
                               ExecutionStrategyParameters executionStrategyParameters,
                               DataFetcher<?> dataFetcher,
-                              Object fetchedValue) {
+                              Object fetchedValue, Supplier<DataFetchingEnvironment> dataFetchingEnvironment) {
 
     }
 

diff --git a/src/main/java/graphql/execution/Execution.java b/src/main/java/graphql/execution/Execution.java
@@ -29,8 +29,8 @@
 import graphql.schema.GraphQLObjectType;
 import graphql.schema.GraphQLSchema;
 import graphql.schema.impl.SchemaUtil;
-import org.jspecify.annotations.NonNull;
 import graphql.util.FpKit;
+import org.jspecify.annotations.NonNull;
 import org.reactivestreams.Publisher;
 
 import java.util.Collections;
@@ -58,6 +58,8 @@ public class Execution {
     private final ValueUnboxer valueUnboxer;
     private final boolean doNotAutomaticallyDispatchDataLoader;
 
+    public static final String EXECUTION_CONTEXT_KEY = "__GraphQL_Java_ExecutionContext";
+
     public Execution(ExecutionStrategy queryStrategy,
                      ExecutionStrategy mutationStrategy,
                      ExecutionStrategy subscriptionStrategy,
@@ -114,6 +116,8 @@ public CompletableFuture<ExecutionResult> execute(Document document, GraphQLSche
                 .build();
 
         executionContext.getGraphQLContext().put(ResultNodesInfo.RESULT_NODES_INFO, executionContext.getResultNodesInfo());
+        executionContext.getGraphQLContext().put(EXECUTION_CONTEXT_KEY, executionContext);
+
 
         InstrumentationExecutionParameters parameters = new InstrumentationExecutionParameters(
                 executionInput, graphQLSchema

diff --git a/src/main/java/graphql/execution/ExecutionContext.java b/src/main/java/graphql/execution/ExecutionContext.java
@@ -28,6 +28,7 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.function.Supplier;
@@ -63,6 +64,8 @@ public class ExecutionContext {
     private final Supplier<ExecutableNormalizedOperation> queryTree;
     private final boolean propagateErrorsOnNonNullContractFailure;
 
+    private final AtomicInteger isRunning = new AtomicInteger(0);
+
     // this is modified after creation so it needs to be volatile to ensure visibility across Threads
     private volatile DataLoaderDispatchStrategy dataLoaderDispatcherStrategy = DataLoaderDispatchStrategy.NO_OP;
 
@@ -349,4 +352,16 @@ public ExecutionContext transform(Consumer<ExecutionContextBuilder> builderConsu
     public ResultNodesInfo getResultNodesInfo() {
         return resultNodesInfo;
     }
+
+    public boolean isRunning() {
+        return isRunning.get() > 0;
+    }
+
+    public void running() {
+        isRunning.incrementAndGet();
+    }
+
+    public void finished() {
+        isRunning.decrementAndGet();
+    }
 }
diff --git a/src/main/java/graphql/execution/ExecutionStrategy.java b/src/main/java/graphql/execution/ExecutionStrategy.java
@@ -201,6 +201,7 @@ public static String mkNameForPath(List<Field> currentField) {
     @SuppressWarnings("unchecked")
     @DuckTyped(shape = "CompletableFuture<Map<String, Object>> | Map<String, Object>")
     protected Object executeObject(ExecutionContext executionContext, ExecutionStrategyParameters parameters) throws NonNullableFieldWasNullException {
+        executionContext.running();
         DataLoaderDispatchStrategy dataLoaderDispatcherStrategy = executionContext.getDataLoaderDispatcherStrategy();
         dataLoaderDispatcherStrategy.executeObject(executionContext, parameters);
         Instrumentation instrumentation = executionContext.getInstrumentation();
@@ -225,15 +226,18 @@ protected Object executeObject(ExecutionContext executionContext, ExecutionStrat
         if (fieldValueInfosResult instanceof CompletableFuture) {
             CompletableFuture<List<FieldValueInfo>> fieldValueInfos = (CompletableFuture<List<FieldValueInfo>>) fieldValueInfosResult;
             fieldValueInfos.whenComplete((completeValueInfos, throwable) -> {
+                executionContext.running();
                 if (throwable != null) {
                     handleResultsConsumer.accept(null, throwable);
+                    executionContext.finished();
                     return;
                 }
 
                 Async.CombinedBuilder<Object> resultFutures = fieldValuesCombinedBuilder(completeValueInfos);
                 dataLoaderDispatcherStrategy.executeObjectOnFieldValuesInfo(completeValueInfos, parameters);
                 resolveObjectCtx.onFieldValuesInfo(completeValueInfos);
                 resultFutures.await().whenComplete(handleResultsConsumer);
+                executionContext.finished();
             }).exceptionally((ex) -> {
                 // if there are any issues with combining/handling the field results,
                 // complete the future at all costs and bubble up any thrown exception so
@@ -244,6 +248,7 @@ protected Object executeObject(ExecutionContext executionContext, ExecutionStrat
                 return null;
             });
             overallResult.whenComplete(resolveObjectCtx::onCompleted);
+            executionContext.finished();
             return overallResult;
         } else {
             List<FieldValueInfo> completeValueInfos = (List<FieldValueInfo>) fieldValueInfosResult;
@@ -257,10 +262,12 @@ protected Object executeObject(ExecutionContext executionContext, ExecutionStrat
                 CompletableFuture<List<Object>> completedValues = (CompletableFuture<List<Object>>) completedValuesObject;
                 completedValues.whenComplete(handleResultsConsumer);
                 overallResult.whenComplete(resolveObjectCtx::onCompleted);
+                executionContext.finished();
                 return overallResult;
             } else {
                 Map<String, Object> fieldValueMap = buildFieldValueMap(fieldsExecutedOnInitialResult, (List<Object>) completedValuesObject);
                 resolveObjectCtx.onCompleted(fieldValueMap, null);
+                executionContext.finished();
                 return fieldValueMap;
             }
         }
@@ -276,12 +283,15 @@ protected Object executeObject(ExecutionContext executionContext, ExecutionStrat
 
     private BiConsumer<List<Object>, Throwable> buildFieldValueMap(List<String> fieldNames, CompletableFuture<Map<String, Object>> overallResult, ExecutionContext executionContext) {
         return (List<Object> results, Throwable exception) -> {
+            executionContext.running();
             if (exception != null) {
                 handleValueException(overallResult, exception, executionContext);
+                executionContext.finished();
                 return;
             }
             Map<String, Object> resolvedValuesByField = buildFieldValueMap(fieldNames, results);
             overallResult.complete(resolvedValuesByField);
+            executionContext.finished();
         };
     }
 
@@ -496,7 +506,7 @@ private Object fetchField(GraphQLFieldDefinition fieldDef, ExecutionContext exec
         dataFetcher = instrumentation.instrumentDataFetcher(dataFetcher, instrumentationFieldFetchParams, executionContext.getInstrumentationState());
         dataFetcher = executionContext.getDataLoaderDispatcherStrategy().modifyDataFetcher(dataFetcher);
         Object fetchedObject = invokeDataFetcher(executionContext, parameters, fieldDef, dataFetchingEnvironment, dataFetcher);
-        executionContext.getDataLoaderDispatcherStrategy().fieldFetched(executionContext, parameters, dataFetcher, fetchedObject);
+        executionContext.getDataLoaderDispatcherStrategy().fieldFetched(executionContext, parameters, dataFetcher, fetchedObject, dataFetchingEnvironment);
         fetchCtx.onDispatched();
         fetchCtx.onFetchedValue(fetchedObject);
         // if it's a subscription, leave any reactive objects alone
@@ -509,11 +519,15 @@ private Object fetchField(GraphQLFieldDefinition fieldDef, ExecutionContext exec
             CompletableFuture<Object> fetchedValue = (CompletableFuture<Object>) fetchedObject;
             return fetchedValue
                     .handle((result, exception) -> {
+                        executionContext.running();
                         fetchCtx.onCompleted(result, exception);
                         if (exception != null) {
-                            return handleFetchingException(dataFetchingEnvironment.get(), parameters, exception);
+                            CompletableFuture<Object> handleFetchingExceptionResult = handleFetchingException(dataFetchingEnvironment.get(), parameters, exception);
+                            executionContext.finished();
+                            return handleFetchingExceptionResult;
                         } else {
                             // we can simply return the fetched value CF and avoid a allocation
+                            executionContext.finished();
                             return fetchedValue;
                         }
                     })
@@ -553,7 +567,7 @@ protected Supplier<ExecutableNormalizedField> getNormalizedField(ExecutionContex
     protected FetchedValue unboxPossibleDataFetcherResult(ExecutionContext executionContext,
                                                           ExecutionStrategyParameters parameters,
                                                           Object result) {
-
+        executionContext.running();
         if (result instanceof DataFetcherResult) {
             DataFetcherResult<?> dataFetcherResult = (DataFetcherResult<?>) result;
 
@@ -567,9 +581,11 @@ protected FetchedValue unboxPossibleDataFetcherResult(ExecutionContext execution
                 localContext = parameters.getLocalContext();
             }
             Object unBoxedValue = executionContext.getValueUnboxer().unbox(dataFetcherResult.getData());
+            executionContext.finished();
             return new FetchedValue(unBoxedValue, dataFetcherResult.getErrors(), localContext);
         } else {
             Object unBoxedValue = executionContext.getValueUnboxer().unbox(result);
+            executionContext.finished();
             return new FetchedValue(unBoxedValue, ImmutableList.of(), parameters.getLocalContext());
         }
     }
@@ -638,6 +654,7 @@ protected FieldValueInfo completeField(ExecutionContext executionContext, Execut
     }
 
     private FieldValueInfo completeField(GraphQLFieldDefinition fieldDef, ExecutionContext executionContext, ExecutionStrategyParameters parameters, FetchedValue fetchedValue) {
+        executionContext.running();
         GraphQLObjectType parentType = (GraphQLObjectType) parameters.getExecutionStepInfo().getUnwrappedNonNullType();
         ExecutionStepInfo executionStepInfo = createExecutionStepInfo(executionContext, parameters, fieldDef, parentType);
 
@@ -661,6 +678,7 @@ private FieldValueInfo completeField(GraphQLFieldDefinition fieldDef, ExecutionC
         CompletableFuture<Object> executionResultFuture = fieldValueInfo.getFieldValueFuture();
         ctxCompleteField.onDispatched();
         executionResultFuture.whenComplete(ctxCompleteField::onCompleted);
+        executionContext.finished();
         return fieldValueInfo;
     }
 
@@ -833,13 +851,16 @@ protected FieldValueInfo completeValueForList(ExecutionContext executionContext,
             overallResult.whenComplete(completeListCtx::onCompleted);
 
             resultsFuture.whenComplete((results, exception) -> {
+                executionContext.running();
                 if (exception != null) {
+                    executionContext.finished();
                     handleValueException(overallResult, exception, executionContext);
                     return;
                 }
                 List<Object> completedResults = new ArrayList<>(results.size());
                 completedResults.addAll(results);
                 overallResult.complete(completedResults);
+                executionContext.finished();
             });
             listOrPromiseToList = overallResult;
         } else {