wip

andimarek · Sep 28, 2025 · Sep 27, 2025 · Sep 27, 2025 · Sep 27, 2025 · Sep 27, 2025
commit 2cb5354da49cf94f8ff959aeca9235c8bf0b004d
diff --git a/src/main/java/graphql/execution/AsyncExecutionStrategy.java b/src/main/java/graphql/execution/AsyncExecutionStrategy.java
@@ -55,8 +55,9 @@ public CompletableFuture<ExecutionResult> execute(ExecutionContext executionCont
         DeferredExecutionSupport deferredExecutionSupport = createDeferredExecutionSupport(executionContext, parameters);
 
         dataLoaderDispatcherStrategy.executionStrategy(executionContext, parameters, deferredExecutionSupport.getNonDeferredFieldNames(fieldNames).size());
-
         Async.CombinedBuilder<FieldValueInfo> futures = getAsyncFieldValueInfo(executionContext, parameters, deferredExecutionSupport);
+        dataLoaderDispatcherStrategy.finishedFetching(executionContext, parameters);
+
 
         CompletableFuture<ExecutionResult> overallResult = new CompletableFuture<>();
         executionStrategyCtx.onDispatched();

diff --git a/src/main/java/graphql/execution/AsyncSerialExecutionStrategy.java b/src/main/java/graphql/execution/AsyncSerialExecutionStrategy.java
@@ -71,6 +71,7 @@ private Object resolveSerialField(ExecutionContext executionContext,
         dataLoaderDispatcherStrategy.executionSerialStrategy(executionContext, newParameters);
 
         Object fieldWithInfo = resolveFieldWithInfo(executionContext, newParameters);
+        dataLoaderDispatcherStrategy.finishedFetching(executionContext, newParameters);
         if (fieldWithInfo instanceof CompletableFuture) {
             //noinspection unchecked
             return ((CompletableFuture<FieldValueInfo>) fieldWithInfo).thenCompose(fvi -> {

diff --git a/src/main/java/graphql/execution/DataLoaderDispatchStrategy.java b/src/main/java/graphql/execution/DataLoaderDispatchStrategy.java
@@ -64,4 +64,8 @@ default void newSubscriptionExecution(AlternativeCallContext alternativeCallCont
     default void subscriptionEventCompletionDone(AlternativeCallContext alternativeCallContext) {
 
     }
+
+    default void finishedFetching(ExecutionContext executionContext, ExecutionStrategyParameters newParameters) {
+
+    }
 }
diff --git a/src/main/java/graphql/execution/Execution.java b/src/main/java/graphql/execution/Execution.java
@@ -15,7 +15,7 @@
 import graphql.execution.instrumentation.Instrumentation;
 import graphql.execution.instrumentation.InstrumentationContext;
 import graphql.execution.instrumentation.InstrumentationState;
-import graphql.execution.instrumentation.dataloader.PerLevelDataLoaderDispatchStrategy;
+import graphql.execution.instrumentation.dataloader.ExhaustedDataLoaderDispatchStrategy;
 import graphql.execution.instrumentation.parameters.InstrumentationExecuteOperationParameters;
 import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters;
 import graphql.extensions.ExtensionsBuilder;
@@ -262,7 +262,8 @@ private DataLoaderDispatchStrategy createDataLoaderDispatchStrategy(ExecutionCon
         if (executionContext.getDataLoaderRegistry() == EMPTY_DATALOADER_REGISTRY || doNotAutomaticallyDispatchDataLoader) {
             return DataLoaderDispatchStrategy.NO_OP;
         }
-        return new PerLevelDataLoaderDispatchStrategy(executionContext);
+//        return new PerLevelDataLoaderDispatchStrategy(executionContext);
+        return new ExhaustedDataLoaderDispatchStrategy(executionContext);
     }
 
 

diff --git a/src/main/java/graphql/execution/ExecutionStrategy.java b/src/main/java/graphql/execution/ExecutionStrategy.java
@@ -210,6 +210,7 @@ protected Object executeObject(ExecutionContext executionContext, ExecutionStrat
         List<String> fieldsExecutedOnInitialResult = deferredExecutionSupport.getNonDeferredFieldNames(fieldNames);
         dataLoaderDispatcherStrategy.executeObject(executionContext, parameters, fieldsExecutedOnInitialResult.size());
         Async.CombinedBuilder<FieldValueInfo> resolvedFieldFutures = getAsyncFieldValueInfo(executionContext, parameters, deferredExecutionSupport);
+        dataLoaderDispatcherStrategy.finishedFetching(executionContext, parameters);
 
         CompletableFuture<Map<String, Object>> overallResult = new CompletableFuture<>();
         BiConsumer<List<Object>, Throwable> handleResultsConsumer = buildFieldValueMap(fieldsExecutedOnInitialResult, overallResult, executionContext);

diff --git a/...ava/graphql/execution/instrumentation/dataloader/ExhaustedDataLoaderDispatchStrategy.java b/...ava/graphql/execution/instrumentation/dataloader/ExhaustedDataLoaderDispatchStrategy.java
@@ -0,0 +1,271 @@
+package graphql.execution.instrumentation.dataloader;
+
+import graphql.Assert;
+import graphql.Internal;
+import graphql.Profiler;
+import graphql.execution.DataLoaderDispatchStrategy;
+import graphql.execution.ExecutionContext;
+import graphql.execution.ExecutionStrategyParameters;
+import graphql.execution.FieldValueInfo;
+import graphql.execution.incremental.AlternativeCallContext;
+import org.dataloader.DataLoader;
+import org.dataloader.DataLoaderRegistry;
+import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
+
+@Internal
+@NullMarked
+public class ExhaustedDataLoaderDispatchStrategy implements DataLoaderDispatchStrategy {
+
+    private final CallStack initialCallStack;
+    private final ExecutionContext executionContext;
+
+    private final Profiler profiler;
+
+    private final Map<AlternativeCallContext, CallStack> alternativeCallContextMap = new ConcurrentHashMap<>();
+
+
+    private static class CallStack {
+
+
+        static class State {
+            final int objectRunningCount;
+            final boolean dataLoaderToDispatch;
+            final boolean currentlyDispatching;
+
+            State(int objectRunningCount, boolean dataLoaderToDispatch, boolean currentlyDispatching) {
+                this.objectRunningCount = objectRunningCount;
+                this.dataLoaderToDispatch = dataLoaderToDispatch;
+                this.currentlyDispatching = currentlyDispatching;
+            }
+
+            public State copy() {
+                return new State(objectRunningCount, dataLoaderToDispatch, currentlyDispatching);
+            }
+
+            public State incrementObjectRunningCount() {
+                return new State(objectRunningCount + 1, dataLoaderToDispatch, currentlyDispatching);
+            }
+
+            public State decrementObjetRunningCount() {
+                return new State(objectRunningCount - 1, dataLoaderToDispatch, currentlyDispatching);
+            }
+
+            public State dataLoaderToDispatch() {
+                return new State(objectRunningCount, true, currentlyDispatching);
+            }
+
+            public State startDispatching() {
+                return new State(objectRunningCount, false, true);
+            }
+
+            public State stopDispatching() {
+                return new State(objectRunningCount, false, false);
+            }
+
+
+            @Override
+            public String toString() {
+                return "State{" +
+                       "objectRunningCount=" + objectRunningCount +
+                       ", dataLoaderToDispatch=" + dataLoaderToDispatch +
+                       '}';
+            }
+        }
+
+        private final AtomicLong state = new AtomicLong();
+        private final AtomicReference<State> stateRef = new AtomicReference<>(new State(0, false, false));
+
+        public State getState() {
+            return Assert.assertNotNull(stateRef.get());
+        }
+
+        public boolean tryUpdateState(State oldState, State newState) {
+            System.out.println("updateState: " + oldState + " -> " + newState);
+            return stateRef.compareAndSet(oldState, newState);
+        }
+
+        private final AtomicInteger deferredFragmentRootFieldsCompleted = new AtomicInteger();
+
+        public CallStack() {
+        }
+
+
+        public void clear() {
+            deferredFragmentRootFieldsCompleted.set(0);
+            stateRef.set(new State(0, false, false));
+        }
+    }
+
+    public ExhaustedDataLoaderDispatchStrategy(ExecutionContext executionContext) {
+        this.initialCallStack = new CallStack();
+        this.executionContext = executionContext;
+
+        this.profiler = executionContext.getProfiler();
+    }
+
+
+    @Override
+    public void executionStrategy(ExecutionContext executionContext, ExecutionStrategyParameters parameters, int fieldCount) {
+        Assert.assertTrue(parameters.getExecutionStepInfo().getPath().isRootPath());
+        // no concurrency access happening
+        CallStack.State state = initialCallStack.getState();
+        Assert.assertTrue(initialCallStack.tryUpdateState(state, state.incrementObjectRunningCount()));
+    }
+
+    @Override
+    public void finishedFetching(ExecutionContext executionContext, ExecutionStrategyParameters newParameters) {
+        CallStack callStack = getCallStack(newParameters);
+        decrementObjectRunningAndMaybeDispatch(callStack);
+    }
+
+    @Override
+    public void executionSerialStrategy(ExecutionContext executionContext, ExecutionStrategyParameters parameters) {
+        CallStack callStack = getCallStack(parameters);
+        callStack.clear();
+        CallStack.State state = callStack.getState();
+        // no concurrency access happening
+        Assert.assertTrue(callStack.tryUpdateState(state, state.incrementObjectRunningCount()));
+    }
+
+
+    @Override
+    public void executeObject(ExecutionContext executionContext, ExecutionStrategyParameters parameters, int fieldCount) {
+        CallStack callStack = getCallStack(parameters);
+        while (true) {
+            CallStack.State state = callStack.getState();
+            if (callStack.tryUpdateState(state, state.incrementObjectRunningCount())) {
+                break;
+            }
+        }
+    }
+
+
+    @Override
+    public void newSubscriptionExecution(AlternativeCallContext alternativeCallContext) {
+        CallStack callStack = new CallStack();
+        alternativeCallContextMap.put(alternativeCallContext, callStack);
+    }
+
+    @Override
+    public void deferredOnFieldValue(String resultKey, FieldValueInfo fieldValueInfo, Throwable throwable, ExecutionStrategyParameters parameters) {
+        CallStack callStack = getCallStack(parameters);
+        int deferredFragmentRootFieldsCompleted = callStack.deferredFragmentRootFieldsCompleted.incrementAndGet();
+        Assert.assertNotNull(parameters.getDeferredCallContext());
+        if (deferredFragmentRootFieldsCompleted == parameters.getDeferredCallContext().getFields()) {
+            decrementObjectRunningAndMaybeDispatch(callStack);
+        }
+
+    }
+
+    private CallStack getCallStack(ExecutionStrategyParameters parameters) {
+        return getCallStack(parameters.getDeferredCallContext());
+    }
+
+    private CallStack getCallStack(@Nullable AlternativeCallContext alternativeCallContext) {
+        if (alternativeCallContext == null) {
+            return this.initialCallStack;
+        } else {
+            return alternativeCallContextMap.computeIfAbsent(alternativeCallContext, k -> {
+                /*
+                  This is only for handling deferred cases. Subscription cases will also get a new callStack, but
+                  it is explicitly created in `newSubscriptionExecution`.
+                  The reason we are doing this lazily is, because we don't have explicit startDeferred callback.
+                 */
+                CallStack callStack = new CallStack();
+                return callStack;
+            });
+        }
+    }
+
+
+    private void decrementObjectRunningAndMaybeDispatch(CallStack callStack) {
+        CallStack.State oldState;
+        CallStack.State newState;
+        while (true) {
+            oldState = callStack.getState();
+            newState = oldState.decrementObjetRunningCount();
+            if (callStack.tryUpdateState(oldState, newState)) {
+                break;
+            }
+        }
+        // this means we have not fetching running and we can execute
+        if (newState.objectRunningCount == 0 && !newState.currentlyDispatching) {
+            dispatchImpl(callStack);
+        }
+    }
+
+    private void newDataLoaderInvocationMaybeDispatch(CallStack callStack) {
+        CallStack.State oldState;
+        CallStack.State newState;
+        while (true) {
+            oldState = callStack.getState();
+            newState = oldState.dataLoaderToDispatch();
+            if (callStack.tryUpdateState(oldState, newState)) {
+                break;
+            }
+        }
+//        System.out.println("new data loader invocation maybe with state: " + newState);
+        // this means we are not waiting for some fetching to be finished and we need to dispatch
+        if (newState.objectRunningCount == 0 && !newState.currentlyDispatching) {
+            dispatchImpl(callStack);
+        }
+
+    }
+
+
+    private void dispatchImpl(CallStack callStack) {
+
+        CallStack.State oldState;
+        while (true) {
+            oldState = callStack.getState();
+            if (!oldState.dataLoaderToDispatch) {
+                CallStack.State newState = oldState.stopDispatching();
+                if (callStack.tryUpdateState(oldState, newState)) {
+                    return;
+                }
+            }
+            CallStack.State newState = oldState.startDispatching();
+            if (callStack.tryUpdateState(oldState, newState)) {
+                break;
+            }
+        }
+
+        DataLoaderRegistry dataLoaderRegistry = executionContext.getDataLoaderRegistry();
+        List<DataLoader<?, ?>> dataLoaders = dataLoaderRegistry.getDataLoaders();
+        List<CompletableFuture<? extends List<?>>> allDispatchedCFs = new ArrayList<>();
+        for (DataLoader<?, ?> dataLoader : dataLoaders) {
+            CompletableFuture<? extends List<?>> dispatch = dataLoader.dispatch();
+            allDispatchedCFs.add(dispatch);
+        }
+        CompletableFuture.allOf(allDispatchedCFs.toArray(new CompletableFuture[0]))
+                .whenComplete((unused, throwable) -> {
+                    dispatchImpl(callStack);
+                });
+
+    }
+
+
+    public void newDataLoaderInvocation(String resultPath,
+                                        int level,
+                                        DataLoader dataLoader,
+                                        String dataLoaderName,
+                                        Object key,
+                                        @Nullable AlternativeCallContext alternativeCallContext) {
+//        DataLoaderInvocation dataLoaderInvocation = new DataLoaderInvocation(resultPath, level, dataLoader, dataLoaderName, key);
+        CallStack callStack = getCallStack(alternativeCallContext);
+        newDataLoaderInvocationMaybeDispatch(callStack);
+    }
+
+
+}
+
diff --git a/...java/graphql/execution/instrumentation/dataloader/PerLevelDataLoaderDispatchStrategy.java b/...java/graphql/execution/instrumentation/dataloader/PerLevelDataLoaderDispatchStrategy.java
@@ -391,8 +391,7 @@ public void subscriptionEventCompletionDone(AlternativeCallContext alternativeCa
     }
 
     @Override
-    public void deferredOnFieldValue(String resultKey, FieldValueInfo fieldValueInfo, Throwable
-            throwable, ExecutionStrategyParameters parameters) {
+    public void deferredOnFieldValue(String resultKey, FieldValueInfo fieldValueInfo, Throwable throwable, ExecutionStrategyParameters parameters) {
         CallStack callStack = getCallStack(parameters);
         int deferredFragmentRootFieldsCompleted = callStack.deferredFragmentRootFieldsCompleted.incrementAndGet();
         Assert.assertNotNull(parameters.getDeferredCallContext());

diff --git a/src/main/java/graphql/schema/DataFetchingEnvironmentImpl.java b/src/main/java/graphql/schema/DataFetchingEnvironmentImpl.java
@@ -15,7 +15,6 @@
 import graphql.execution.MergedField;
 import graphql.execution.directives.QueryDirectives;
 import graphql.execution.incremental.AlternativeCallContext;
-import graphql.execution.instrumentation.dataloader.DataLoaderDispatchingContextKeys;
 import graphql.language.Document;
 import graphql.language.Field;
 import graphql.language.FragmentDefinition;
@@ -233,9 +232,9 @@ public ExecutionStepInfo getExecutionStepInfo() {
         if (dataLoader == null) {
             return null;
         }
-        if (!graphQLContext.getBoolean(DataLoaderDispatchingContextKeys.ENABLE_DATA_LOADER_CHAINING, false)) {
-            return dataLoader;
-        }
+//        if (!graphQLContext.getBoolean(DataLoaderDispatchingContextKeys.ENABLE_DATA_LOADER_CHAINING, false)) {
+//            return dataLoader;
+//        }
         return new DataLoaderWithContext<>(this, dataLoaderName, dataLoader);
     }
 

diff --git a/src/main/java/graphql/schema/DataLoaderWithContext.java b/src/main/java/graphql/schema/DataLoaderWithContext.java
@@ -2,7 +2,7 @@
 
 import graphql.Internal;
 import graphql.execution.incremental.AlternativeCallContext;
-import graphql.execution.instrumentation.dataloader.PerLevelDataLoaderDispatchStrategy;
+import graphql.execution.instrumentation.dataloader.ExhaustedDataLoaderDispatchStrategy;
 import org.dataloader.DataLoader;
 import org.dataloader.DelegatingDataLoader;
 import org.jspecify.annotations.NonNull;
@@ -32,11 +32,11 @@ public CompletableFuture<V> load(@NonNull K key, @Nullable Object keyContext) {
         DataFetchingEnvironmentImpl dfeImpl = (DataFetchingEnvironmentImpl) dfe;
         DataFetchingEnvironmentImpl.DFEInternalState dfeInternalState = (DataFetchingEnvironmentImpl.DFEInternalState) dfeImpl.toInternal();
         dfeInternalState.getProfiler().dataLoaderUsed(dataLoaderName);
-        if (dfeInternalState.getDataLoaderDispatchStrategy() instanceof PerLevelDataLoaderDispatchStrategy) {
+        if (dfeInternalState.getDataLoaderDispatchStrategy() instanceof ExhaustedDataLoaderDispatchStrategy) {
             AlternativeCallContext alternativeCallContext = dfeInternalState.getDeferredCallContext();
             int level = dfe.getExecutionStepInfo().getPath().getLevel();
             String path = dfe.getExecutionStepInfo().getPath().toString();
-            ((PerLevelDataLoaderDispatchStrategy) dfeInternalState.dataLoaderDispatchStrategy).newDataLoaderInvocation(path, level, delegate, dataLoaderName, key, alternativeCallContext);
+            ((ExhaustedDataLoaderDispatchStrategy) dfeInternalState.dataLoaderDispatchStrategy).newDataLoaderInvocation(path, level, delegate, dataLoaderName, key, alternativeCallContext);
         }
         return result;
     }
-Original file line number
+Diff line change
@@ Expand Up @@
         default void subscriptionEventCompletionDone(AlternativeCallContext alternativeCallContext) {
         }
+        default void finishedFetching(ExecutionContext executionContext, ExecutionStrategyParameters newParameters) {
+        }
     }