From 26eba39cfea39532cb2df656df7d76661aa01809 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=91=A3=E8=8F=B2?= Date: Mon, 13 Aug 2018 14:47:49 +0800 Subject: [PATCH] make DataLoaderDispatcherInstrumentation support request scope for servlet usage --- .../DataLoaderDispatcherInstrumentation.java | 51 ++++++++++++++----- ...aLoaderDispatcherInstrumentationState.java | 30 +++++++++++ .../FieldLevelTrackingApproach.java | 2 +- 3 files changed, 69 insertions(+), 14 deletions(-) diff --git a/src/main/java/graphql/execution/instrumentation/dataloader/DataLoaderDispatcherInstrumentation.java b/src/main/java/graphql/execution/instrumentation/dataloader/DataLoaderDispatcherInstrumentation.java index 5dc14160a7..6da23a9f60 100644 --- 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 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 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 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) 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,17 +144,20 @@ 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 beginFieldFetch(InstrumentationFieldFetchParameters parameters) { - return fieldLevelTrackingApproach.beginFieldFetch(parameters); + DataLoaderDispatcherInstrumentationState state = parameters.getInstrumentationState(); + return state.getApproach().beginFieldFetch(parameters.withNewState(state.getState())); } @Override @@ -142,10 +165,11 @@ public CompletableFuture instrumentExecutionResult(ExecutionRes if (!options.isIncludeStatistics()) { return CompletableFuture.completedFuture(executionResult); } + DataLoaderDispatcherInstrumentationState state = parameters.getInstrumentationState(); Map currentExt = executionResult.getExtensions(); Map statsMap = new LinkedHashMap<>(); statsMap.putAll(currentExt == null ? Collections.emptyMap() : currentExt); - Map dataLoaderStats = buildStatsMap(); + Map dataLoaderStats = buildStatsMap(state); statsMap.put("dataloader", dataLoaderStats); log.debug("Data loader stats : {}", dataLoaderStats); @@ -153,7 +177,8 @@ public CompletableFuture instrumentExecutionResult(ExecutionRes return CompletableFuture.completedFuture(new ExecutionResultImpl(executionResult.getData(), executionResult.getErrors(), statsMap)); } - private Map buildStatsMap() { + private Map buildStatsMap(DataLoaderDispatcherInstrumentationState state ) { + DataLoaderRegistry dataLoaderRegistry = state.getDataLoaderRegistry(); Statistics allStats = dataLoaderRegistry.getStatistics(); Map 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 index d12b34428e..39ebd07b71 100644 --- 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 index dcd4ec0a7c..8da1fda0e7 100644 --- 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 expectedFetchCountPerLevel = new LinkedHashMap<>(); private final Map fetchCountPerLevel = new LinkedHashMap<>();