@@ -35,7 +35,9 @@ public class ConfigurableDispatchInstrumentation extends DataLoaderDispatcherIns
3535
3636 private final Function <DataLoaderRegistry , TrackingApproach > approachFunction ;
3737
38- /** Creates a DataLoaderDispatcherInstrumentation with the default options */
38+ /**
39+ * Creates a DataLoaderDispatcherInstrumentation with the default options
40+ */
3941 public ConfigurableDispatchInstrumentation (
4042 Function <DataLoaderRegistry , TrackingApproach > approachFunction ) {
4143 this (DataLoaderDispatcherInstrumentationOptions .newOptions (), approachFunction );
@@ -46,8 +48,7 @@ public ConfigurableDispatchInstrumentation(
4648 *
4749 * @param options the options to control the behaviour
4850 */
49- public ConfigurableDispatchInstrumentation (
50- DataLoaderDispatcherInstrumentationOptions options ,
51+ public ConfigurableDispatchInstrumentation (DataLoaderDispatcherInstrumentationOptions options ,
5152 Function <DataLoaderRegistry , TrackingApproach > approachFunction ) {
5253 this .options = options ;
5354 this .approachFunction = approachFunction ;
@@ -59,26 +60,27 @@ public InstrumentationState createState(InstrumentationCreateStateParameters par
5960 return new DataLoaderDispatcherInstrumentationState (
6061 registry ,
6162 approachFunction .apply (registry ),
62- parameters .getExecutionInput ().getExecutionId ());
63+ parameters .getExecutionInput ().getExecutionId ()
64+ );
6365 }
6466
6567 @ Override
66- public DataFetcher <?> instrumentDataFetcher (
67- DataFetcher <?> dataFetcher , InstrumentationFieldFetchParameters parameters ) {
68- DataLoaderDispatcherInstrumentationState state = parameters .getInstrumentationState ();
68+ public DataFetcher <?> instrumentDataFetcher (DataFetcher <?> dataFetcher ,
69+ InstrumentationFieldFetchParameters parameters , InstrumentationState instrumentationState ) {
70+ DataLoaderDispatcherInstrumentationState state = InstrumentationState .ofState (
71+ instrumentationState );
6972 if (state .isAggressivelyBatching ()) {
7073 return dataFetcher ;
7174 }
7275 //
7376 // currently only AsyncExecutionStrategy with DataLoader and hence this allows us to "dispatch"
74- // on every object if its not using aggressive batching for other execution strategies
77+ // on every object if it's not using aggressive batching for other execution strategies
7578 // which allows them to work if used.
76- return (DataFetcher <Object >)
77- environment -> {
78- Object obj = dataFetcher .get (environment );
79- doImmediatelyDispatch (state );
80- return obj ;
81- };
79+ return (DataFetcher <Object >) environment -> {
80+ Object obj = dataFetcher .get (environment );
81+ doImmediatelyDispatch (state );
82+ return obj ;
83+ };
8284 }
8385
8486 private void doImmediatelyDispatch (DataLoaderDispatcherInstrumentationState state ) {
@@ -87,12 +89,14 @@ private void doImmediatelyDispatch(DataLoaderDispatcherInstrumentationState stat
8789
8890 @ Override
8991 public InstrumentationContext <ExecutionResult > beginExecuteOperation (
90- InstrumentationExecuteOperationParameters parameters ) {
92+ InstrumentationExecuteOperationParameters parameters ,
93+ InstrumentationState instrumentationState ) {
9194 if (!isDataLoaderCompatible (parameters .getExecutionContext ())) {
92- DataLoaderDispatcherInstrumentationState state = parameters .getInstrumentationState ();
95+ DataLoaderDispatcherInstrumentationState state = InstrumentationState .ofState (
96+ instrumentationState );
9397 state .setAggressivelyBatching (false );
9498 }
95- return new SimpleInstrumentationContext <> ();
99+ return SimpleInstrumentationContext . noOp ();
96100 }
97101
98102 private boolean isDataLoaderCompatible (ExecutionContext executionContext ) {
@@ -111,8 +115,10 @@ private boolean isDataLoaderCompatible(ExecutionContext executionContext) {
111115
112116 @ Override
113117 public ExecutionStrategyInstrumentationContext beginExecutionStrategy (
114- InstrumentationExecutionStrategyParameters parameters ) {
115- DataLoaderDispatcherInstrumentationState state = parameters .getInstrumentationState ();
118+ InstrumentationExecutionStrategyParameters parameters ,
119+ InstrumentationState instrumentationState ) {
120+ DataLoaderDispatcherInstrumentationState state = InstrumentationState .ofState (
121+ instrumentationState );
116122 //
117123 // if there are no data loaders, there is nothing to do
118124 //
@@ -134,36 +140,39 @@ public void onCompleted(ExecutionResult result, Throwable t) {
134140
135141 @ Override
136142 public InstrumentationContext <Object > beginFieldFetch (
137- InstrumentationFieldFetchParameters parameters ) {
138- DataLoaderDispatcherInstrumentationState state = parameters .getInstrumentationState ();
143+ InstrumentationFieldFetchParameters parameters , InstrumentationState instrumentationState ) {
144+ DataLoaderDispatcherInstrumentationState state = InstrumentationState .ofState (
145+ instrumentationState );
139146 //
140147 // if there are no data loaders, there is nothing to do
141148 //
142149 if (state .hasNoDataLoaders ()) {
143- return new SimpleInstrumentationContext <> ();
150+ return SimpleInstrumentationContext . noOp ();
144151 }
145152 return state .getApproach ().beginFieldFetch (parameters );
146153 }
147154
148155 @ Override
149156 public CompletableFuture <ExecutionResult > instrumentExecutionResult (
150- ExecutionResult executionResult , InstrumentationExecutionParameters parameters ) {
151- DataLoaderDispatcherInstrumentationState state = parameters .getInstrumentationState ();
157+ ExecutionResult executionResult , InstrumentationExecutionParameters parameters ,
158+ InstrumentationState instrumentationState ) {
159+ DataLoaderDispatcherInstrumentationState state = InstrumentationState .ofState (
160+ instrumentationState );
152161 state .getApproach ().removeTracking (parameters .getExecutionInput ().getExecutionId ());
153162 if (!options .isIncludeStatistics ()) {
154163 return CompletableFuture .completedFuture (executionResult );
155164 } else {
156165 Map <Object , Object > currentExt = executionResult .getExtensions ();
157- Map <Object , Object > statsMap =
158- new LinkedHashMap <>( currentExt == null ? Collections .emptyMap () : currentExt );
166+ Map <Object , Object > statsMap = new LinkedHashMap <>(
167+ currentExt == null ? Collections .emptyMap () : currentExt );
159168 Map <Object , Object > dataLoaderStats = buildStatisticsMap (state );
160169 statsMap .put ("dataloader" , dataLoaderStats );
161170
162171 log .debug ("Data loader stats : {}" , dataLoaderStats );
163172
164173 return CompletableFuture .completedFuture (
165- new ExecutionResultImpl (
166- executionResult . getData (), executionResult . getErrors (), statsMap ));
174+ new ExecutionResultImpl (executionResult . getData (), executionResult . getErrors (),
175+ statsMap ));
167176 }
168177 }
169178
0 commit comments