Merge pull request graphql-java#3639 from felipe-gdr/defer-dataloader-integration

Defer dataloader integration

Author: andimarek

src/main/java/graphql/execution/DataLoaderDispatchStrategy.java

@@ -4,7 +4,6 @@
 import graphql.schema.DataFetcher;
 
 import java.util.List;
-import java.util.concurrent.CompletableFuture;
 
 @Internal
 public interface DataLoaderDispatchStrategy {
@@ -50,7 +49,7 @@ default DataFetcher<?> modifyDataFetcher(DataFetcher<?> dataFetcher) {
         return dataFetcher;
     }
 
-    default void deferredField(ExecutionContext executionContext, MergedField currentField) {
+    default void executeDeferredOnFieldValueInfo(FieldValueInfo fieldValueInfo, ExecutionStrategyParameters executionStrategyParameters) {
 
     }
 }

src/main/java/graphql/execution/Execution.java

@@ -14,6 +14,7 @@
 import graphql.execution.instrumentation.InstrumentationState;
 import graphql.execution.instrumentation.dataloader.FallbackDataLoaderDispatchStrategy;
 import graphql.execution.instrumentation.dataloader.PerLevelDataLoaderDispatchStrategy;
+import graphql.execution.instrumentation.dataloader.PerLevelDataLoaderDispatchStrategyWithDeferAlwaysDispatch;
 import graphql.execution.instrumentation.parameters.InstrumentationExecuteOperationParameters;
 import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters;
 import graphql.extensions.ExtensionsBuilder;
@@ -228,7 +229,14 @@ private DataLoaderDispatchStrategy createDataLoaderDispatchStrategy(ExecutionCon
             return DataLoaderDispatchStrategy.NO_OP;
         }
         if (executionStrategy instanceof AsyncExecutionStrategy) {
-            return new PerLevelDataLoaderDispatchStrategy(executionContext);
+            boolean deferEnabled = Optional.ofNullable(executionContext.getGraphQLContext())
+                    .map(graphqlContext -> graphqlContext.getBoolean(ExperimentalApi.ENABLE_INCREMENTAL_SUPPORT))
+                    .orElse(false);
+
+            // Dedicated strategy for defer support, for safety purposes.
+            return deferEnabled ?
+                    new PerLevelDataLoaderDispatchStrategyWithDeferAlwaysDispatch(executionContext) :
+                    new PerLevelDataLoaderDispatchStrategy(executionContext);
         } else {
             return new FallbackDataLoaderDispatchStrategy(executionContext);
         }

src/main/java/graphql/execution/ExecutionStrategy.java

@@ -318,7 +318,7 @@ Async.CombinedBuilder<FieldValueInfo> getAsyncFieldValueInfo(
     ) {
         MergedSelectionSet fields = parameters.getFields();
 
-        executionContext.getIncrementalCallState().enqueue(deferredExecutionSupport.createCalls());
+        executionContext.getIncrementalCallState().enqueue(deferredExecutionSupport.createCalls(parameters));
 
         // Only non-deferred fields should be considered for calculating the expected size of futures.
         Async.CombinedBuilder<FieldValueInfo> futures = Async

src/main/java/graphql/execution/incremental/DeferredExecutionSupport.java

@@ -44,7 +44,7 @@ public interface DeferredExecutionSupport {
 
     List<String> getNonDeferredFieldNames(List<String> allFieldNames);
 
-    Set<IncrementalCall<? extends IncrementalPayload>> createCalls();
+    Set<IncrementalCall<? extends IncrementalPayload>> createCalls(ExecutionStrategyParameters executionStrategyParameters);
 
     DeferredExecutionSupport NOOP = new DeferredExecutionSupport.NoOp();
 
@@ -105,19 +105,19 @@ public List<String> getNonDeferredFieldNames(List<String> allFieldNames) {
         }
 
         @Override
-        public Set<IncrementalCall<? extends IncrementalPayload>> createCalls() {
+        public Set<IncrementalCall<? extends IncrementalPayload>> createCalls(ExecutionStrategyParameters executionStrategyParameters) {
             return deferredExecutionToFields.keySet().stream()
-                    .map(this::createDeferredFragmentCall)
+                    .map(deferredExecution -> this.createDeferredFragmentCall(deferredExecution, executionStrategyParameters))
                     .collect(Collectors.toSet());
         }
 
-        private DeferredFragmentCall createDeferredFragmentCall(DeferredExecution deferredExecution) {
+        private DeferredFragmentCall createDeferredFragmentCall(DeferredExecution deferredExecution, ExecutionStrategyParameters executionStrategyParameters) {
             DeferredCallContext deferredCallContext = new DeferredCallContext();
 
             List<MergedField> mergedFields = deferredExecutionToFields.get(deferredExecution);
 
             List<Supplier<CompletableFuture<DeferredFragmentCall.FieldWithExecutionResult>>> calls = mergedFields.stream()
-                    .map(currentField -> this.createResultSupplier(currentField, deferredCallContext))
+                    .map(currentField -> this.createResultSupplier(currentField, deferredCallContext, executionStrategyParameters))
                     .collect(Collectors.toList());
 
             return new DeferredFragmentCall(
@@ -130,7 +130,8 @@ private DeferredFragmentCall createDeferredFragmentCall(DeferredExecution deferr
 
         private Supplier<CompletableFuture<DeferredFragmentCall.FieldWithExecutionResult>> createResultSupplier(
                 MergedField currentField,
-                DeferredCallContext deferredCallContext
+                DeferredCallContext deferredCallContext,
+                ExecutionStrategyParameters executionStrategyParameters
         ) {
             Map<String, MergedField> fields = new LinkedHashMap<>();
             fields.put(currentField.getResultKey(), currentField);
@@ -149,7 +150,6 @@ private Supplier<CompletableFuture<DeferredFragmentCall.FieldWithExecutionResult
 
             Instrumentation instrumentation = executionContext.getInstrumentation();
 
-            executionContext.getDataLoaderDispatcherStrategy().deferredField(executionContext, currentField);
             instrumentation.beginDeferredField(executionContext.getInstrumentationState());
 
             return dfCache.computeIfAbsent(
@@ -160,12 +160,14 @@ private Supplier<CompletableFuture<DeferredFragmentCall.FieldWithExecutionResult
                                 CompletableFuture<FieldValueInfo> fieldValueResult = resolveFieldWithInfoFn
                                         .apply(executionContext, callParameters);
 
-                                // Create a reference to the CompletableFuture that resolves an ExecutionResult
-                                // so we can pass it to the Instrumentation "onDispatched" callback.
                                 CompletableFuture<ExecutionResult> executionResultCF = fieldValueResult
-                                        .thenCompose(fvi -> fvi
-                                                .getFieldValueFuture()
-                                                .thenApply(fv -> ExecutionResultImpl.newExecutionResult().data(fv).build())
+                                        .thenCompose(fvi -> {
+                                                    executionContext.getDataLoaderDispatcherStrategy().executeDeferredOnFieldValueInfo(fvi, executionStrategyParameters);
+
+                                                    return fvi
+                                                            .getFieldValueFuture()
+                                                            .thenApply(fv -> ExecutionResultImpl.newExecutionResult().data(fv).build());
+                                                }
                                         );
 
                                 return executionResultCF
@@ -199,7 +201,7 @@ public List<String> getNonDeferredFieldNames(List<String> allFieldNames) {
         }
 
         @Override
-        public Set<IncrementalCall<? extends IncrementalPayload>> createCalls() {
+        public Set<IncrementalCall<? extends IncrementalPayload>> createCalls(ExecutionStrategyParameters executionStrategyParameters) {
             return Collections.emptySet();
         }
     }

src/main/java/graphql/execution/instrumentation/dataloader/DataLoaderDispatcherInstrumentation.java

@@ -98,7 +98,7 @@ public PerLevelDataLoaderDispatchStrategy(ExecutionContext executionContext) {
     }
 
     @Override
-    public void deferredField(ExecutionContext executionContext, MergedField currentField) {
+    public void executeDeferredOnFieldValueInfo(FieldValueInfo fieldValueInfo, ExecutionStrategyParameters executionStrategyParameters) {
         throw new UnsupportedOperationException("Data Loaders cannot be used to resolve deferred fields");
     }