@@ -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,29 +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 ,
68
- InstrumentationFieldFetchParameters parameters ,
69
- InstrumentationState instrumentationState ) {
70
- DataLoaderDispatcherInstrumentationState state =
71
- InstrumentationState .ofState (instrumentationState );
68
+ public DataFetcher <?> instrumentDataFetcher (DataFetcher <?> dataFetcher ,
69
+ InstrumentationFieldFetchParameters parameters , InstrumentationState instrumentationState ) {
70
+ DataLoaderDispatcherInstrumentationState state = InstrumentationState .ofState (
71
+ instrumentationState );
72
72
if (state .isAggressivelyBatching ()) {
73
73
return dataFetcher ;
74
74
}
75
75
//
76
76
// currently only AsyncExecutionStrategy with DataLoader and hence this allows us to "dispatch"
77
77
// on every object if it's not using aggressive batching for other execution strategies
78
78
// which allows them to work if used.
79
- return (DataFetcher <Object >)
80
- environment -> {
81
- Object obj = dataFetcher .get (environment );
82
- doImmediatelyDispatch (state );
83
- return obj ;
84
- };
79
+ return (DataFetcher <Object >) environment -> {
80
+ Object obj = dataFetcher .get (environment );
81
+ doImmediatelyDispatch (state );
82
+ return obj ;
83
+ };
85
84
}
86
85
87
86
private void doImmediatelyDispatch (DataLoaderDispatcherInstrumentationState state ) {
@@ -93,8 +92,8 @@ public InstrumentationContext<ExecutionResult> beginExecuteOperation(
93
92
InstrumentationExecuteOperationParameters parameters ,
94
93
InstrumentationState instrumentationState ) {
95
94
if (!isDataLoaderCompatible (parameters .getExecutionContext ())) {
96
- DataLoaderDispatcherInstrumentationState state =
97
- InstrumentationState . ofState ( instrumentationState );
95
+ DataLoaderDispatcherInstrumentationState state = InstrumentationState . ofState (
96
+ instrumentationState );
98
97
state .setAggressivelyBatching (false );
99
98
}
100
99
return SimpleInstrumentationContext .noOp ();
@@ -118,8 +117,8 @@ private boolean isDataLoaderCompatible(ExecutionContext executionContext) {
118
117
public ExecutionStrategyInstrumentationContext beginExecutionStrategy (
119
118
InstrumentationExecutionStrategyParameters parameters ,
120
119
InstrumentationState instrumentationState ) {
121
- DataLoaderDispatcherInstrumentationState state =
122
- InstrumentationState . ofState ( instrumentationState );
120
+ DataLoaderDispatcherInstrumentationState state = InstrumentationState . ofState (
121
+ instrumentationState );
123
122
//
124
123
// if there are no data loaders, there is nothing to do
125
124
//
@@ -142,8 +141,8 @@ public void onCompleted(ExecutionResult result, Throwable t) {
142
141
@ Override
143
142
public InstrumentationContext <Object > beginFieldFetch (
144
143
InstrumentationFieldFetchParameters parameters , InstrumentationState instrumentationState ) {
145
- DataLoaderDispatcherInstrumentationState state =
146
- InstrumentationState . ofState ( instrumentationState );
144
+ DataLoaderDispatcherInstrumentationState state = InstrumentationState . ofState (
145
+ instrumentationState );
147
146
//
148
147
// if there are no data loaders, there is nothing to do
149
148
//
@@ -155,26 +154,25 @@ public InstrumentationContext<Object> beginFieldFetch(
155
154
156
155
@ Override
157
156
public CompletableFuture <ExecutionResult > instrumentExecutionResult (
158
- ExecutionResult executionResult ,
159
- InstrumentationExecutionParameters parameters ,
157
+ ExecutionResult executionResult , InstrumentationExecutionParameters parameters ,
160
158
InstrumentationState instrumentationState ) {
161
- DataLoaderDispatcherInstrumentationState state =
162
- InstrumentationState . ofState ( instrumentationState );
159
+ DataLoaderDispatcherInstrumentationState state = InstrumentationState . ofState (
160
+ instrumentationState );
163
161
state .getApproach ().removeTracking (parameters .getExecutionInput ().getExecutionId ());
164
162
if (!options .isIncludeStatistics ()) {
165
163
return CompletableFuture .completedFuture (executionResult );
166
164
} else {
167
165
Map <Object , Object > currentExt = executionResult .getExtensions ();
168
- Map <Object , Object > statsMap =
169
- new LinkedHashMap <>( currentExt == null ? Collections .emptyMap () : currentExt );
166
+ Map <Object , Object > statsMap = new LinkedHashMap <>(
167
+ currentExt == null ? Collections .emptyMap () : currentExt );
170
168
Map <Object , Object > dataLoaderStats = buildStatisticsMap (state );
171
169
statsMap .put ("dataloader" , dataLoaderStats );
172
170
173
171
log .debug ("Data loader stats : {}" , dataLoaderStats );
174
172
175
173
return CompletableFuture .completedFuture (
176
- new ExecutionResultImpl (
177
- executionResult . getData (), executionResult . getErrors (), statsMap ));
174
+ new ExecutionResultImpl (executionResult . getData (), executionResult . getErrors (),
175
+ statsMap ));
178
176
}
179
177
}
180
178
0 commit comments