introducing of a special CompletableFuture enabling scheduling of nested/chained DataLoader calls

andimarek · andimarek · commit 0a27e76bc164 · 2025-03-30T18:21:55.000+10:00
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
@@ -496,7 +496,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/DataLoaderCF.java b/src/main/java/graphql/execution/instrumentation/dataloader/DataLoaderCF.java
@@ -0,0 +1,99 @@
+package graphql.execution.instrumentation.dataloader;
+
+import graphql.ExperimentalApi;
+import graphql.Internal;
+import graphql.execution.DataLoaderDispatchStrategy;
+import graphql.execution.ExecutionContext;
+import graphql.schema.DataFetchingEnvironment;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
+import java.util.function.Supplier;
+
+import static graphql.execution.Execution.EXECUTION_CONTEXT_KEY;
+
+@Internal
+public class DataLoaderCF<T> extends CompletableFuture<T> {
+    final DataFetchingEnvironment dfe;
+    final String dataLoaderName;
+    final Object key;
+    final CompletableFuture<Object> dataLoaderCF;
+
+    volatile CountDownLatch latch;
+
+    public DataLoaderCF(DataFetchingEnvironment dfe, String dataLoaderName, Object key) {
+        this.dfe = dfe;
+        this.dataLoaderName = dataLoaderName;
+        this.key = key;
+        if (dataLoaderName != null) {
+            dataLoaderCF = dfe.getDataLoaderRegistry().getDataLoader(dataLoaderName).load(key);
+            dataLoaderCF.whenComplete((value, throwable) -> {
+                System.out.println("underlying DataLoader completed");
+                if (throwable != null) {
+                    completeExceptionally(throwable);
+                } else {
+                    complete((T) value);
+                }
+                // post completion hook
+                if (latch != null) {
+                    latch.countDown();
+                }
+            });
+        } else {
+            dataLoaderCF = null;
+        }
+    }
+
+    DataLoaderCF() {
+        this.dfe = null;
+        this.dataLoaderName = null;
+        this.key = null;
+        dataLoaderCF = null;
+    }
+
+    @Override
+    public <U> CompletableFuture<U> newIncompleteFuture() {
+        return new DataLoaderCF<>();
+    }
+
+    public static boolean isDataLoaderCF(Object object) {
+        return object instanceof DataLoaderCF;
+    }
+
+    @ExperimentalApi
+    public static <T> CompletableFuture<T> newDataLoaderCF(DataFetchingEnvironment dfe, String dataLoaderName, Object key) {
+        DataLoaderCF<T> result = new DataLoaderCF<>(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;
+    }
+
+
+    @ExperimentalApi
+    public static <U> CompletableFuture<U> supplyAsyncDataLoaderCF(DataFetchingEnvironment env, Supplier<U> supplier) {
+        DataLoaderCF<U> d = new DataLoaderCF<>(env, null, null);
+        d.defaultExecutor().execute(() -> {
+            d.complete(supplier.get());
+        });
+        return d;
+
+    }
+
+    @ExperimentalApi
+    public static <U> CompletableFuture<U> wrap(DataFetchingEnvironment env, CompletableFuture<U> completableFuture) {
+        DataLoaderCF<U> d = new DataLoaderCF<>(env, null, null);
+        completableFuture.whenComplete((u, ex) -> {
+            if (ex != null) {
+                d.completeExceptionally(ex);
+            } else {
+                d.complete(u);
+            }
+        });
+        return d;
+    }
+
+
+}
diff --git a/src/main/java/graphql/execution/instrumentation/dataloader/PerLevelDataLoaderDispatchStrategy.java b/src/main/java/graphql/execution/instrumentation/dataloader/PerLevelDataLoaderDispatchStrategy.java
@@ -7,19 +7,33 @@
 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.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.CountDownLatch;
+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 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 +48,27 @@ private static class CallStack {
 
         private final Set<Integer> dispatchedLevels = new LinkedHashSet<>();
 
+        // fields only relevant when a DataLoaderCF is involved
+        private final List<DataLoaderCF<?>> allDataLoaderCF = 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 +265,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 (DataLoaderCF.isDataLoaderCF(fetchedValue)) {
+                callStack.addDataLoaderDFE(level, dataFetchingEnvironment.get());
+            }
             callStack.increaseFetchCount(level);
             return dispatchIfNeeded(level);
         });
@@ -275,9 +310,89 @@ private boolean levelReady(int level) {
     }
 
     void dispatch(int level) {
-        DataLoaderRegistry dataLoaderRegistry = executionContext.getDataLoaderRegistry();
-        dataLoaderRegistry.dispatchAll();
+        if (callStack.levelToDFEWithDataLoaderCF.size() > 0) {
+            dispatchDLCFImpl(callStack.levelToDFEWithDataLoaderCF.get(level));
+        } else {
+            DataLoaderRegistry dataLoaderRegistry = executionContext.getDataLoaderRegistry();
+            dataLoaderRegistry.dispatchAll();
+        }
     }
 
+
+    public void dispatchDLCFImpl(Set<DataFetchingEnvironment> dfeToDispatchSet) {
+
+        // filter out all DataLoaderCFS that are matching the fields we want to dispatch
+        List<DataLoaderCF<?>> relevantDataLoaderCFs = new ArrayList<>();
+        for (DataLoaderCF<?> dataLoaderCF : callStack.allDataLoaderCF) {
+            if (dfeToDispatchSet.contains(dataLoaderCF.dfe)) {
+                relevantDataLoaderCFs.add(dataLoaderCF);
+            }
+        }
+        // we are cleaning up the list of all DataLoadersCFs
+        callStack.allDataLoaderCF.removeAll(relevantDataLoaderCFs);
+
+        // means we are all done dispatching the fields
+        if (relevantDataLoaderCFs.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
+        CountDownLatch countDownLatch = new CountDownLatch(relevantDataLoaderCFs.size());
+        for (DataLoaderCF dlCF : relevantDataLoaderCFs) {
+            dlCF.latch = countDownLatch;
+        }
+        // TODO: this should be done async or in a more regulated way with a configurable thread pool or so
+        new Thread(() -> {
+            try {
+                // waiting until all sync codes for all DL CFs are run
+                countDownLatch.await();
+            } catch (InterruptedException e) {
+                throw new RuntimeException(e);
+            }
+            // now we handle all new DataLoaders
+            dispatchDLCFImpl(dfeToDispatchSet);
+        }).start();
+        // Only dispatching relevant data loaders
+        for (DataLoaderCF dlCF : relevantDataLoaderCFs) {
+            dlCF.dfe.getDataLoader(dlCF.dataLoaderName).dispatch();
+        }
+//        executionContext.getDataLoaderRegistry().dispatchAll();
+    }
+
+
+    public void newDataLoaderCF(DataLoaderCF<?> dataLoaderCF) {
+        System.out.println("newDataLoaderCF");
+        callStack.lock.runLocked(() -> {
+            callStack.allDataLoaderCF.add(dataLoaderCF);
+        });
+        if (callStack.fieldsFinishedDispatching.contains(dataLoaderCF.dfe)) {
+            System.out.println("isolated dispatch");
+            dispatchIsolatedDataLoader(dataLoaderCF);
+        }
+
+    }
+
+    private void dispatchIsolatedDataLoader(DataLoaderCF<?> 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());
+                };
+                isolatedDLCFBatchWindowScheduler.schedule(runnable, BATCH_WINDOW_NANO_SECONDS, TimeUnit.NANOSECONDS);
+            }
+
+        });
+    }
+
+
 }
 
diff --git a/src/main/java/graphql/execution/instrumentation/dataloader/PerLevelDataLoaderDispatchStrategyWithDeferAlwaysDispatch.java b/src/main/java/graphql/execution/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/test/groovy/graphql/DataLoaderCFTest.groovy b/src/test/groovy/graphql/DataLoaderCFTest.groovy
