scheduled executor can be configured for delayed dispatching

andimarek · May 16, 2025 · Mar 30, 2025 · Mar 30, 2025 · Mar 30, 2025 · Mar 31, 2025
commit 401d09d40da0964916ed47ad39a35b05b491c8ad
diff --git a/...phql/execution/instrumentation/dataloader/DelayedDataLoaderDispatcherExecutorFactory.java b/...phql/execution/instrumentation/dataloader/DelayedDataLoaderDispatcherExecutorFactory.java
@@ -0,0 +1,26 @@
+package graphql.execution.instrumentation.dataloader;
+
+import graphql.ExperimentalApi;
+import graphql.GraphQLContext;
+import graphql.execution.ExecutionId;
+import org.jspecify.annotations.NullMarked;
+
+import java.util.concurrent.ScheduledExecutorService;
+
+@ExperimentalApi
+@NullMarked
+@FunctionalInterface
+public interface DelayedDataLoaderDispatcherExecutorFactory {
+
+    /**
+     * Called once per execution to create the {@link ScheduledExecutorService} for the delayed DataLoader dispatching.
+     *
+     * Will only called if needed, i.e. if there are delayed DataLoaders.
+     *
+     * @param executionId
+     * @param graphQLContext
+     *
+     * @return
+     */
+    ScheduledExecutorService createExecutor(ExecutionId executionId, GraphQLContext graphQLContext);
+}
diff --git a/src/main/java/graphql/execution/instrumentation/dataloader/DispatchingContextKeys.java b/src/main/java/graphql/execution/instrumentation/dataloader/DispatchingContextKeys.java
@@ -2,8 +2,10 @@
 
 
 import graphql.ExperimentalApi;
+import org.jspecify.annotations.NullMarked;
 
 @ExperimentalApi
+@NullMarked
 public final class DispatchingContextKeys {
     private DispatchingContextKeys() {
     }
@@ -12,6 +14,18 @@ private DispatchingContextKeys() {
      * In nano seconds, the batch window size for delayed DataLoaders.
      * That is for DataLoaders, that are not batched as part of the normal per level
      * dispatching, because they were created after the level was already dispatched.
+     *
+     * Expect Integer values
+     *
+     * Default is 500_000 (0.5 ms)
      */
-    public static final String BATCH_WINDOW_DELAYED_DL_NANO_SECONDS = "__batch_window_delayed_dl_nano_seconds";
+    public static final String DELAYED_DATA_LOADER_BATCH_WINDOW_SIZE_NANO_SECONDS = "__GJ_delayed_data_loader_batch_window_size_nano_seconds";
+
+    /**
+     * An instance of {@link DelayedDataLoaderDispatcherExecutorFactory} that is used to create the
+     * {@link java.util.concurrent.ScheduledExecutorService} for the delayed DataLoader dispatching.
+     *
+     * Default is one static executor thread pool with a single thread.
+     */
+    public static final String DELAYED_DATA_LOADER_DISPATCHING_EXECUTOR_FACTORY = "__GJ_delayed_data_loader_dispatching_executor_factory";
 }
diff --git a/...java/graphql/execution/instrumentation/dataloader/PerLevelDataLoaderDispatchStrategy.java b/...java/graphql/execution/instrumentation/dataloader/PerLevelDataLoaderDispatchStrategy.java
@@ -8,6 +8,7 @@
 import graphql.execution.FieldValueInfo;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
+import graphql.util.InterThreadMemoizedSupplier;
 import graphql.util.LockKit;
 import org.dataloader.DataLoader;
 import org.dataloader.DataLoaderRegistry;
@@ -34,8 +35,12 @@ public class PerLevelDataLoaderDispatchStrategy implements DataLoaderDispatchStr
     private final ExecutionContext executionContext;
     private final int batchWindowNs;
 
-    static final ScheduledExecutorService delayedDLCFBatchWindowScheduler = Executors.newSingleThreadScheduledExecutor();
-    static final int BATCH_WINDOW_NANO_SECONDS_DEFAULT = 500_000;
+    private final InterThreadMemoizedSupplier<ScheduledExecutorService> delayedDataLoaderDispatchExecutor;
+
+    static final InterThreadMemoizedSupplier<ScheduledExecutorService> defaultDelayedDLCFBatchWindowScheduler
+            = new InterThreadMemoizedSupplier<>(Executors::newSingleThreadScheduledExecutor);
+
+    static final int DEFAULT_BATCH_WINDOW_NANO_SECONDS_DEFAULT = 500_000;
 
 
     private static class CallStack {
@@ -154,12 +159,20 @@ public PerLevelDataLoaderDispatchStrategy(ExecutionContext executionContext) {
         this.callStack = new CallStack();
         this.executionContext = executionContext;
 
-        Integer batchWindowNs = executionContext.getGraphQLContext().get(DispatchingContextKeys.BATCH_WINDOW_DELAYED_DL_NANO_SECONDS);
+        Integer batchWindowNs = executionContext.getGraphQLContext().get(DispatchingContextKeys.DELAYED_DATA_LOADER_BATCH_WINDOW_SIZE_NANO_SECONDS);
         if (batchWindowNs != null) {
             this.batchWindowNs = batchWindowNs;
         } else {
-            this.batchWindowNs = BATCH_WINDOW_NANO_SECONDS_DEFAULT;
+            this.batchWindowNs = DEFAULT_BATCH_WINDOW_NANO_SECONDS_DEFAULT;
         }
+
+        this.delayedDataLoaderDispatchExecutor = new InterThreadMemoizedSupplier<>(() -> {
+            DelayedDataLoaderDispatcherExecutorFactory delayedDataLoaderDispatcherExecutorFactory = executionContext.getGraphQLContext().get(DispatchingContextKeys.DELAYED_DATA_LOADER_DISPATCHING_EXECUTOR_FACTORY);
+            if (delayedDataLoaderDispatcherExecutorFactory != null) {
+                return delayedDataLoaderDispatcherExecutorFactory.createExecutor(executionContext.getExecutionId(), executionContext.getGraphQLContext());
+            }
+            return defaultDelayedDLCFBatchWindowScheduler.get();
+        });
     }
 
     @Override
@@ -320,11 +333,9 @@ private boolean levelReady(int level) {
     }
 
     void dispatch(int level) {
-        System.out.println("dispatching level " + level);
         // if we have any DataLoaderCFs => use new Algorithm
         Set<ResultPathWithDataLoader> resultPathWithDataLoaders = callStack.levelToResultPathWithDataLoader.get(level);
         if (resultPathWithDataLoaders != null) {
-            System.out.println("dispatching level " + level + " with " + resultPathWithDataLoaders.size() + " DataLoaderCFs" + " this: " + this);
             dispatchDLCFImpl(resultPathWithDataLoaders
                     .stream()
                     .map(resultPathWithDataLoader -> resultPathWithDataLoader.resultPath)
@@ -375,22 +386,15 @@ public void dispatchDLCFImpl(Set<String> resultPathsToDispatch, boolean dispatch
 
 
     public void newDataLoaderCF(String resultPath, int level, DataLoader dataLoader) {
-        System.out.println("newDataLoaderCF");
         ResultPathWithDataLoader resultPathWithDataLoader = new ResultPathWithDataLoader(resultPath, level, dataLoader);
         callStack.lock.runLocked(() -> {
             callStack.allResultPathWithDataLoader.add(resultPathWithDataLoader);
             if (!callStack.dispatchedLevels.contains(level)) {
-                System.out.println("not finished dispatching level " + level);
                 callStack.addDataLoaderDFE(level, resultPathWithDataLoader);
-            } else {
-                System.out.println("already finished dispatching level " + level);
             }
         });
         if (callStack.dispatchedLevels.contains(level)) {
-            System.out.println("delayed dispatch");
             dispatchDelayedDataLoader(resultPathWithDataLoader);
-        } else {
-            System.out.println("normal dispatch");
         }
 
 
@@ -408,16 +412,18 @@ private void dispatchDelayedDataLoader(ResultPathWithDataLoader resultPathWithDa
                         callStack.batchWindowOfDelayedDataLoaderToDispatch.clear();
                         callStack.batchWindowOpen = false;
                     });
-                    System.out.println("start dispatch with " + dfesToDispatch.get().size());
                     dispatchDLCFImpl(dfesToDispatch.get(), false);
                 };
-                delayedDLCFBatchWindowScheduler.schedule(runnable, this.batchWindowNs, TimeUnit.NANOSECONDS);
+                try {
+                    delayedDataLoaderDispatchExecutor.get().schedule(runnable, this.batchWindowNs, TimeUnit.NANOSECONDS);
+                } catch (Exception e) {
+                    e.printStackTrace();
+                }
             }
 
         });
     }
 
-
     private static class ResultPathWithDataLoader {
         final String resultPath;
         final int level;

diff --git a/src/test/groovy/graphql/ChainedDataLoaderTest.groovy b/src/test/groovy/graphql/ChainedDataLoaderTest.groovy
@@ -1,5 +1,8 @@
 package graphql
 
+
+import graphql.execution.ExecutionId
+import graphql.execution.instrumentation.dataloader.DelayedDataLoaderDispatcherExecutorFactory
 import graphql.execution.instrumentation.dataloader.DispatchingContextKeys
 import graphql.schema.DataFetcher
 import org.dataloader.BatchLoader
@@ -8,6 +11,9 @@ import org.dataloader.DataLoaderFactory
 import org.dataloader.DataLoaderRegistry
 import spock.lang.Specification
 
+import java.util.concurrent.Executors
+import java.util.concurrent.ScheduledExecutorService
+import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicInteger
 
 import static graphql.ExecutionInput.newExecutionInput
@@ -276,7 +282,7 @@ class ChainedDataLoaderTest extends Specification {
         def ei = newExecutionInput(query).dataLoaderRegistry(dataLoaderRegistry).build()
 
         // make the window large enough to avoid flaky tests
-        ei.getGraphQLContext().put(DispatchingContextKeys.BATCH_WINDOW_DELAYED_DL_NANO_SECONDS, 2_000_000)
+        ei.getGraphQLContext().put(DispatchingContextKeys.DELAYED_DATA_LOADER_BATCH_WINDOW_SIZE_NANO_SECONDS, 2_000_000)
 
         when:
         def er = graphQL.execute(ei)
@@ -285,5 +291,63 @@ class ChainedDataLoaderTest extends Specification {
         batchLoadCalls.get() == 1
     }
 
+    def "executor for delayed dispatching can be configured"() {
+        given:
+        def sdl = '''
+
+        type Query {
+          foo: String
+         bar: String
+        }
+        '''
+        BatchLoader<String, String> batchLoader = { keys ->
+            return supplyAsync {
+                Thread.sleep(250)
+                return keys;
+            }
+        }
+
+        DataLoader<String, String> nameDataLoader = DataLoaderFactory.newDataLoader(batchLoader);
+
+        DataLoaderRegistry dataLoaderRegistry = new DataLoaderRegistry();
+        dataLoaderRegistry.register("dl", nameDataLoader);
+
+        def fooDF = { env ->
+            return supplyAsync {
+                Thread.sleep(1000)
+                return "fooFirstValue"
+            }.thenCompose {
+                return env.getDataLoader("dl").load(it)
+            }
+        } as DataFetcher
+
+
+        def fetchers = ["Query": ["foo": fooDF]]
+        def schema = TestUtil.schema(sdl, fetchers)
+        def graphQL = GraphQL.newGraphQL(schema).build()
+
+        def query = "{ foo } "
+        def ei = newExecutionInput(query).dataLoaderRegistry(dataLoaderRegistry).build()
+
+
+        ScheduledExecutorService scheduledExecutorService = Mock()
+        ei.getGraphQLContext().put(DispatchingContextKeys.DELAYED_DATA_LOADER_DISPATCHING_EXECUTOR_FACTORY, new DelayedDataLoaderDispatcherExecutorFactory() {
+            @Override
+            ScheduledExecutorService createExecutor(ExecutionId executionId, GraphQLContext graphQLContext) {
+                return scheduledExecutorService
+            }
+        })
+
+
+        when:
+        def er = graphQL.execute(ei)
+
+        then:
+        er.data == [foo: "fooFirstValue"]
+        1 * scheduledExecutorService.schedule(_ as Runnable, _ as Long, _ as TimeUnit) >> { Runnable runnable, Long delay, TimeUnit timeUnit ->
+            return Executors.newSingleThreadScheduledExecutor().schedule(runnable, delay, timeUnit)
+        }
+
+    }
 
 }