@@ -35,9 +35,7 @@ public class ConfigurableDispatchInstrumentation extends DataLoaderDispatcherIns
35
35
36
36
private final Function <DataLoaderRegistry , TrackingApproach > approachFunction ;
37
37
38
- /**
39
- * Creates a DataLoaderDispatcherInstrumentation with the default options
40
- */
38
+ /** Creates a DataLoaderDispatcherInstrumentation with the default options */
41
39
public ConfigurableDispatchInstrumentation (
42
40
Function <DataLoaderRegistry , TrackingApproach > approachFunction ) {
43
41
this (DataLoaderDispatcherInstrumentationOptions .newOptions (), approachFunction );
@@ -48,7 +46,8 @@ public ConfigurableDispatchInstrumentation(
48
46
*
49
47
* @param options the options to control the behaviour
50
48
*/
51
- public ConfigurableDispatchInstrumentation (DataLoaderDispatcherInstrumentationOptions options ,
49
+ public ConfigurableDispatchInstrumentation (
50
+ DataLoaderDispatcherInstrumentationOptions options ,
52
51
Function <DataLoaderRegistry , TrackingApproach > approachFunction ) {
53
52
this .options = options ;
54
53
this .approachFunction = approachFunction ;
@@ -60,27 +59,29 @@ public InstrumentationState createState(InstrumentationCreateStateParameters par
60
59
return new DataLoaderDispatcherInstrumentationState (
61
60
registry ,
62
61
approachFunction .apply (registry ),
63
- parameters .getExecutionInput ().getExecutionId ()
64
- );
62
+ parameters .getExecutionInput ().getExecutionId ());
65
63
}
66
64
67
65
@ Override
68
- public DataFetcher <?> instrumentDataFetcher (DataFetcher <?> dataFetcher ,
69
- InstrumentationFieldFetchParameters parameters , InstrumentationState instrumentationState ) {
70
- DataLoaderDispatcherInstrumentationState state = InstrumentationState .ofState (
71
- instrumentationState );
66
+ public DataFetcher <?> instrumentDataFetcher (
67
+ DataFetcher <?> dataFetcher ,
68
+ InstrumentationFieldFetchParameters parameters ,
69
+ InstrumentationState instrumentationState ) {
70
+ DataLoaderDispatcherInstrumentationState state =
71
+ InstrumentationState .ofState (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 >) environment -> {
80
- Object obj = dataFetcher .get (environment );
81
- doImmediatelyDispatch (state );
82
- return obj ;
83
- };
79
+ return (DataFetcher <Object >)
80
+ environment -> {
81
+ Object obj = dataFetcher .get (environment );
82
+ doImmediatelyDispatch (state );
83
+ return obj ;
84
+ };
84
85
}
85
86
86
87
private void doImmediatelyDispatch (DataLoaderDispatcherInstrumentationState state ) {
@@ -92,8 +93,8 @@ public InstrumentationContext<ExecutionResult> beginExecuteOperation(
92
93
InstrumentationExecuteOperationParameters parameters ,
93
94
InstrumentationState instrumentationState ) {
94
95
if (!isDataLoaderCompatible (parameters .getExecutionContext ())) {
95
- DataLoaderDispatcherInstrumentationState state = InstrumentationState . ofState (
96
- instrumentationState );
96
+ DataLoaderDispatcherInstrumentationState state =
97
+ InstrumentationState . ofState ( instrumentationState );
97
98
state .setAggressivelyBatching (false );
98
99
}
99
100
return SimpleInstrumentationContext .noOp ();
@@ -117,8 +118,8 @@ private boolean isDataLoaderCompatible(ExecutionContext executionContext) {
117
118
public ExecutionStrategyInstrumentationContext beginExecutionStrategy (
118
119
InstrumentationExecutionStrategyParameters parameters ,
119
120
InstrumentationState instrumentationState ) {
120
- DataLoaderDispatcherInstrumentationState state = InstrumentationState . ofState (
121
- instrumentationState );
121
+ DataLoaderDispatcherInstrumentationState state =
122
+ InstrumentationState . ofState ( instrumentationState );
122
123
//
123
124
// if there are no data loaders, there is nothing to do
124
125
//
@@ -141,8 +142,8 @@ public void onCompleted(ExecutionResult result, Throwable t) {
141
142
@ Override
142
143
public InstrumentationContext <Object > beginFieldFetch (
143
144
InstrumentationFieldFetchParameters parameters , InstrumentationState instrumentationState ) {
144
- DataLoaderDispatcherInstrumentationState state = InstrumentationState . ofState (
145
- instrumentationState );
145
+ DataLoaderDispatcherInstrumentationState state =
146
+ InstrumentationState . ofState ( instrumentationState );
146
147
//
147
148
// if there are no data loaders, there is nothing to do
148
149
//
@@ -154,25 +155,26 @@ public InstrumentationContext<Object> beginFieldFetch(
154
155
155
156
@ Override
156
157
public CompletableFuture <ExecutionResult > instrumentExecutionResult (
157
- ExecutionResult executionResult , InstrumentationExecutionParameters parameters ,
158
+ ExecutionResult executionResult ,
159
+ InstrumentationExecutionParameters parameters ,
158
160
InstrumentationState instrumentationState ) {
159
- DataLoaderDispatcherInstrumentationState state = InstrumentationState . ofState (
160
- instrumentationState );
161
+ DataLoaderDispatcherInstrumentationState state =
162
+ InstrumentationState . ofState ( instrumentationState );
161
163
state .getApproach ().removeTracking (parameters .getExecutionInput ().getExecutionId ());
162
164
if (!options .isIncludeStatistics ()) {
163
165
return CompletableFuture .completedFuture (executionResult );
164
166
} else {
165
167
Map <Object , Object > currentExt = executionResult .getExtensions ();
166
- Map <Object , Object > statsMap = new LinkedHashMap <>(
167
- currentExt == null ? Collections .emptyMap () : currentExt );
168
+ Map <Object , Object > statsMap =
169
+ new LinkedHashMap <>( currentExt == null ? Collections .emptyMap () : currentExt );
168
170
Map <Object , Object > dataLoaderStats = buildStatisticsMap (state );
169
171
statsMap .put ("dataloader" , dataLoaderStats );
170
172
171
173
log .debug ("Data loader stats : {}" , dataLoaderStats );
172
174
173
175
return CompletableFuture .completedFuture (
174
- new ExecutionResultImpl (executionResult . getData (), executionResult . getErrors (),
175
- statsMap ));
176
+ new ExecutionResultImpl (
177
+ executionResult . getData (), executionResult . getErrors (), statsMap ));
176
178
}
177
179
}
178
180
0 commit comments