graphql-java · Alex079 · Apr 2, 2022 · Apr 2, 2022 · Apr 8, 2022 · Apr 20, 2022
diff --git a/build.gradle b/build.gradle
@@ -87,7 +87,7 @@ dependencies {
     compileOnly 'org.jetbrains:annotations:23.0.0'
     implementation 'org.antlr:antlr4-runtime:' + antlrVersion
     implementation 'org.slf4j:slf4j-api:' + slf4jVersion
-    api 'com.graphql-java:java-dataloader:3.1.2'
+    api 'com.graphql-java:java-dataloader:3.1.3-SNAPSHOT'
     api 'org.reactivestreams:reactive-streams:' + reactiveStreamsVersion
     antlr 'org.antlr:antlr4:' + antlrVersion
     implementation 'com.google.guava:guava:31.0.1-jre'

diff --git a/src/main/java/graphql/execution/AsyncExecutionStrategy.java b/src/main/java/graphql/execution/AsyncExecutionStrategy.java
@@ -70,8 +70,9 @@ public CompletableFuture<ExecutionResult> execute(ExecutionContext executionCont
                 return;
             }
             List<CompletableFuture<ExecutionResult>> executionResultFuture = map(completeValueInfos, FieldValueInfo::getFieldValue);
-            executionStrategyCtx.onFieldValuesInfo(completeValueInfos);
-            Async.each(executionResultFuture).whenComplete(handleResultsConsumer);
+            Async.each(executionResultFuture)
+                .thenCombine(executionStrategyCtx.onFieldValuesInfo(completeValueInfos), (result, __) -> result)
+                .whenComplete(handleResultsConsumer);
         }).exceptionally((ex) -> {
             // if there are any issues with combining/handling the field results,
             // complete the future at all costs and bubble up any thrown exception so

diff --git a/src/main/java/graphql/execution/ExecutionStrategy.java b/src/main/java/graphql/execution/ExecutionStrategy.java
@@ -286,8 +286,7 @@ protected CompletableFuture<FetchedValue> fetchField(ExecutionContext executionC
             fetchedValue = new CompletableFuture<>();
             fetchedValue.completeExceptionally(e);
         }
-        fetchCtx.onDispatched(fetchedValue);
-        return fetchedValue
+        return fetchCtx.onDispatched(fetchedValue)
                 .handle((result, exception) -> {
                     fetchCtx.onCompleted(result, exception);
                     if (exception != null) {

diff --git a/src/main/java/graphql/execution/instrumentation/ChainedInstrumentation.java b/src/main/java/graphql/execution/instrumentation/ChainedInstrumentation.java
@@ -230,8 +230,11 @@ private static class ChainedInstrumentationContext<T> implements Instrumentation
         }
 
         @Override
-        public void onDispatched(CompletableFuture<T> result) {
-            contexts.forEach(context -> context.onDispatched(result));
+        public CompletableFuture<T> onDispatched(CompletableFuture<T> result) {
+            return CompletableFuture.allOf(contexts.stream()
+                    .map(context -> context.onDispatched(result))
+                    .toArray(CompletableFuture<?>[]::new))
+                .thenCompose(__ -> result);
         }
 
         @Override
@@ -249,8 +252,11 @@ private static class ChainedExecutionStrategyInstrumentationContext implements E
         }
 
         @Override
-        public void onDispatched(CompletableFuture<ExecutionResult> result) {
-            contexts.forEach(context -> context.onDispatched(result));
+        public CompletableFuture<ExecutionResult> onDispatched(CompletableFuture<ExecutionResult> result) {
+            return CompletableFuture.allOf(contexts.stream()
+                    .map(context -> context.onDispatched(result))
+                    .toArray(CompletableFuture<?>[]::new))
+                .thenCompose(__ -> result);
         }
 
         @Override
@@ -259,8 +265,10 @@ public void onCompleted(ExecutionResult result, Throwable t) {
         }
 
         @Override
-        public void onFieldValuesInfo(List<FieldValueInfo> fieldValueInfoList) {
-            contexts.forEach(context -> context.onFieldValuesInfo(fieldValueInfoList));
+        public CompletableFuture<Void> onFieldValuesInfo(List<FieldValueInfo> fieldValueInfoList) {
+            return CompletableFuture.allOf(contexts.stream()
+                    .map(context -> context.onFieldValuesInfo(fieldValueInfoList))
+                    .toArray(CompletableFuture<?>[]::new));
         }
 
     }

diff --git a/src/main/java/graphql/execution/instrumentation/ExecutionStrategyInstrumentationContext.java b/src/main/java/graphql/execution/instrumentation/ExecutionStrategyInstrumentationContext.java
@@ -5,12 +5,13 @@
 import graphql.execution.FieldValueInfo;
 
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
 
 @PublicSpi
 public interface ExecutionStrategyInstrumentationContext extends InstrumentationContext<ExecutionResult> {
 
-    default void onFieldValuesInfo(List<FieldValueInfo> fieldValueInfoList) {
-
+    default CompletableFuture<Void> onFieldValuesInfo(List<FieldValueInfo> fieldValueInfoList) {
+        return CompletableFuture.completedFuture(null);
     }
 
     default void onFieldValuesException() {

diff --git a/src/main/java/graphql/execution/instrumentation/InstrumentationContext.java b/src/main/java/graphql/execution/instrumentation/InstrumentationContext.java
@@ -19,8 +19,9 @@ public interface InstrumentationContext<T> {
      * This is invoked when the instrumentation step is initially dispatched
      *
      * @param result the result of the step as a completable future
+     * @return instrumented or unmodified result
      */
-    void onDispatched(CompletableFuture<T> result);
+    CompletableFuture<T> onDispatched(CompletableFuture<T> result);
 
     /**
      * This is invoked when the instrumentation step is fully completed

diff --git a/src/main/java/graphql/execution/instrumentation/SimpleInstrumentation.java b/src/main/java/graphql/execution/instrumentation/SimpleInstrumentation.java
@@ -45,8 +45,8 @@ public InstrumentationContext<List<ValidationError>> beginValidation(Instrumenta
     public ExecutionStrategyInstrumentationContext beginExecutionStrategy(InstrumentationExecutionStrategyParameters parameters) {
         return new ExecutionStrategyInstrumentationContext() {
             @Override
-            public void onDispatched(CompletableFuture<ExecutionResult> result) {
-
+            public CompletableFuture<ExecutionResult> onDispatched(CompletableFuture<ExecutionResult> result) {
+                return result;
             }
 
             @Override

diff --git a/src/main/java/graphql/execution/instrumentation/SimpleInstrumentationContext.java b/src/main/java/graphql/execution/instrumentation/SimpleInstrumentationContext.java
@@ -39,10 +39,11 @@ private SimpleInstrumentationContext(Consumer<CompletableFuture<T>> codeToRunOnD
     }
 
     @Override
-    public void onDispatched(CompletableFuture<T> result) {
+    public CompletableFuture<T> onDispatched(CompletableFuture<T> result) {
         if (codeToRunOnDispatch != null) {
             codeToRunOnDispatch.accept(result);
         }
+        return result;
     }
 
     @Override

diff --git a/...ava/graphql/execution/instrumentation/dataloader/DataLoaderDispatcherInstrumentation.java b/...ava/graphql/execution/instrumentation/dataloader/DataLoaderDispatcherInstrumentation.java
@@ -80,18 +80,18 @@ public DataFetcher<?> instrumentDataFetcher(DataFetcher<?> dataFetcher, Instrume
             return dataFetcher;
         }
         //
-        // currently only AsyncExecutionStrategy with DataLoader and hence this allows us to "dispatch"
-        // on every object if its not using aggressive batching for other execution strategies
+        // currently, only AsyncExecutionStrategy with DataLoader and hence this allows us to "dispatch"
+        // on every object if it's not using aggressive batching for other execution strategies
         // which allows them to work if used.
         return (DataFetcher<Object>) environment -> {
             Object obj = dataFetcher.get(environment);
-            immediatelyDispatch(state);
-            return obj;
+            return Async.toCompletableFuture(obj)
+                .thenCombine(immediatelyDispatch(state), (result, __) -> result);
         };
     }
 
-    private void immediatelyDispatch(DataLoaderDispatcherInstrumentationState state) {
-        state.getApproach().dispatch();
+    private CompletableFuture<Void> immediatelyDispatch(DataLoaderDispatcherInstrumentationState state) {
+        return state.getApproach().dispatch();
     }
 
     @Override
@@ -128,7 +128,8 @@ public ExecutionStrategyInstrumentationContext beginExecutionStrategy(Instrument
         if (state.hasNoDataLoaders()) {
             return new ExecutionStrategyInstrumentationContext() {
                 @Override
-                public void onDispatched(CompletableFuture<ExecutionResult> result) {
+                public CompletableFuture<ExecutionResult> onDispatched(CompletableFuture<ExecutionResult> result) {
+                    return result;
                 }
 
                 @Override

diff --git a/src/main/java/graphql/execution/instrumentation/dataloader/FieldLevelTrackingApproach.java b/src/main/java/graphql/execution/instrumentation/dataloader/FieldLevelTrackingApproach.java
@@ -131,8 +131,8 @@ ExecutionStrategyInstrumentationContext beginExecutionStrategy(InstrumentationEx
 
         return new ExecutionStrategyInstrumentationContext() {
             @Override
-            public void onDispatched(CompletableFuture<ExecutionResult> result) {
-
+            public CompletableFuture<ExecutionResult> onDispatched(CompletableFuture<ExecutionResult> result) {
+                return result;
             }
 
             @Override
@@ -141,14 +141,15 @@ public void onCompleted(ExecutionResult result, Throwable t) {
             }
 
             @Override
-            public void onFieldValuesInfo(List<FieldValueInfo> fieldValueInfoList) {
+            public CompletableFuture<Void> onFieldValuesInfo(List<FieldValueInfo> fieldValueInfoList) {
                 boolean dispatchNeeded;
                 synchronized (callStack) {
                     dispatchNeeded = handleOnFieldValuesInfo(fieldValueInfoList, callStack, curLevel);
                 }
                 if (dispatchNeeded) {
-                    dispatch();
+                    return dispatch();
                 }
+                return CompletableFuture.completedFuture(null);
             }
 
             @Override
@@ -197,16 +198,16 @@ public InstrumentationContext<Object> beginFieldFetch(InstrumentationFieldFetchP
         return new InstrumentationContext<Object>() {
 
             @Override
-            public void onDispatched(CompletableFuture result) {
+            public CompletableFuture<Object> onDispatched(CompletableFuture<Object> result) {
                 boolean dispatchNeeded;
                 synchronized (callStack) {
                     callStack.increaseFetchCount(level);
                     dispatchNeeded = dispatchIfNeeded(callStack, level);
                 }
                 if (dispatchNeeded) {
-                    dispatch();
+                    return result.thenCombine(dispatch(), (value, __) -> value);
                 }
-
+                return result;
             }
 
             @Override
@@ -241,12 +242,17 @@ private boolean levelReady(CallStack callStack, int level) {
         return false;
     }
 
-    void dispatch() {
+    CompletableFuture<Void> dispatch() {
         DataLoaderRegistry dataLoaderRegistry = getDataLoaderRegistry();
         if (log.isDebugEnabled()) {
             log.debug("Dispatching data loaders ({})", dataLoaderRegistry.getKeys());
         }
-        dataLoaderRegistry.dispatchAll();
+        return dataLoaderRegistry.dispatch().thenApply(count -> {
+            if (log.isDebugEnabled()) {
+                log.debug("Dispatched {} keys to load", count);
+            }
+            return null;
+        });
     }
 
     private DataLoaderRegistry getDataLoaderRegistry() {

diff --git a/src/test/groovy/graphql/ComposedDataLoadersTest.groovy b/src/test/groovy/graphql/ComposedDataLoadersTest.groovy
@@ -0,0 +1,69 @@
+package graphql
+
+import graphql.schema.GraphQLSchema
+import org.dataloader.BatchLoader
+import org.dataloader.DataLoader
+import org.dataloader.DataLoaderFactory
+import org.dataloader.DataLoaderRegistry
+import spock.lang.Specification
+
+import static graphql.Scalars.GraphQLInt
+import static graphql.Scalars.GraphQLString
+import static graphql.schema.GraphQLFieldDefinition.newFieldDefinition
+import static graphql.schema.GraphQLObjectType.newObject
+import static java.util.concurrent.CompletableFuture.completedFuture
+
+class ComposedDataLoadersTest extends Specification {
+
+    private static final def QUERY = '{ level1field1 { level2field1 level2field2 } }'
+
+    private static final BatchLoader<Integer, Integer> LOADER = v -> completedFuture(v.collect() {++it})
+
+    private static final def REGISTRY = DataLoaderRegistry.newRegistry()
+            .register('loader', DataLoaderFactory.newDataLoader(LOADER))
+            .build()
+
+    private static final def DUMMY_QUERY_SUBTYPE = newObject()
+            .name('Subtype')
+            .field(newFieldDefinition()
+                    .name('level2field1')
+                    .type(GraphQLString)
+                    .dataFetcher(environment -> {
+                        DataLoader<Integer, Integer> loader = environment.getDataLoader('loader')
+                        loader.load(1).thenCompose(loader::load)
+                    }))
+            .field(newFieldDefinition()
+                    .name('level2field2')
+                    .type(GraphQLInt)
+                    .dataFetcher(environment -> {
+                        DataLoader<Integer, Integer> loader = environment.getDataLoader('loader')
+                        loader.load(5).thenCompose(loader::load)
+                    }))
+            .build()
+
+    private static final def DUMMY_QUERY_TYPE = newObject()
+            .name('Query')
+            .field(newFieldDefinition()
+                    .name('level1field1')
+                    .type(DUMMY_QUERY_SUBTYPE)
+                    .dataFetcher(environment -> {
+                        DataLoader<Integer, Integer> loader = environment.getDataLoader('loader')
+                        loader.load(11).thenCompose(loader::load).thenApply {[:]}
+                    })
+            .build())
+
+    def 'composed data loaders must complete to finish field fetching level'() {
+        given:
+        def graphQL = GraphQL.newGraphQL(GraphQLSchema.newSchema().query(DUMMY_QUERY_TYPE).build()).build()
+        when:
+        def result = graphQL.execute(ExecutionInput.newExecutionInput(QUERY).dataLoaderRegistry(REGISTRY).build()).data
+        then:
+        result == [
+                level1field1: [
+                        level2field1: '3',
+                        level2field2: 7
+                ]
+        ]
+    }
+
+}
diff --git a/src/test/groovy/graphql/execution/AsyncExecutionStrategyTest.groovy b/src/test/groovy/graphql/execution/AsyncExecutionStrategyTest.groovy
@@ -246,17 +246,17 @@ class AsyncExecutionStrategyTest extends Specification {
                         return new ExecutionStrategyInstrumentationContext() {
 
                             @Override
-                            void onFieldValuesInfo(List<FieldValueInfo> fieldValueInfoList) {
+                            CompletableFuture<Void> onFieldValuesInfo(List<FieldValueInfo> fieldValueInfoList) {
                                 throw new RuntimeException("Exception raised from instrumentation")
                             }
 
                             @Override
-                            public void onDispatched(CompletableFuture<ExecutionResult> result) {
-
+                            CompletableFuture<ExecutionResult> onDispatched(CompletableFuture<ExecutionResult> result) {
+                                return result;
                             }
 
                             @Override
-                            public void onCompleted(ExecutionResult result, Throwable t) {
+                            void onCompleted(ExecutionResult result, Throwable t) {
 
                             }
                         }

diff --git a/src/test/groovy/graphql/execution/instrumentation/InstrumentationTest.groovy b/src/test/groovy/graphql/execution/instrumentation/InstrumentationTest.groovy
@@ -146,7 +146,7 @@ class InstrumentationTest extends Specification {
 
         then:
         instrumentation.throwableList.size() == 1
-        instrumentation.throwableList[0].getMessage() == "DF BANG!"
+        instrumentation.throwableList[0].getCause().getMessage() == "DF BANG!"
     }
 
     /**
@@ -169,9 +169,10 @@ class InstrumentationTest extends Specification {
             return new ExecutionStrategyInstrumentationContext() {
 
                 @Override
-                void onDispatched(CompletableFuture<ExecutionResult> result) {
+                CompletableFuture<ExecutionResult> onDispatched(CompletableFuture<ExecutionResult> result) {
                     System.out.println(String.format("t%s setting go signal on", Thread.currentThread().getId()))
                     goSignal.set(true)
+                    return result;
                 }
 
                 @Override

diff --git a/src/test/groovy/graphql/execution/instrumentation/TestingInstrumentContext.groovy b/src/test/groovy/graphql/execution/instrumentation/TestingInstrumentContext.groovy
@@ -25,10 +25,11 @@ class TestingInstrumentContext<T> implements InstrumentationContext<T> {
     }
 
     @Override
-    void onDispatched(CompletableFuture<T> result) {
+    CompletableFuture<T> onDispatched(CompletableFuture<T> result) {
         if (useOnDispatch) {
             this.executionList << "onDispatched:$op"
         }
+        result
     }
 
     @Override

diff --git a/...aphql/execution/instrumentation/dataloader/DataLoaderDispatcherInstrumentationTest.groovy b/...aphql/execution/instrumentation/dataloader/DataLoaderDispatcherInstrumentationTest.groovy
@@ -99,9 +99,9 @@ class DataLoaderDispatcherInstrumentationTest extends Specification {
         def dispatchedCalled = false
         def dataLoaderRegistry = new DataLoaderRegistry() {
             @Override
-            void dispatchAll() {
+            CompletableFuture<?> dispatch() {
                 dispatchedCalled = true
-                super.dispatchAll()
+                super.dispatch()
             }
         }
         def dataLoader = DataLoader.newDataLoader(new BatchLoader() {