add helper methods to configure new dispatching strategy

andimarek · May 16, 2025 · Mar 30, 2025 · Mar 30, 2025 · Mar 30, 2025 · Mar 31, 2025
commit 34bc3d7f0533c74196cfeb76c6df3331c13a9a32
diff --git a/...n/java/graphql/execution/instrumentation/dataloader/DataLoaderDispatchingContextKeys.java b/...n/java/graphql/execution/instrumentation/dataloader/DataLoaderDispatchingContextKeys.java
@@ -2,8 +2,12 @@
 
 
 import graphql.ExperimentalApi;
+import graphql.GraphQLContext;
 import org.jspecify.annotations.NullMarked;
 
+/**
+ * GraphQLContext keys related to DataLoader dispatching.
+ */
 @ExperimentalApi
 @NullMarked
 public final class DataLoaderDispatchingContextKeys {
@@ -39,4 +43,39 @@ private DataLoaderDispatchingContextKeys() {
      * Expects a boolean value.
      */
     public static final String ENABLE_DATA_LOADER_CHAINING = "__GJ_enable_data_loader_chaining";
+
+
+    /**
+     * Enables the ability that chained DataLoaders are dispatched automatically.
+     *
+     * @param graphQLContext
+     */
+    public static void enableDataLoaderChaining(GraphQLContext graphQLContext) {
+        graphQLContext.put(ENABLE_DATA_LOADER_CHAINING, true);
+    }
+
+
+    /**
+     * Sets in nanoseconds 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.
+     *
+     * @param graphQLContext
+     * @param batchWindowSizeNanoSeconds
+     */
+    public static void setDelayedDataLoaderBatchWindowSizeNanoSeconds(GraphQLContext graphQLContext, long batchWindowSizeNanoSeconds) {
+        graphQLContext.put(DELAYED_DATA_LOADER_BATCH_WINDOW_SIZE_NANO_SECONDS, batchWindowSizeNanoSeconds);
+    }
+
+    /**
+     * Sets the instance of {@link DelayedDataLoaderDispatcherExecutorFactory} that is used to create the
+     * {@link java.util.concurrent.ScheduledExecutorService} for the delayed DataLoader dispatching.
+     * <p>
+     *
+     * @param graphQLContext
+     * @param delayedDataLoaderDispatcherExecutorFactory
+     */
+    public static void setDelayedDataLoaderDispatchingExecutorFactory(GraphQLContext graphQLContext, DelayedDataLoaderDispatcherExecutorFactory delayedDataLoaderDispatcherExecutorFactory) {
+        graphQLContext.put(DELAYED_DATA_LOADER_DISPATCHING_EXECUTOR_FACTORY, delayedDataLoaderDispatcherExecutorFactory);
+    }
 }
diff --git a/...phql/execution/instrumentation/dataloader/DelayedDataLoaderDispatcherExecutorFactory.java b/...phql/execution/instrumentation/dataloader/DelayedDataLoaderDispatcherExecutorFactory.java
@@ -7,6 +7,9 @@
 
 import java.util.concurrent.ScheduledExecutorService;
 
+/**
+ * See {@link DataLoaderDispatchingContextKeys} for how to set it.
+ */
 @ExperimentalApi
 @NullMarked
 @FunctionalInterface

diff --git a/src/test/groovy/graphql/ChainedDataLoaderTest.groovy b/src/test/groovy/graphql/ChainedDataLoaderTest.groovy
@@ -1,8 +1,8 @@
 package graphql
 
 import graphql.execution.ExecutionId
-import graphql.execution.instrumentation.dataloader.DelayedDataLoaderDispatcherExecutorFactory
 import graphql.execution.instrumentation.dataloader.DataLoaderDispatchingContextKeys
+import graphql.execution.instrumentation.dataloader.DelayedDataLoaderDispatcherExecutorFactory
 import graphql.schema.DataFetcher
 import org.awaitility.Awaitility
 import org.dataloader.BatchLoader
@@ -72,7 +72,8 @@ class ChainedDataLoaderTest extends Specification {
         def graphQL = GraphQL.newGraphQL(schema).build()
 
         def query = "{ dogName catName } "
-        def ei = newExecutionInput(query).graphQLContext([(DataLoaderDispatchingContextKeys.ENABLE_DATA_LOADER_CHAINING): true]).dataLoaderRegistry(dataLoaderRegistry).build()
+        def ei = newExecutionInput(query).dataLoaderRegistry(dataLoaderRegistry).build()
+        DataLoaderDispatchingContextKeys.enableDataLoaderChaining(ei.graphQLContext)
 
         when:
         def efCF = graphQL.executeAsync(ei)
@@ -383,8 +384,8 @@ class ChainedDataLoaderTest extends Specification {
         def query = "{ foo bar } "
         def ei = newExecutionInput(query).graphQLContext([(DataLoaderDispatchingContextKeys.ENABLE_DATA_LOADER_CHAINING): true]).dataLoaderRegistry(dataLoaderRegistry).build()
 
-        // make the window to 50ms
-        ei.getGraphQLContext().put(DataLoaderDispatchingContextKeys.DELAYED_DATA_LOADER_BATCH_WINDOW_SIZE_NANO_SECONDS, 1_000_000L * 250)
+        // make the window  250ms
+        DataLoaderDispatchingContextKeys.setDelayedDataLoaderBatchWindowSizeNanoSeconds(ei.graphQLContext, 1_000_000L * 250)
 
         when:
         def efCF = graphQL.executeAsync(ei)
@@ -435,7 +436,7 @@ class ChainedDataLoaderTest extends Specification {
 
 
         ScheduledExecutorService scheduledExecutorService = Mock()
-        ei.getGraphQLContext().put(DataLoaderDispatchingContextKeys.DELAYED_DATA_LOADER_DISPATCHING_EXECUTOR_FACTORY, new DelayedDataLoaderDispatcherExecutorFactory() {
+        DataLoaderDispatchingContextKeys.setDelayedDataLoaderDispatchingExecutorFactory(ei.getGraphQLContext(), new DelayedDataLoaderDispatcherExecutorFactory() {
             @Override
             ScheduledExecutorService createExecutor(ExecutionId executionId, GraphQLContext graphQLContext) {
                 return scheduledExecutorService