make DataLoaderDispatcherInstrumentation support request scope for servlet usage

kdlan · bbakerman · commit fd2454fd02d7 · 2018-08-14T17:44:37.000+10:00
diff --git a/src/main/java/graphql/execution/instrumentation/dataloader/DataLoaderDispatcherInstrumentation.java b/src/main/java/graphql/execution/instrumentation/dataloader/DataLoaderDispatcherInstrumentation.java
@@ -28,6 +28,7 @@
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
+import java.util.function.Supplier;
 
 /**
  * This graphql {@link graphql.execution.instrumentation.Instrumentation} will dispatch
@@ -44,9 +45,8 @@ public class DataLoaderDispatcherInstrumentation extends SimpleInstrumentation {
 
     private static final Logger log = LoggerFactory.getLogger(DataLoaderDispatcherInstrumentation.class);
 
-    private final DataLoaderRegistry dataLoaderRegistry;
     private final DataLoaderDispatcherInstrumentationOptions options;
-    private final FieldLevelTrackingApproach fieldLevelTrackingApproach;
+    private final Supplier<DataLoaderRegistry> supplier;
 
     /**
      * You pass in a registry of N data loaders which will be {@link org.dataloader.DataLoader#dispatch() dispatched} as
@@ -66,15 +66,35 @@ public DataLoaderDispatcherInstrumentation(DataLoaderRegistry dataLoaderRegistry
      * @param options            the options to control the behaviour
      */
     public DataLoaderDispatcherInstrumentation(DataLoaderRegistry dataLoaderRegistry, DataLoaderDispatcherInstrumentationOptions options) {
-        this.dataLoaderRegistry = dataLoaderRegistry;
+        this(() -> dataLoaderRegistry, options);
+    }
+    
+    /**
+     * You pass in supplier of a registry of N data loaders which will be {@link org.dataloader.DataLoader#dispatch() dispatched} as
+     * each level of the query executes.
+     *
+     * @param supplier the supplier of registry of data loaders that will be dispatched
+     */
+    public DataLoaderDispatcherInstrumentation(Supplier<DataLoaderRegistry> supplier) {
+        this(supplier, DataLoaderDispatcherInstrumentationOptions.newOptions());
+    }
+    
+    /**
+     * You pass in a supplier of registry of N data loaders which will be {@link org.dataloader.DataLoader#dispatch() dispatched} as
+     * each level of the query executes.
+     *
+     * @param dataLoaderRegistry the registry of data loaders that will be dispatched
+     * @param options            the options to control the behaviour
+     */
+    public DataLoaderDispatcherInstrumentation(Supplier<DataLoaderRegistry> supplier, DataLoaderDispatcherInstrumentationOptions options) {
+        this.supplier = supplier;
         this.options = options;
-        this.fieldLevelTrackingApproach = new FieldLevelTrackingApproach(log, dataLoaderRegistry);
     }
 
 
     @Override
     public InstrumentationState createState() {
-        return fieldLevelTrackingApproach.createState();
+        return new DataLoaderDispatcherInstrumentationState(log, supplier.get());
     }
 
 
@@ -90,13 +110,13 @@ public DataFetcher<?> instrumentDataFetcher(DataFetcher<?> dataFetcher, Instrume
         // which allows them to work if used.
         return (DataFetcher<Object>) environment -> {
             Object obj = dataFetcher.get(environment);
-            immediatelyDispatch();
+            immediatelyDispatch(state);
             return obj;
         };
     }
 
-    private void immediatelyDispatch() {
-        fieldLevelTrackingApproach.dispatch();
+    private void immediatelyDispatch(DataLoaderDispatcherInstrumentationState state) {
+        state.getApproach().dispatch();
     }
 
     @Override
@@ -124,36 +144,41 @@ private boolean isDataLoaderCompatibleExecution(ExecutionContext executionContex
 
     @Override
     public ExecutionStrategyInstrumentationContext beginExecutionStrategy(InstrumentationExecutionStrategyParameters parameters) {
-        return fieldLevelTrackingApproach.beginExecutionStrategy(parameters);
+        DataLoaderDispatcherInstrumentationState state = parameters.getInstrumentationState();
+        return state.getApproach().beginExecutionStrategy(parameters.withNewState(state.getState()));
     }
 
     @Override
     public DeferredFieldInstrumentationContext beginDeferredField(InstrumentationDeferredFieldParameters parameters) {
-        return fieldLevelTrackingApproach.beginDeferredField(parameters);
+        DataLoaderDispatcherInstrumentationState state = parameters.getInstrumentationState();
+        return state.getApproach().beginDeferredField(parameters.withNewState(state.getState()));
     }
 
     @Override
     public InstrumentationContext<Object> beginFieldFetch(InstrumentationFieldFetchParameters parameters) {
-        return fieldLevelTrackingApproach.beginFieldFetch(parameters);
+        DataLoaderDispatcherInstrumentationState state = parameters.getInstrumentationState();
+        return state.getApproach().beginFieldFetch(parameters.withNewState(state.getState()));
     }
 
     @Override
     public CompletableFuture<ExecutionResult> instrumentExecutionResult(ExecutionResult executionResult, InstrumentationExecutionParameters parameters) {
         if (!options.isIncludeStatistics()) {
             return CompletableFuture.completedFuture(executionResult);
         }
+        DataLoaderDispatcherInstrumentationState state = parameters.getInstrumentationState();
         Map<Object, Object> currentExt = executionResult.getExtensions();
         Map<Object, Object> statsMap = new LinkedHashMap<>();
         statsMap.putAll(currentExt == null ? Collections.emptyMap() : currentExt);
-        Map<Object, Object> dataLoaderStats = buildStatsMap();
+        Map<Object, Object> dataLoaderStats = buildStatsMap(state);
         statsMap.put("dataloader", dataLoaderStats);
 
         log.debug("Data loader stats : {}", dataLoaderStats);
 
         return CompletableFuture.completedFuture(new ExecutionResultImpl(executionResult.getData(), executionResult.getErrors(), statsMap));
     }
 
-    private Map<Object, Object> buildStatsMap() {
+    private Map<Object, Object> buildStatsMap(DataLoaderDispatcherInstrumentationState state ) {
+        DataLoaderRegistry dataLoaderRegistry = state.getDataLoaderRegistry();
         Statistics allStats = dataLoaderRegistry.getStatistics();
         Map<Object, Object> statsMap = new LinkedHashMap<>();
         statsMap.put("overall-statistics", allStats.toMap());
diff --git a/src/main/java/graphql/execution/instrumentation/dataloader/DataLoaderDispatcherInstrumentationState.java b/src/main/java/graphql/execution/instrumentation/dataloader/DataLoaderDispatcherInstrumentationState.java
@@ -1,13 +1,30 @@
 package graphql.execution.instrumentation.dataloader;
 
+import org.dataloader.DataLoaderRegistry;
+import org.slf4j.Logger;
+
 import graphql.execution.instrumentation.InstrumentationState;
 
 /**
  * A base class that keeps track of whether aggressive batching can be used
  */
 public class DataLoaderDispatcherInstrumentationState implements InstrumentationState {
+
+    private final FieldLevelTrackingApproach approach;
+
+    private final DataLoaderRegistry dataLoaderRegistry;
+
+    private final InstrumentationState state;
+
     private boolean aggressivelyBatching = true;
 
+    public DataLoaderDispatcherInstrumentationState(Logger log, DataLoaderRegistry dataLoaderRegistry) {
+
+        this.dataLoaderRegistry = dataLoaderRegistry;
+        this.approach = new FieldLevelTrackingApproach(log, dataLoaderRegistry);
+        this.state = approach.createState();
+    }
+
     boolean isAggressivelyBatching() {
         return aggressivelyBatching;
     }
@@ -16,5 +33,18 @@ void setAggressivelyBatching(boolean aggressivelyBatching) {
         this.aggressivelyBatching = aggressivelyBatching;
     }
 
+    FieldLevelTrackingApproach getApproach() {
+        return approach;
+    }
+
+    DataLoaderRegistry getDataLoaderRegistry() {
+        return dataLoaderRegistry;
+    }
+
+    InstrumentationState getState() {
+        return state;
+    }
+    
+    
 
 }
diff --git a/src/main/java/graphql/execution/instrumentation/dataloader/FieldLevelTrackingApproach.java b/src/main/java/graphql/execution/instrumentation/dataloader/FieldLevelTrackingApproach.java
@@ -33,7 +33,7 @@ public class FieldLevelTrackingApproach {
     private final DataLoaderRegistry dataLoaderRegistry;
     private final Logger log;
 
-    private static class CallStack extends DataLoaderDispatcherInstrumentationState {
+    private static class CallStack implements InstrumentationState {
 
         private final Map<Integer, Integer> expectedFetchCountPerLevel = new LinkedHashMap<>();
         private final Map<Integer, Integer> fetchCountPerLevel = new LinkedHashMap<>();
