andimarek · May 16, 2025 · Mar 30, 2025 · Mar 30, 2025 · Mar 30, 2025 · Mar 31, 2025
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/ExecutionStrategy.java b/src/main/java/graphql/execution/ExecutionStrategy.java
@@ -446,6 +446,7 @@ private Object fetchField(GraphQLFieldDefinition fieldDef, ExecutionContext exec
         }
 
         MergedField field = parameters.getField();
+        String pathString = parameters.getPath().toString();
         GraphQLObjectType parentType = (GraphQLObjectType) parameters.getExecutionStepInfo().getUnwrappedNonNullType();
 
         // if the DF (like PropertyDataFetcher) does not use the arguments or execution step info then dont build any
@@ -496,7 +497,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

diff --git a/src/main/java/graphql/execution/instrumentation/dataloader/DataLoaderCompletableFuture.java b/src/main/java/graphql/execution/instrumentation/dataloader/DataLoaderCompletableFuture.java
@@ -0,0 +1,114 @@
+package graphql.execution.instrumentation.dataloader;
+
+import graphql.Internal;
+import graphql.execution.DataLoaderDispatchStrategy;
+import graphql.execution.ExecutionContext;
+import graphql.schema.DataFetchingEnvironment;
+import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
+
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executor;
+import java.util.function.Supplier;
+
+import static graphql.execution.Execution.EXECUTION_CONTEXT_KEY;
+
+
+@Internal
+@NullMarked
+public class DataLoaderCompletableFuture<T> extends CompletableFuture<T> {
+    final DataFetchingEnvironment dfe;
+    final String dataLoaderName;
+    final Object key;
+    private final CompletableFuture<Object> underlyingDataLoaderCompletableFuture;
+
+    final CompletableFuture<Void> finishedSyncDependents = new CompletableFuture<>();
+
+    public DataLoaderCompletableFuture(DataFetchingEnvironment dfe, String dataLoaderName, Object key) {
+        this.dfe = dfe;
+        this.dataLoaderName = dataLoaderName;
+        this.key = key;
+        if (dataLoaderName != null) {
+            underlyingDataLoaderCompletableFuture = dfe.getDataLoaderRegistry().getDataLoader(dataLoaderName).load(key);
+            underlyingDataLoaderCompletableFuture.whenComplete((value, throwable) -> {
+                if (throwable != null) {
+                    completeExceptionally(throwable);
+                } else {
+                    complete((T) value); // causing all sync dependent code to run
+                }
+                // post completion hook
+                finishedSyncDependents.complete(null);
+            });
+        } else {
+            underlyingDataLoaderCompletableFuture = null;
+        }
+    }
+
+    DataLoaderCompletableFuture() {
+        this.dfe = null;
+        this.dataLoaderName = null;
+        this.key = null;
+        underlyingDataLoaderCompletableFuture = null;
+    }
+
+
+    @Override
+    public <U> CompletableFuture<U> newIncompleteFuture() {
+        return new DataLoaderCompletableFuture<>();
+    }
+
+    public static boolean isDataLoaderCompletableFuture(Object object) {
+        return object instanceof DataLoaderCompletableFuture;
+    }
+
+    public static <T> CompletableFuture<T> newDLCF(DataFetchingEnvironment dfe, String dataLoaderName, Object key) {
+        DataLoaderCompletableFuture<T> result = new DataLoaderCompletableFuture<>(dfe, dataLoaderName, key);
+        ExecutionContext executionContext = dfe.getGraphQlContext().get(EXECUTION_CONTEXT_KEY);
+        DataLoaderDispatchStrategy dataLoaderDispatcherStrategy = executionContext.getDataLoaderDispatcherStrategy();
+        if (dataLoaderDispatcherStrategy instanceof PerLevelDataLoaderDispatchStrategy) {
+            ((PerLevelDataLoaderDispatchStrategy) dataLoaderDispatcherStrategy).newDataLoaderCF(result);
+        }
+        return result;
+    }
+
+
+    public static <U> CompletableFuture<U> supplyDLCF(DataFetchingEnvironment env, Supplier<U> supplier, @Nullable Executor executor) {
+        DataLoaderCompletableFuture<U> d = new DataLoaderCompletableFuture<>(env, null, null);
+        if (executor == null) {
+            executor = d.defaultExecutor();
+        }
+        executor.execute(() -> {
+            d.complete(supplier.get());
+        });
+        return d;
+
+    }
+
+
+    public static <U> CompletableFuture<U> wrap(DataFetchingEnvironment env, CompletableFuture<U> completableFuture) {
+        if (completableFuture instanceof DataLoaderCompletableFuture) {
+            return completableFuture;
+        }
+        DataLoaderCompletableFuture<U> d = new DataLoaderCompletableFuture<>(env, null, null);
+        completableFuture.whenComplete((u, ex) -> {
+            if (ex != null) {
+                d.completeExceptionally(ex);
+            } else {
+                d.complete(u);
+            }
+        });
+        return d;
+    }
+
+    public static CompletableFuture<Void> waitUntilAllSyncDependentsComplete(List<DataLoaderCompletableFuture<?>> dataLoaderCompletableFutureList) {
+        CompletableFuture<?>[] finishedSyncArray = dataLoaderCompletableFutureList
+                .stream()
+                .map(dataLoaderCompletableFuture -> dataLoaderCompletableFuture.finishedSyncDependents)
+                .toArray(CompletableFuture[]::new);
+        return CompletableFuture.allOf(finishedSyncArray);
+    }
+
+
+}
+
diff --git a/...java/graphql/execution/instrumentation/dataloader/PerLevelDataLoaderDispatchStrategy.java b/...java/graphql/execution/instrumentation/dataloader/PerLevelDataLoaderDispatchStrategy.java
@@ -7,19 +7,35 @@
 import graphql.execution.ExecutionStrategyParameters;
 import graphql.execution.FieldValueInfo;
 import graphql.schema.DataFetcher;
+import graphql.schema.DataFetchingEnvironment;
 import graphql.util.LockKit;
 import org.dataloader.DataLoaderRegistry;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
+import java.util.ArrayList;
 import java.util.LinkedHashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
 
 @Internal
 public class PerLevelDataLoaderDispatchStrategy implements DataLoaderDispatchStrategy {
 
+    private static final Logger log = LoggerFactory.getLogger(PerLevelDataLoaderDispatchStrategy.class);
     private final CallStack callStack;
     private final ExecutionContext executionContext;
 
+    static final ScheduledExecutorService isolatedDLCFBatchWindowScheduler = Executors.newSingleThreadScheduledExecutor();
+    static final int BATCH_WINDOW_NANO_SECONDS = 500_000;
+
 
     private static class CallStack {
 
@@ -34,10 +50,27 @@ private static class CallStack {
 
         private final Set<Integer> dispatchedLevels = new LinkedHashSet<>();
 
+        // fields only relevant when a DataLoaderCF is involved
+        private final List<DataLoaderCompletableFuture<?>> allDataLoaderCompletableFuture = new CopyOnWriteArrayList<>();
+        //TODO: maybe this should be cleaned up once the CF returned by these fields are completed
+        // otherwise this will stick around until the whole request is finished
+        private final Set<DataFetchingEnvironment> fieldsFinishedDispatching = ConcurrentHashMap.newKeySet();
+        private final Map<Integer, Set<DataFetchingEnvironment>> levelToDFEWithDataLoaderCF = new ConcurrentHashMap<>();
+
+        private final Set<DataFetchingEnvironment> batchWindowOfIsolatedDfeToDispatch = ConcurrentHashMap.newKeySet();
+
+        private boolean batchWindowOpen = false;
+
+
         public CallStack() {
             expectedExecuteObjectCallsPerLevel.set(1, 1);
         }
 
+        public void addDataLoaderDFE(int level, DataFetchingEnvironment dfe) {
+            levelToDFEWithDataLoaderCF.computeIfAbsent(level, k -> new LinkedHashSet<>()).add(dfe);
+        }
+
+
         void increaseExpectedFetchCount(int level, int count) {
             expectedFetchCountPerLevel.increment(level, count);
         }
@@ -234,9 +267,13 @@ private int getObjectCountForList(List<FieldValueInfo> fieldValueInfos) {
     public void fieldFetched(ExecutionContext executionContext,
                              ExecutionStrategyParameters executionStrategyParameters,
                              DataFetcher<?> dataFetcher,
-                             Object fetchedValue) {
+                             Object fetchedValue,
+                             Supplier<DataFetchingEnvironment> dataFetchingEnvironment) {
         int level = executionStrategyParameters.getPath().getLevel();
         boolean dispatchNeeded = callStack.lock.callLocked(() -> {
+            if (DataLoaderCompletableFuture.isDataLoaderCompletableFuture(fetchedValue)) {
+                callStack.addDataLoaderDFE(level, dataFetchingEnvironment.get());
+            }
             callStack.increaseFetchCount(level);
             return dispatchIfNeeded(level);
         });
@@ -275,9 +312,84 @@ private boolean levelReady(int level) {
     }
 
     void dispatch(int level) {
-        DataLoaderRegistry dataLoaderRegistry = executionContext.getDataLoaderRegistry();
-        dataLoaderRegistry.dispatchAll();
+        // if we have any DataLoaderCFs => use new Algorithm
+        if (callStack.levelToDFEWithDataLoaderCF.get(level) != null) {
+            dispatchDLCFImpl(callStack.levelToDFEWithDataLoaderCF.get(level), true);
+        } else {
+            // otherwise dispatch all DataLoaders
+            DataLoaderRegistry dataLoaderRegistry = executionContext.getDataLoaderRegistry();
+            dataLoaderRegistry.dispatchAll();
+        }
+    }
+
+
+    public void dispatchDLCFImpl(Set<DataFetchingEnvironment> dfeToDispatchSet, boolean dispatchAll) {
+
+        // filter out all DataLoaderCFS that are matching the fields we want to dispatch
+        List<DataLoaderCompletableFuture<?>> relevantDataLoaderCompletableFutures = new ArrayList<>();
+        for (DataLoaderCompletableFuture<?> dataLoaderCompletableFuture : callStack.allDataLoaderCompletableFuture) {
+            if (dfeToDispatchSet.contains(dataLoaderCompletableFuture.dfe)) {
+                relevantDataLoaderCompletableFutures.add(dataLoaderCompletableFuture);
+            }
+        }
+        // we are cleaning up the list of all DataLoadersCFs
+        callStack.allDataLoaderCompletableFuture.removeAll(relevantDataLoaderCompletableFutures);
+
+        // means we are all done dispatching the fields
+        if (relevantDataLoaderCompletableFutures.size() == 0) {
+            callStack.fieldsFinishedDispatching.addAll(dfeToDispatchSet);
+            return;
+        }
+        // we are dispatching all data loaders and waiting for all dataLoaderCFs to complete
+        // and to finish their sync actions
+        DataLoaderCompletableFuture.waitUntilAllSyncDependentsComplete(relevantDataLoaderCompletableFutures)
+                .whenComplete((unused, throwable) ->
+                        dispatchDLCFImpl(dfeToDispatchSet, false)
+                );
+        if (dispatchAll) {
+            // if we have a mixed world with old and new DataLoaderCFs we dispatch all DataLoaders to retain compatibility
+            executionContext.getDataLoaderRegistry().dispatchAll();
+        } else {
+            // Only dispatching relevant data loaders
+            for (DataLoaderCompletableFuture<?> dlCF : relevantDataLoaderCompletableFutures) {
+                dlCF.dfe.getDataLoader(dlCF.dataLoaderName).dispatch();
+            }
+        }
     }
 
+
+    public void newDataLoaderCF(DataLoaderCompletableFuture<?> dataLoaderCompletableFuture) {
+        System.out.println("newDataLoaderCF");
+        callStack.lock.runLocked(() -> {
+            callStack.allDataLoaderCompletableFuture.add(dataLoaderCompletableFuture);
+        });
+        if (callStack.fieldsFinishedDispatching.contains(dataLoaderCompletableFuture.dfe)) {
+            System.out.println("isolated dispatch");
+            dispatchIsolatedDataLoader(dataLoaderCompletableFuture);
+        }
+
+    }
+
+    private void dispatchIsolatedDataLoader(DataLoaderCompletableFuture<?> dlCF) {
+        callStack.lock.runLocked(() -> {
+            callStack.batchWindowOfIsolatedDfeToDispatch.add(dlCF.dfe);
+            if (!callStack.batchWindowOpen) {
+                callStack.batchWindowOpen = true;
+                AtomicReference<Set<DataFetchingEnvironment>> dfesToDispatch = new AtomicReference<>();
+                Runnable runnable = () -> {
+                    callStack.lock.runLocked(() -> {
+                        dfesToDispatch.set(new LinkedHashSet<>(callStack.batchWindowOfIsolatedDfeToDispatch));
+                        callStack.batchWindowOfIsolatedDfeToDispatch.clear();
+                        callStack.batchWindowOpen = false;
+                    });
+                    dispatchDLCFImpl(dfesToDispatch.get(), false);
+                };
+                isolatedDLCFBatchWindowScheduler.schedule(runnable, BATCH_WINDOW_NANO_SECONDS, TimeUnit.NANOSECONDS);
+            }
+
+        });
+    }
+
+
 }
 
diff --git a/...instrumentation/dataloader/PerLevelDataLoaderDispatchStrategyWithDeferAlwaysDispatch.java b/...instrumentation/dataloader/PerLevelDataLoaderDispatchStrategyWithDeferAlwaysDispatch.java
@@ -7,13 +7,15 @@
 import graphql.execution.ExecutionStrategyParameters;
 import graphql.execution.FieldValueInfo;
 import graphql.schema.DataFetcher;
+import graphql.schema.DataFetchingEnvironment;
 import graphql.util.LockKit;
 import org.dataloader.DataLoaderRegistry;
 
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Supplier;
 
 /**
  * The execution of a query can be divided into 2 phases: first, the non-deferred fields are executed and only once
@@ -173,7 +175,7 @@ public void executeObjectOnFieldValuesException(Throwable t, ExecutionStrategyPa
     public void fieldFetched(ExecutionContext executionContext,
                              ExecutionStrategyParameters parameters,
                              DataFetcher<?> dataFetcher,
-                             Object fetchedValue) {
+                             Object fetchedValue, Supplier<DataFetchingEnvironment> dataFetchingEnvironment) {
 
         final boolean dispatchNeeded;
 

diff --git a/src/main/java/graphql/schema/DataFetchingEnvironment.java b/src/main/java/graphql/schema/DataFetchingEnvironment.java
@@ -237,6 +237,9 @@ public interface DataFetchingEnvironment extends IntrospectionDataFetchingEnviro
     @Nullable
     <K, V> DataLoader<K, V> getDataLoader(String dataLoaderName);
 
+
+    DataLoaderChain getDataLoaderChain();
+
     /**
      * @return the {@link org.dataloader.DataLoaderRegistry} in play
      */
@@ -270,4 +273,6 @@ public interface DataFetchingEnvironment extends IntrospectionDataFetchingEnviro
      * @return the coerced variables that have been passed to the query that is being executed
      */
     Map<String, Object> getVariables();
+
+
 }