WIP: hacky defer is working, some new tests for basic stuff, but a lo…

…t still missing
andimarek · Feb 13, 2024 · Dec 20, 2023 · Dec 20, 2023 · Dec 20, 2023 · Jan 2, 2024
commit e86286e2fd0839abb2ce3f7d174bfae18fbaf357
diff --git a/src/main/java/graphql/ExecutionInput.java b/src/main/java/graphql/ExecutionInput.java
@@ -29,6 +29,7 @@ public class ExecutionInput {
     private final DataLoaderRegistry dataLoaderRegistry;
     private final ExecutionId executionId;
     private final Locale locale;
+    private final boolean incrementalSupport;
 
 
     @Internal
@@ -44,6 +45,7 @@ private ExecutionInput(Builder builder) {
         this.locale = builder.locale != null ? builder.locale : Locale.getDefault(); // always have a locale in place
         this.localContext = builder.localContext;
         this.extensions = builder.extensions;
+        this.incrementalSupport = builder.incrementalSupport;
     }
 
     /**
@@ -139,6 +141,16 @@ public Map<String, Object> getExtensions() {
         return extensions;
     }
 
+    /**
+     * TODO: Javadoc
+     * @return
+     */
+    @ExperimentalApi
+    public boolean isIncrementalSupport() {
+        return incrementalSupport;
+    }
+
+
     /**
      * This helps you transform the current ExecutionInput object into another one by starting a builder with all
      * the current values and allows you to transform it how you want.
@@ -159,7 +171,8 @@ public ExecutionInput transform(Consumer<Builder> builderConsumer) {
                 .variables(this.rawVariables.toMap())
                 .extensions(this.extensions)
                 .executionId(this.executionId)
-                .locale(this.locale);
+                .locale(this.locale)
+                .incrementalSupport(this.incrementalSupport);
 
         builderConsumer.accept(builder);
 
@@ -216,6 +229,7 @@ public static class Builder {
         private DataLoaderRegistry dataLoaderRegistry = DataLoaderDispatcherInstrumentationState.EMPTY_DATALOADER_REGISTRY;
         private Locale locale = Locale.getDefault();
         private ExecutionId executionId;
+        public boolean incrementalSupport = false;
 
         public Builder query(String query) {
             this.query = assertNotNull(query, () -> "query can't be null");
@@ -379,8 +393,19 @@ public Builder dataLoaderRegistry(DataLoaderRegistry dataLoaderRegistry) {
             return this;
         }
 
+        /**
+         * TODO: Javadoc
+         * @param incrementalSupport
+         * @return
+         */
+        @ExperimentalApi
+        public Builder incrementalSupport(boolean incrementalSupport) {
+            this.incrementalSupport = incrementalSupport;
+            return this;
+        }
+
         public ExecutionInput build() {
             return new ExecutionInput(this);
         }
     }
-}
+}
diff --git a/src/main/java/graphql/execution/AsyncExecutionStrategy.java b/src/main/java/graphql/execution/AsyncExecutionStrategy.java
@@ -1,25 +1,31 @@
 package graphql.execution;
 
+import com.google.common.collect.ImmutableListMultimap;
+import com.google.common.collect.ImmutableSet;
 import graphql.ExecutionResult;
 import graphql.PublicApi;
-import graphql.execution.defer.DeferSupport;
+import graphql.execution.defer.DeferExecutionSupport;
 import graphql.execution.defer.DeferredCall;
 import graphql.execution.defer.DeferredErrorSupport;
-import graphql.execution.instrumentation.DeferredFieldInstrumentationContext;
+import graphql.execution.incremental.DeferExecution;
 import graphql.execution.instrumentation.ExecutionStrategyInstrumentationContext;
 import graphql.execution.instrumentation.Instrumentation;
 import graphql.execution.instrumentation.InstrumentationContext;
 import graphql.execution.instrumentation.parameters.InstrumentationDeferredFieldParameters;
 import graphql.execution.instrumentation.parameters.InstrumentationExecutionStrategyParameters;
 import graphql.schema.GraphQLFieldDefinition;
 import graphql.util.FpKit;
+import graphql.util.Pair;
 
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.BiConsumer;
+import java.util.function.BiFunction;
 import java.util.function.Supplier;
+import java.util.stream.Collectors;
 
 import static graphql.execution.MergedSelectionSet.newMergedSelectionSet;
 
@@ -56,7 +62,20 @@ public CompletableFuture<ExecutionResult> execute(ExecutionContext executionCont
 
         MergedSelectionSet fields = parameters.getFields();
         List<String> fieldNames = fields.getKeys();
-        Async.CombinedBuilder<FieldValueInfo> futures = Async.ofExpectedSize(fields.size());
+
+//        if(true /* check if incremental support is enabled*/) {
+            SomethingDefer somethingDefer = new SomethingDefer(
+                    fields,
+                    parameters,
+                    executionContext,
+                    this::resolveFieldWithInfo
+            );
+
+            executionContext.getDeferSupport().enqueue(somethingDefer.createCalls());
+//        }
+
+        Async.CombinedBuilder<FieldValueInfo> futures = Async.ofExpectedSize(fields.size() - somethingDefer.deferredFields.size());
+
         for (String fieldName : fieldNames) {
             MergedField currentField = fields.getSubField(fieldName);
 
@@ -66,13 +85,14 @@ public CompletableFuture<ExecutionResult> execute(ExecutionContext executionCont
 
             CompletableFuture<FieldValueInfo> future;
 
-            if (isDeferred(executionContext, newParameters, currentField)) {
+            if (somethingDefer.isDeferredField(currentField)) {
                 executionStrategyCtx.onDeferredField(currentField);
-                future = resolveFieldWithInfoToNull(executionContext, newParameters);
+//                future = resolveFieldWithInfoToNull(executionContext, newParameters);
             } else {
                 future = resolveFieldWithInfo(executionContext, newParameters);
+                futures.add(future);
             }
-            futures.add(future);
+
         }
         CompletableFuture<ExecutionResult> overallResult = new CompletableFuture<>();
         executionStrategyCtx.onDispatched(overallResult);
@@ -104,8 +124,9 @@ public CompletableFuture<ExecutionResult> execute(ExecutionContext executionCont
     }
 
     private boolean isDeferred(ExecutionContext executionContext, ExecutionStrategyParameters parameters, MergedField currentField) {
-        DeferSupport deferSupport = executionContext.getDeferSupport();
-        if (deferSupport.checkForDeferDirective(currentField, executionContext)) {
+        DeferExecutionSupport deferSupport = executionContext.getDeferSupport();
+
+        if (currentField.getDeferExecutions() != null && !currentField.getDeferExecutions().isEmpty()) {
             DeferredErrorSupport errorSupport = new DeferredErrorSupport();
 
             // with a deferred field we are really resetting where we execute from, that is from this current field onwards
@@ -122,8 +143,8 @@ private boolean isDeferred(ExecutionContext executionContext, ExecutionStrategyP
                     }
             );
 
-            DeferredCall call = new DeferredCall(parameters.getPath(), deferredExecutionResult(executionContext, callParameters), errorSupport);
-            deferSupport.enqueue(call);
+//            DeferredCall call = new DeferredCall(null /* TODO extract label somehow*/, parameters.getPath(), deferredExecutionResult(executionContext, callParameters), errorSupport);
+//            deferSupport.enqueue(call);
             return true;
         }
         return false;
@@ -160,4 +181,112 @@ private Supplier<CompletableFuture<ExecutionResult>> deferredExecutionResult(Exe
             return result;
         };
     }
+
+
+    private static class SomethingDefer {
+        private final ImmutableListMultimap<DeferExecution, MergedField> deferExecutionToFields;
+        private final ImmutableSet<MergedField> deferredFields;
+        private final ExecutionStrategyParameters parameters;
+        private final ExecutionContext executionContext;
+        private final BiFunction<ExecutionContext, ExecutionStrategyParameters, CompletableFuture<FieldValueInfo>> resolveFieldWithInfoFn;
+
+        private SomethingDefer(
+                MergedSelectionSet mergedSelectionSet,
+                ExecutionStrategyParameters parameters,
+                ExecutionContext executionContext, BiFunction<ExecutionContext, ExecutionStrategyParameters, CompletableFuture<FieldValueInfo>> resolveFieldWithInfoFn
+        ) {
+            this.executionContext = executionContext;
+            this.resolveFieldWithInfoFn = resolveFieldWithInfoFn;
+            ImmutableListMultimap.Builder<DeferExecution, MergedField> deferExecutionToFieldsBuilder = ImmutableListMultimap.builder();
+            ImmutableSet.Builder<MergedField> deferredFieldsBuilder = ImmutableSet.builder();
+
+            mergedSelectionSet.getSubFields().values().forEach(mergedField -> {
+                mergedField.getDeferExecutions().forEach(de -> {
+                    deferExecutionToFieldsBuilder.put(de, mergedField);
+                    deferredFieldsBuilder.add(mergedField);
+                });
+            });
+
+            this.deferExecutionToFields = deferExecutionToFieldsBuilder.build();
+            this.deferredFields = deferredFieldsBuilder.build();
+            this.parameters = parameters;
+        }
+
+        private boolean isDeferredField(MergedField mergedField) {
+            return deferredFields.contains(mergedField);
+        }
+
+        private Set<DeferredCall> createCalls() {
+            return deferExecutionToFields.keySet().stream().map(deferExecution -> {
+                        DeferredErrorSupport errorSupport = new DeferredErrorSupport();
+
+                        List<MergedField> mergedFields = deferExecutionToFields.get(deferExecution);
+
+                        List<Supplier<CompletableFuture<DeferredCall.FieldWithExecutionResult>>> collect = mergedFields.stream()
+                                .map(currentField -> {
+                                    Map<String, MergedField> fields = new LinkedHashMap<>();
+                                    fields.put(currentField.getName(), currentField);
+
+                                    ExecutionStrategyParameters callParameters = parameters.transform(builder ->
+                                            {
+                                                MergedSelectionSet mergedSelectionSet = newMergedSelectionSet().subFields(fields).build();
+                                                builder.deferredErrorSupport(errorSupport)
+                                                        .field(currentField)
+                                                        .fields(mergedSelectionSet)
+                                                        .parent(null); // this is a break in the parent -> child chain - its a new start effectively
+                                            }
+                                    );
+
+
+                                    return (Supplier<CompletableFuture<DeferredCall.FieldWithExecutionResult>>) () -> resolveFieldWithInfoFn
+                                            .apply(executionContext, callParameters)
+                                            .thenCompose(FieldValueInfo::getFieldValue)
+                                            .thenApply(executionResult -> new DeferredCall.FieldWithExecutionResult(currentField.getName(), executionResult));
+
+                                })
+                                .collect(Collectors.toList());
+
+                        // with a deferred field we are really resetting where we execute from, that is from this current field onwards
+                        return new DeferredCall(
+                                deferExecution.getLabel(),
+                                this.parameters.getPath(),
+                                collect,
+                                errorSupport
+                        );
+                    })
+                    .collect(Collectors.toSet());
+        }
+
+        @SuppressWarnings("FutureReturnValueIgnored")
+        private Supplier<CompletableFuture<ExecutionResult>> deferredExecutionResult(ExecutionContext executionContext, ExecutionStrategyParameters parameters) {
+            return () -> {
+//                GraphQLFieldDefinition fieldDef = getFieldDef(executionContext, parameters, parameters.getField().getSingleField());
+                // TODO: freis: This is suddenly not needed anymore
+//            GraphQLObjectType fieldContainer = (GraphQLObjectType) parameters.getExecutionStepInfo().getUnwrappedNonNullType();
+
+//                Instrumentation instrumentation = executionContext.getInstrumentation();
+
+//                Supplier<ExecutionStepInfo> executionStepInfo = FpKit.intraThreadMemoize(() -> createExecutionStepInfo(executionContext, parameters, fieldDef, null));
+
+//                InstrumentationContext<ExecutionResult> fieldCtx = instrumentation.beginDeferredField(
+//                        new InstrumentationDeferredFieldParameters(executionContext, executionStepInfo, parameters),
+//                        executionContext.getInstrumentationState()
+//                );
+
+                CompletableFuture<ExecutionResult> result = new CompletableFuture<>();
+//                fieldCtx.onDispatched(result);
+                CompletableFuture<FieldValueInfo> fieldValueInfoFuture = resolveFieldWithInfoFn.apply(executionContext, parameters);
+
+                fieldValueInfoFuture.whenComplete((fieldValueInfo, throwable) -> {
+                    // TODO:
+//                fieldCtx.onFieldValueInfo(fieldValueInfo);
+
+                    CompletableFuture<ExecutionResult> execResultFuture = fieldValueInfo.getFieldValue();
+//                    execResultFuture = execResultFuture.whenComplete(fieldCtx::onCompleted);
+                    Async.copyResults(execResultFuture, result);
+                });
+                return result;
+            };
+        }
+    }
 }
diff --git a/src/main/java/graphql/execution/Execution.java b/src/main/java/graphql/execution/Execution.java
@@ -1,21 +1,21 @@
 package graphql.execution;
 
 
-import graphql.DeferredExecutionResult;
 import graphql.ExecutionInput;
 import graphql.ExecutionResult;
 import graphql.ExecutionResultImpl;
-import graphql.GraphQL;
 import graphql.GraphQLContext;
 import graphql.GraphQLError;
 import graphql.Internal;
-import graphql.execution.defer.DeferSupport;
+import graphql.execution.defer.DeferExecutionSupport;
 import graphql.execution.instrumentation.Instrumentation;
 import graphql.execution.instrumentation.InstrumentationContext;
 import graphql.execution.instrumentation.InstrumentationState;
 import graphql.execution.instrumentation.parameters.InstrumentationExecuteOperationParameters;
 import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters;
 import graphql.extensions.ExtensionsBuilder;
+import graphql.incremental.DelayedIncrementalExecutionResult;
+import graphql.incremental.IncrementalExecutionResultImpl;
 import graphql.language.Document;
 import graphql.language.FragmentDefinition;
 import graphql.language.NodeUtil;
@@ -186,18 +186,26 @@ private CompletableFuture<ExecutionResult> executeOperation(ExecutionContext exe
     }
 
     /*
-     * Adds the deferred publisher if its needed at the end of the query.  This is also a good time for the deferred code to start running
+     * Adds the deferred publisher if it's needed at the end of the query.  This is also a good time for the deferred code to start running
      */
     private CompletableFuture<ExecutionResult> deferSupport(ExecutionContext executionContext, CompletableFuture<ExecutionResult> result) {
         return result.thenApply(er -> {
-            DeferSupport deferSupport = executionContext.getDeferSupport();
+            DeferExecutionSupport deferSupport = executionContext.getDeferSupport();
             if (deferSupport.isDeferDetected()) {
                 // we start the rest of the query now to maximize throughput.  We have the initial important results
                 // and now we can start the rest of the calls as early as possible (even before some one subscribes)
-                Publisher<DeferredExecutionResult> publisher = deferSupport.startDeferredCalls();
-                return ExecutionResultImpl.newExecutionResult().from(er)
-                        .addExtension(GraphQL.DEFERRED_RESULTS, publisher)
+                Publisher<DelayedIncrementalExecutionResult> publisher = deferSupport.startDeferredCalls();
+
+                return IncrementalExecutionResultImpl.fromExecutionResult(er)
+                        // TODO: would `hasNext` ever be false?
+                        .hasNext(true)
+                        .incrementalItemPublisher(publisher)
                         .build();
+//
+//
+//                return ExecutionResultImpl.newExecutionResult().from(er)
+//                        .addExtension(GraphQL.DEFERRED_RESULTS, publisher)
+//                        .build();
             }
             return er;
         });

diff --git a/src/main/java/graphql/execution/ExecutionContext.java b/src/main/java/graphql/execution/ExecutionContext.java
@@ -8,7 +8,7 @@
 import graphql.GraphQLError;
 import graphql.PublicApi;
 import graphql.collect.ImmutableKit;
-import graphql.execution.defer.DeferSupport;
+import graphql.execution.defer.DeferExecutionSupport;
 import graphql.execution.instrumentation.Instrumentation;
 import graphql.execution.instrumentation.InstrumentationState;
 import graphql.language.Document;
@@ -54,7 +54,7 @@ public class ExecutionContext {
     private final Set<ResultPath> errorPaths = new HashSet<>();
     private final DataLoaderRegistry dataLoaderRegistry;
     private final Locale locale;
-    private final DeferSupport deferSupport = new DeferSupport();
+    private final DeferExecutionSupport deferSupport = new DeferExecutionSupport();
     private final ValueUnboxer valueUnboxer;
     private final ExecutionInput executionInput;
     private final Supplier<ExecutableNormalizedOperation> queryTree;
@@ -257,7 +257,7 @@ public ExecutionStrategy getSubscriptionStrategy() {
         return subscriptionStrategy;
     }
 
-    public DeferSupport getDeferSupport() {
+    public DeferExecutionSupport getDeferSupport() {
         return deferSupport;
     }
 

diff --git a/src/main/java/graphql/execution/ExecutionStrategy.java b/src/main/java/graphql/execution/ExecutionStrategy.java
@@ -696,7 +696,11 @@ protected CompletableFuture<ExecutionResult> completeValueForObject(ExecutionCon
                 .variables(executionContext.getCoercedVariables().toMap())
                 .build();
 
-        MergedSelectionSet subFields = fieldCollector.collectFields(collectorParameters, parameters.getField());
+        MergedSelectionSet subFields = fieldCollector.collectFields(
+                collectorParameters,
+                parameters.getField(),
+                executionContext.getExecutionInput().isIncrementalSupport()
+        );
 
         ExecutionStepInfo newExecutionStepInfo = executionStepInfo.changeTypeWithPreservedNonNull(resolvedObjectType);
         NonNullableFieldValidator nonNullableFieldValidator = new NonNullableFieldValidator(executionContext, newExecutionStepInfo);