@@ -35,7 +35,9 @@ public class ConfigurableDispatchInstrumentation extends DataLoaderDispatcherIns
35
35
36
36
private final Function <DataLoaderRegistry , TrackingApproach > approachFunction ;
37
37
38
- /** Creates a DataLoaderDispatcherInstrumentation with the default options */
38
+ /**
39
+ * Creates a DataLoaderDispatcherInstrumentation with the default options
40
+ */
39
41
public ConfigurableDispatchInstrumentation (
40
42
Function <DataLoaderRegistry , TrackingApproach > approachFunction ) {
41
43
this (DataLoaderDispatcherInstrumentationOptions .newOptions (), approachFunction );
@@ -46,8 +48,7 @@ public ConfigurableDispatchInstrumentation(
46
48
*
47
49
* @param options the options to control the behaviour
48
50
*/
49
- public ConfigurableDispatchInstrumentation (
50
- DataLoaderDispatcherInstrumentationOptions options ,
51
+ public ConfigurableDispatchInstrumentation (DataLoaderDispatcherInstrumentationOptions options ,
51
52
Function <DataLoaderRegistry , TrackingApproach > approachFunction ) {
52
53
this .options = options ;
53
54
this .approachFunction = approachFunction ;
@@ -59,26 +60,27 @@ public InstrumentationState createState(InstrumentationCreateStateParameters par
59
60
return new DataLoaderDispatcherInstrumentationState (
60
61
registry ,
61
62
approachFunction .apply (registry ),
62
- parameters .getExecutionInput ().getExecutionId ());
63
+ parameters .getExecutionInput ().getExecutionId ()
64
+ );
63
65
}
64
66
65
67
@ 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 );
69
72
if (state .isAggressivelyBatching ()) {
70
73
return dataFetcher ;
71
74
}
72
75
//
73
76
// 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
75
78
// 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
+ };
82
84
}
83
85
84
86
private void doImmediatelyDispatch (DataLoaderDispatcherInstrumentationState state ) {
@@ -87,12 +89,14 @@ private void doImmediatelyDispatch(DataLoaderDispatcherInstrumentationState stat
87
89
88
90
@ Override
89
91
public InstrumentationContext <ExecutionResult > beginExecuteOperation (
90
- InstrumentationExecuteOperationParameters parameters ) {
92
+ InstrumentationExecuteOperationParameters parameters ,
93
+ InstrumentationState instrumentationState ) {
91
94
if (!isDataLoaderCompatible (parameters .getExecutionContext ())) {
92
- DataLoaderDispatcherInstrumentationState state = parameters .getInstrumentationState ();
95
+ DataLoaderDispatcherInstrumentationState state = InstrumentationState .ofState (
96
+ instrumentationState );
93
97
state .setAggressivelyBatching (false );
94
98
}
95
- return new SimpleInstrumentationContext <> ();
99
+ return SimpleInstrumentationContext . noOp ();
96
100
}
97
101
98
102
private boolean isDataLoaderCompatible (ExecutionContext executionContext ) {
@@ -111,8 +115,10 @@ private boolean isDataLoaderCompatible(ExecutionContext executionContext) {
111
115
112
116
@ Override
113
117
public ExecutionStrategyInstrumentationContext beginExecutionStrategy (
114
- InstrumentationExecutionStrategyParameters parameters ) {
115
- DataLoaderDispatcherInstrumentationState state = parameters .getInstrumentationState ();
118
+ InstrumentationExecutionStrategyParameters parameters ,
119
+ InstrumentationState instrumentationState ) {
120
+ DataLoaderDispatcherInstrumentationState state = InstrumentationState .ofState (
121
+ instrumentationState );
116
122
//
117
123
// if there are no data loaders, there is nothing to do
118
124
//
@@ -134,36 +140,39 @@ public void onCompleted(ExecutionResult result, Throwable t) {
134
140
135
141
@ Override
136
142
public InstrumentationContext <Object > beginFieldFetch (
137
- InstrumentationFieldFetchParameters parameters ) {
138
- DataLoaderDispatcherInstrumentationState state = parameters .getInstrumentationState ();
143
+ InstrumentationFieldFetchParameters parameters , InstrumentationState instrumentationState ) {
144
+ DataLoaderDispatcherInstrumentationState state = InstrumentationState .ofState (
145
+ instrumentationState );
139
146
//
140
147
// if there are no data loaders, there is nothing to do
141
148
//
142
149
if (state .hasNoDataLoaders ()) {
143
- return new SimpleInstrumentationContext <> ();
150
+ return SimpleInstrumentationContext . noOp ();
144
151
}
145
152
return state .getApproach ().beginFieldFetch (parameters );
146
153
}
147
154
148
155
@ Override
149
156
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 );
152
161
state .getApproach ().removeTracking (parameters .getExecutionInput ().getExecutionId ());
153
162
if (!options .isIncludeStatistics ()) {
154
163
return CompletableFuture .completedFuture (executionResult );
155
164
} else {
156
165
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 );
159
168
Map <Object , Object > dataLoaderStats = buildStatisticsMap (state );
160
169
statsMap .put ("dataloader" , dataLoaderStats );
161
170
162
171
log .debug ("Data loader stats : {}" , dataLoaderStats );
163
172
164
173
return CompletableFuture .completedFuture (
165
- new ExecutionResultImpl (
166
- executionResult . getData (), executionResult . getErrors (), statsMap ));
174
+ new ExecutionResultImpl (executionResult . getData (), executionResult . getErrors (),
175
+ statsMap ));
167
176
}
168
177
}
169
178
0 commit comments