Merge pull request graphql-java#3447 from graphql-java/native-dispatc…

…h-strategy Native DataLoader dispatch strategy without Instrumentation
felipe-gdr · Feb 19, 2024 · 4cfbb35 · 4cfbb35
2 parents 06b6844 + 6bf834b
commit 4cfbb35
Show file tree

Hide file tree

Showing 35 changed files with 492 additions and 844 deletions.
diff --git a/src/main/java/graphql/ExecutionInput.java b/src/main/java/graphql/ExecutionInput.java
@@ -3,7 +3,6 @@
 import graphql.collect.ImmutableKit;
 import graphql.execution.ExecutionId;
 import graphql.execution.RawVariables;
-import graphql.execution.instrumentation.dataloader.DataLoaderDispatcherInstrumentationState;
 import org.dataloader.DataLoaderRegistry;
 
 import java.util.Locale;
@@ -12,6 +11,7 @@
 import java.util.function.UnaryOperator;
 
 import static graphql.Assert.assertNotNull;
+import static graphql.execution.instrumentation.dataloader.EmptyDataLoaderRegistryInstance.EMPTY_DATALOADER_REGISTRY;
 
 /**
  * This represents the series of values that can be input on a graphql query execution
@@ -213,7 +213,7 @@ public static class Builder {
         // this is important - it allows code to later known if we never really set a dataloader and hence it can optimize
         // dataloader field tracking away.
         //
-        private DataLoaderRegistry dataLoaderRegistry = DataLoaderDispatcherInstrumentationState.EMPTY_DATALOADER_REGISTRY;
+        private DataLoaderRegistry dataLoaderRegistry = EMPTY_DATALOADER_REGISTRY;
         private Locale locale = Locale.getDefault();
         private ExecutionId executionId;
 

diff --git a/src/main/java/graphql/GraphQL.java b/src/main/java/graphql/GraphQL.java
@@ -12,14 +12,11 @@
 import graphql.execution.SimpleDataFetcherExceptionHandler;
 import graphql.execution.SubscriptionExecutionStrategy;
 import graphql.execution.ValueUnboxer;
-import graphql.execution.instrumentation.ChainedInstrumentation;
 import graphql.execution.instrumentation.DocumentAndVariables;
 import graphql.execution.instrumentation.Instrumentation;
 import graphql.execution.instrumentation.InstrumentationContext;
 import graphql.execution.instrumentation.InstrumentationState;
-import graphql.execution.instrumentation.NoContextChainedInstrumentation;
 import graphql.execution.instrumentation.SimplePerformantInstrumentation;
-import graphql.execution.instrumentation.dataloader.DataLoaderDispatcherInstrumentation;
 import graphql.execution.instrumentation.parameters.InstrumentationCreateStateParameters;
 import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters;
 import graphql.execution.instrumentation.parameters.InstrumentationValidationParameters;
@@ -30,7 +27,6 @@
 import graphql.schema.GraphQLSchema;
 import graphql.validation.ValidationError;
 
-import java.util.ArrayList;
 import java.util.List;
 import java.util.Locale;
 import java.util.Optional;
@@ -96,6 +92,7 @@ public class GraphQL {
     private final Instrumentation instrumentation;
     private final PreparsedDocumentProvider preparsedDocumentProvider;
     private final ValueUnboxer valueUnboxer;
+    private final boolean doNotAutomaticallyDispatchDataLoader;
 
 
     private GraphQL(Builder builder) {
@@ -107,6 +104,7 @@ private GraphQL(Builder builder) {
         this.instrumentation = assertNotNull(builder.instrumentation, () -> "instrumentation must not be null");
         this.preparsedDocumentProvider = assertNotNull(builder.preparsedDocumentProvider, () -> "preparsedDocumentProvider must be non null");
         this.valueUnboxer = assertNotNull(builder.valueUnboxer, () -> "valueUnboxer must not be null");
+        this.doNotAutomaticallyDispatchDataLoader = builder.doNotAutomaticallyDispatchDataLoader;
     }
 
     /**
@@ -151,6 +149,10 @@ public Instrumentation getInstrumentation() {
         return instrumentation;
     }
 
+    public boolean isDoNotAutomaticallyDispatchDataLoader() {
+        return doNotAutomaticallyDispatchDataLoader;
+    }
+
     /**
      * @return the PreparsedDocumentProvider for this {@link GraphQL} instance
      */
@@ -209,7 +211,7 @@ public static class Builder {
         private ExecutionIdProvider idProvider = DEFAULT_EXECUTION_ID_PROVIDER;
         private Instrumentation instrumentation = null; // deliberate default here
         private PreparsedDocumentProvider preparsedDocumentProvider = NoOpPreparsedDocumentProvider.INSTANCE;
-        private boolean doNotAddDefaultInstrumentations = false;
+        private boolean doNotAutomaticallyDispatchDataLoader = false;
         private ValueUnboxer valueUnboxer = ValueUnboxer.DEFAULT;
 
 
@@ -265,20 +267,15 @@ public Builder executionIdProvider(ExecutionIdProvider executionIdProvider) {
             return this;
         }
 
+
         /**
-         * For performance reasons you can opt into situation where the default instrumentations (such
-         * as {@link graphql.execution.instrumentation.dataloader.DataLoaderDispatcherInstrumentation} will not be
-         * automatically added into the graphql instance.
-         * <p>
-         * For most situations this is not needed unless you are really pushing the boundaries of performance
-         * <p>
-         * By default a certain graphql instrumentations will be added to the mix to more easily enable certain functionality.  This
-         * allows you to stop this behavior
+         * Deactivates the automatic dispatching of DataLoaders.
+         * If deactivated the user is responsible for dispatching the DataLoaders manually.
          *
          * @return this builder
          */
-        public Builder doNotAddDefaultInstrumentations() {
-            this.doNotAddDefaultInstrumentations = true;
+        public Builder doNotAutomaticallyDispatchDataLoader() {
+            this.doNotAutomaticallyDispatchDataLoader = true;
             return this;
         }
 
@@ -299,7 +296,9 @@ public GraphQL build() {
                 this.subscriptionExecutionStrategy = new SubscriptionExecutionStrategy(this.defaultExceptionHandler);
             }
 
-            this.instrumentation = checkInstrumentationDefaultState(this.instrumentation, this.doNotAddDefaultInstrumentations);
+            if (instrumentation == null) {
+                this.instrumentation = SimplePerformantInstrumentation.INSTANCE;
+            }
             return new GraphQL(this);
         }
     }
@@ -540,42 +539,16 @@ private List<ValidationError> validate(ExecutionInput executionInput, Document d
         return validationErrors;
     }
 
-    private CompletableFuture<ExecutionResult> execute(ExecutionInput executionInput, Document document, GraphQLSchema graphQLSchema, InstrumentationState instrumentationState) {
+    private CompletableFuture<ExecutionResult> execute(ExecutionInput executionInput,
+                                                       Document document,
+                                                       GraphQLSchema graphQLSchema,
+                                                       InstrumentationState instrumentationState
+    ) {
 
-        Execution execution = new Execution(queryStrategy, mutationStrategy, subscriptionStrategy, instrumentation, valueUnboxer);
+        Execution execution = new Execution(queryStrategy, mutationStrategy, subscriptionStrategy, instrumentation, valueUnboxer, doNotAutomaticallyDispatchDataLoader);
         ExecutionId executionId = executionInput.getExecutionId();
 
         return execution.execute(document, graphQLSchema, executionId, executionInput, instrumentationState);
     }
 
-    private static Instrumentation checkInstrumentationDefaultState(Instrumentation instrumentation, boolean doNotAddDefaultInstrumentations) {
-        if (doNotAddDefaultInstrumentations) {
-            return instrumentation == null ? SimplePerformantInstrumentation.INSTANCE : instrumentation;
-        }
-        if (instrumentation instanceof DataLoaderDispatcherInstrumentation) {
-            return instrumentation;
-        }
-        if (instrumentation instanceof NoContextChainedInstrumentation) {
-            return instrumentation;
-        }
-        if (instrumentation == null) {
-            return new DataLoaderDispatcherInstrumentation();
-        }
-
-        //
-        // if we don't have a DataLoaderDispatcherInstrumentation in play, we add one.  We want DataLoader to be 1st class in graphql without requiring
-        // people to remember to wire it in.  Later we may decide to have more default instrumentations but for now it's just the one
-        //
-        List<Instrumentation> instrumentationList = new ArrayList<>();
-        if (instrumentation instanceof ChainedInstrumentation) {
-            instrumentationList.addAll(((ChainedInstrumentation) instrumentation).getInstrumentations());
-        } else {
-            instrumentationList.add(instrumentation);
-        }
-        boolean containsDLInstrumentation = instrumentationList.stream().anyMatch(instr -> instr instanceof DataLoaderDispatcherInstrumentation);
-        if (!containsDLInstrumentation) {
-            instrumentationList.add(new DataLoaderDispatcherInstrumentation());
-        }
-        return new ChainedInstrumentation(instrumentationList);
-    }
 }
diff --git a/src/main/java/graphql/execution/AsyncExecutionStrategy.java b/src/main/java/graphql/execution/AsyncExecutionStrategy.java
@@ -36,6 +36,8 @@ public AsyncExecutionStrategy(DataFetcherExceptionHandler exceptionHandler) {
     @Override
     @SuppressWarnings("FutureReturnValueIgnored")
     public CompletableFuture<ExecutionResult> execute(ExecutionContext executionContext, ExecutionStrategyParameters parameters) throws NonNullableFieldWasNullException {
+        DataLoaderDispatchStrategy dataLoaderDispatcherStrategy = executionContext.getDataLoaderDispatcherStrategy();
+        dataLoaderDispatcherStrategy.executionStrategy(executionContext, parameters);
         Instrumentation instrumentation = executionContext.getInstrumentation();
         InstrumentationExecutionStrategyParameters instrumentationParameters = new InstrumentationExecutionStrategyParameters(executionContext, parameters);
 
@@ -63,12 +65,14 @@ public CompletableFuture<ExecutionResult> execute(ExecutionContext executionCont
             for (FieldValueInfo completeValueInfo : completeValueInfos) {
                 fieldValuesFutures.add(completeValueInfo.getFieldValueFuture());
             }
+            dataLoaderDispatcherStrategy.executionStrategyOnFieldValuesInfo(completeValueInfos, parameters);
             executionStrategyCtx.onFieldValuesInfo(completeValueInfos);
             fieldValuesFutures.await().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
             // the execution does not hang.
+            dataLoaderDispatcherStrategy.executionStrategyOnFieldValuesException(ex, parameters);
             executionStrategyCtx.onFieldValuesException();
             overallResult.completeExceptionally(ex);
             return null;

diff --git a/src/main/java/graphql/execution/AsyncSerialExecutionStrategy.java b/src/main/java/graphql/execution/AsyncSerialExecutionStrategy.java
@@ -30,6 +30,7 @@ public AsyncSerialExecutionStrategy(DataFetcherExceptionHandler exceptionHandler
     @Override
     @SuppressWarnings({"TypeParameterUnusedInFormals", "FutureReturnValueIgnored"})
     public CompletableFuture<ExecutionResult> execute(ExecutionContext executionContext, ExecutionStrategyParameters parameters) throws NonNullableFieldWasNullException {
+        executionContext.getDataLoaderDispatcherStrategy().executionStrategy(executionContext, parameters);
 
         Instrumentation instrumentation = executionContext.getInstrumentation();
         InstrumentationExecutionStrategyParameters instrumentationParameters = new InstrumentationExecutionStrategyParameters(executionContext, parameters);

diff --git a/src/main/java/graphql/execution/DataLoaderDispatchStrategy.java b/src/main/java/graphql/execution/DataLoaderDispatchStrategy.java
@@ -0,0 +1,56 @@
+package graphql.execution;
+
+import graphql.Internal;
+import graphql.schema.DataFetcher;
+
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+
+@Internal
+public interface DataLoaderDispatchStrategy {
+
+    DataLoaderDispatchStrategy NO_OP = new DataLoaderDispatchStrategy() {
+    };
+
+
+    default void executionStrategy(ExecutionContext executionContext, ExecutionStrategyParameters parameters) {
+
+    }
+
+    default void executionStrategyOnFieldValuesInfo(List<FieldValueInfo> fieldValueInfoList, ExecutionStrategyParameters parameters) {
+
+    }
+
+    default void executionStrategyOnFieldValuesException(Throwable t, ExecutionStrategyParameters parameters) {
+
+    }
+
+
+    default void executeObject(ExecutionContext executionContext, ExecutionStrategyParameters executionStrategyParameters) {
+
+    }
+
+    default void executeObjectOnFieldValuesInfo(List<FieldValueInfo> fieldValueInfoList, ExecutionStrategyParameters parameters) {
+
+    }
+
+    default void executeObjectOnFieldValuesException(Throwable t, ExecutionStrategyParameters parameters) {
+
+    }
+
+    default void fieldFetched(ExecutionContext executionContext,
+                              ExecutionStrategyParameters executionStrategyParameters,
+                              DataFetcher<?> dataFetcher,
+                              CompletableFuture<Object> fetchedValue) {
+
+    }
+
+
+    default DataFetcher<?> modifyDataFetcher(DataFetcher<?> dataFetcher) {
+        return dataFetcher;
+    }
+
+    default void deferredField(ExecutionContext executionContext, MergedField currentField) {
+
+    }
+}
diff --git a/src/main/java/graphql/execution/Execution.java b/src/main/java/graphql/execution/Execution.java
@@ -12,6 +12,8 @@
 import graphql.execution.instrumentation.Instrumentation;
 import graphql.execution.instrumentation.InstrumentationContext;
 import graphql.execution.instrumentation.InstrumentationState;
+import graphql.execution.instrumentation.dataloader.FallbackDataLoaderDispatchStrategy;
+import graphql.execution.instrumentation.dataloader.PerLevelDataLoaderDispatchStrategy;
 import graphql.execution.instrumentation.parameters.InstrumentationExecuteOperationParameters;
 import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters;
 import graphql.extensions.ExtensionsBuilder;
@@ -37,6 +39,7 @@
 import static graphql.execution.ExecutionStepInfo.newExecutionStepInfo;
 import static graphql.execution.ExecutionStrategyParameters.newParameters;
 import static graphql.execution.instrumentation.SimpleInstrumentationContext.nonNullCtx;
+import static graphql.execution.instrumentation.dataloader.EmptyDataLoaderRegistryInstance.EMPTY_DATALOADER_REGISTRY;
 import static java.util.concurrent.CompletableFuture.completedFuture;
 
 @Internal
@@ -47,13 +50,20 @@ public class Execution {
     private final ExecutionStrategy subscriptionStrategy;
     private final Instrumentation instrumentation;
     private final ValueUnboxer valueUnboxer;
-
-    public Execution(ExecutionStrategy queryStrategy, ExecutionStrategy mutationStrategy, ExecutionStrategy subscriptionStrategy, Instrumentation instrumentation, ValueUnboxer valueUnboxer) {
+    private final boolean doNotAutomaticallyDispatchDataLoader;
+
+    public Execution(ExecutionStrategy queryStrategy,
+                     ExecutionStrategy mutationStrategy,
+                     ExecutionStrategy subscriptionStrategy,
+                     Instrumentation instrumentation,
+                     ValueUnboxer valueUnboxer,
+                     boolean doNotAutomaticallyDispatchDataLoader) {
         this.queryStrategy = queryStrategy != null ? queryStrategy : new AsyncExecutionStrategy();
         this.mutationStrategy = mutationStrategy != null ? mutationStrategy : new AsyncSerialExecutionStrategy();
         this.subscriptionStrategy = subscriptionStrategy != null ? subscriptionStrategy : new AsyncExecutionStrategy();
         this.instrumentation = instrumentation;
         this.valueUnboxer = valueUnboxer;
+        this.doNotAutomaticallyDispatchDataLoader = doNotAutomaticallyDispatchDataLoader;
     }
 
     public CompletableFuture<ExecutionResult> execute(Document document, GraphQLSchema graphQLSchema, ExecutionId executionId, ExecutionInput executionInput, InstrumentationState instrumentationState) {
@@ -160,9 +170,12 @@ private CompletableFuture<ExecutionResult> executeOperation(ExecutionContext exe
                 .path(path)
                 .build();
 
+
         CompletableFuture<ExecutionResult> result;
         try {
             ExecutionStrategy executionStrategy = executionContext.getStrategy(operation);
+            DataLoaderDispatchStrategy dataLoaderDispatchStrategy = createDataLoaderDispatchStrategy(executionContext, executionStrategy);
+            executionContext.setDataLoaderDispatcherStrategy(dataLoaderDispatchStrategy);
             result = executionStrategy.execute(executionContext, parameters);
         } catch (NonNullableFieldWasNullException e) {
             // this means it was non-null types all the way from an offending non-null type
@@ -209,6 +222,18 @@ private CompletableFuture<ExecutionResult> incrementalSupport(ExecutionContext e
         });
     }
 
+    private DataLoaderDispatchStrategy createDataLoaderDispatchStrategy(ExecutionContext executionContext, ExecutionStrategy executionStrategy) {
+        if (executionContext.getDataLoaderRegistry() == EMPTY_DATALOADER_REGISTRY || doNotAutomaticallyDispatchDataLoader) {
+            return DataLoaderDispatchStrategy.NO_OP;
+        }
+        if (executionStrategy instanceof AsyncExecutionStrategy) {
+            return new PerLevelDataLoaderDispatchStrategy(executionContext);
+        } else {
+            return new FallbackDataLoaderDispatchStrategy(executionContext);
+        }
+    }
+
+
     private void addExtensionsBuilderNotPresent(GraphQLContext graphQLContext) {
         Object builder = graphQLContext.get(ExtensionsBuilder.class);
         if (builder == null) {

diff --git a/src/main/java/graphql/execution/ExecutionContext.java b/src/main/java/graphql/execution/ExecutionContext.java
@@ -6,6 +6,7 @@
 import graphql.ExecutionInput;
 import graphql.GraphQLContext;
 import graphql.GraphQLError;
+import graphql.Internal;
 import graphql.PublicApi;
 import graphql.collect.ImmutableKit;
 import graphql.execution.incremental.IncrementalCallState;
@@ -59,6 +60,9 @@ public class ExecutionContext {
     private final ExecutionInput executionInput;
     private final Supplier<ExecutableNormalizedOperation> queryTree;
 
+    // this is modified after creation so it needs to be volatile to ensure visibility across Threads
+    private volatile DataLoaderDispatchStrategy dataLoaderDispatcherStrategy = DataLoaderDispatchStrategy.NO_OP;
+
     ExecutionContext(ExecutionContextBuilder builder) {
         this.graphQLSchema = builder.graphQLSchema;
         this.executionId = builder.executionId;
@@ -247,7 +251,9 @@ public List<GraphQLError> getErrors() {
         return errors.get();
     }
 
-    public ExecutionStrategy getQueryStrategy() { return queryStrategy; }
+    public ExecutionStrategy getQueryStrategy() {
+        return queryStrategy;
+    }
 
     public ExecutionStrategy getMutationStrategy() {
         return mutationStrategy;
@@ -275,6 +281,16 @@ public Supplier<ExecutableNormalizedOperation> getNormalizedQueryTree() {
         return queryTree;
     }
 
+    @Internal
+    public void setDataLoaderDispatcherStrategy(DataLoaderDispatchStrategy dataLoaderDispatcherStrategy) {
+        this.dataLoaderDispatcherStrategy = dataLoaderDispatcherStrategy;
+    }
+
+    @Internal
+    public DataLoaderDispatchStrategy getDataLoaderDispatcherStrategy() {
+        return dataLoaderDispatcherStrategy;
+    }
+
     /**
      * This helps you transform the current ExecutionContext object into another one by starting a builder with all
      * the current values and allows you to transform it how you want.