Skip to content

Commit 3a8adcd

Browse files
bbakermanbbakerman
committed
Instrumentation support for dataloader -added registry instrumentation
1 parent 411b9bd commit 3a8adcd

9 files changed

+421
-75
lines changed

‎src/main/java/org/dataloader/DataLoaderRegistry.java‎

Lines changed: 72 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
package org.dataloader;
22

33
import org.dataloader.annotations.PublicApi;
4+
import org.dataloader.instrumentation.ChainedDataLoaderInstrumentation;
5+
import org.dataloader.instrumentation.DataLoaderInstrumentation;
6+
import org.dataloader.instrumentation.DataLoaderInstrumentationHelper;
47
import org.dataloader.stats.Statistics;
58

69
import java.util.ArrayList;
@@ -21,25 +24,81 @@
2124
@PublicApi
2225
public class DataLoaderRegistry {
2326
protected final Map<String, DataLoader<?, ?>> dataLoaders = new ConcurrentHashMap<>();
27+
protected final DataLoaderInstrumentation instrumentation;
28+
2429

2530
public DataLoaderRegistry() {
31+
instrumentation = null;
2632
}
2733

2834
private DataLoaderRegistry(Builder builder) {
29-
this.dataLoaders.putAll(builder.dataLoaders);
35+
instrument(builder.instrumentation, builder.dataLoaders);
36+
this.instrumentation = builder.instrumentation;
37+
}
38+
39+
private void instrument(DataLoaderInstrumentation registryInstrumentation, Map<String, DataLoader<?, ?>> incomingDataLoaders) {
40+
this.dataLoaders.putAll(incomingDataLoaders);
41+
if (registryInstrumentation != null) {
42+
this.dataLoaders.replaceAll((k, existingDL) -> instrumentDL(registryInstrumentation, existingDL));
43+
}
44+
}
45+
46+
/**
47+
* Can be called to tweak a {@link DataLoader} so that it has the registry {@link DataLoaderInstrumentation} added as the first one.
48+
*
49+
* @param registryInstrumentation the common registry {@link DataLoaderInstrumentation}
50+
* @param existingDL the existing data loader
51+
* @return a new {@link DataLoader} or the same one if there is nothing to change
52+
*/
53+
private static DataLoader<?, ?> instrumentDL(DataLoaderInstrumentation registryInstrumentation, DataLoader<?, ?> existingDL) {
54+
if (registryInstrumentation == null) {
55+
return existingDL;
56+
}
57+
DataLoaderOptions options = existingDL.getOptions();
58+
DataLoaderInstrumentation existingInstrumentation = options.getInstrumentation();
59+
// if they have any instrumentations then add to it
60+
if (existingInstrumentation != null) {
61+
if (existingInstrumentation == registryInstrumentation) {
62+
// nothing to change
63+
return existingDL;
64+
}
65+
if (existingInstrumentation == DataLoaderInstrumentationHelper.NOOP_INSTRUMENTATION) {
66+
// replace it with the registry one
67+
return mkInstrumentedDataLoader(existingDL, options, registryInstrumentation);
68+
}
69+
if (existingInstrumentation instanceof ChainedDataLoaderInstrumentation) {
70+
// avoids calling a chained inside a chained
71+
DataLoaderInstrumentation newInstrumentation = ((ChainedDataLoaderInstrumentation) existingInstrumentation).prepend(registryInstrumentation);
72+
return mkInstrumentedDataLoader(existingDL, options, newInstrumentation);
73+
} else {
74+
DataLoaderInstrumentation newInstrumentation = new ChainedDataLoaderInstrumentation().add(registryInstrumentation).add(existingInstrumentation);
75+
return mkInstrumentedDataLoader(existingDL, options, newInstrumentation);
76+
}
77+
} else {
78+
return mkInstrumentedDataLoader(existingDL, options, registryInstrumentation);
79+
}
80+
}
81+
82+
private static DataLoader<?, ?> mkInstrumentedDataLoader(DataLoader<?, ?> existingDL, DataLoaderOptions options, DataLoaderInstrumentation newInstrumentation) {
83+
return existingDL.transform(builder -> {
84+
options.setInstrumentation(newInstrumentation);
85+
builder.options(options);
86+
});
3087
}
3188

89+
public DataLoaderInstrumentation getInstrumentation() {
90+
return instrumentation;
91+
}
3292

3393
/**
3494
* This will register a new dataloader
3595
*
3696
* @param key the key to put the data loader under
3797
* @param dataLoader the data loader to register
38-
*
3998
* @return this registry
4099
*/
41100
public DataLoaderRegistry register(String key, DataLoader<?, ?> dataLoader) {
42-
dataLoaders.put(key, dataLoader);
101+
dataLoaders.put(key, instrumentDL(instrumentation, dataLoader));
43102
return this;
44103
}
45104

@@ -54,21 +113,22 @@ public DataLoaderRegistry register(String key, DataLoader<?, ?> dataLoader) {
54113
* @param mappingFunction the function to compute a data loader
55114
* @param <K> the type of keys
56115
* @param <V> the type of values
57-
*
58116
* @return a data loader
59117
*/
60118
@SuppressWarnings("unchecked")
61119
public <K, V> DataLoader<K, V> computeIfAbsent(final String key,
62120
final Function<String, DataLoader<?, ?>> mappingFunction) {
63-
return (DataLoader<K, V>) dataLoaders.computeIfAbsent(key, mappingFunction);
121+
return (DataLoader<K, V>) dataLoaders.computeIfAbsent(key, (k) -> {
122+
DataLoader<?, ?> dl = mappingFunction.apply(k);
123+
return instrumentDL(instrumentation, dl);
124+
});
64125
}
65126

66127
/**
67128
* This will combine all the current data loaders in this registry and all the data loaders from the specified registry
68129
* and return a new combined registry
69130
*
70131
* @param registry the registry to combine into this registry
71-
*
72132
* @return a new combined registry
73133
*/
74134
public DataLoaderRegistry combine(DataLoaderRegistry registry) {
@@ -97,7 +157,6 @@ public DataLoaderRegistry combine(DataLoaderRegistry registry) {
97157
* This will unregister a new dataloader
98158
*
99159
* @param key the key of the data loader to unregister
100-
*
101160
* @return this registry
102161
*/
103162
public DataLoaderRegistry unregister(String key) {
@@ -111,7 +170,6 @@ public DataLoaderRegistry unregister(String key) {
111170
* @param key the key of the data loader
112171
* @param <K> the type of keys
113172
* @param <V> the type of values
114-
*
115173
* @return a data loader or null if its not present
116174
*/
117175
@SuppressWarnings("unchecked")
@@ -182,13 +240,13 @@ public static Builder newRegistry() {
182240
public static class Builder {
183241

184242
private final Map<String, DataLoader<?, ?>> dataLoaders = new HashMap<>();
243+
private DataLoaderInstrumentation instrumentation;
185244

186245
/**
187246
* This will register a new dataloader
188247
*
189248
* @param key the key to put the data loader under
190249
* @param dataLoader the data loader to register
191-
*
192250
* @return this builder for a fluent pattern
193251
*/
194252
public Builder register(String key, DataLoader<?, ?> dataLoader) {
@@ -201,14 +259,18 @@ public Builder register(String key, DataLoader<?, ?> dataLoader) {
201259
* from a previous {@link DataLoaderRegistry}
202260
*
203261
* @param otherRegistry the previous {@link DataLoaderRegistry}
204-
*
205262
* @return this builder for a fluent pattern
206263
*/
207264
public Builder registerAll(DataLoaderRegistry otherRegistry) {
208265
dataLoaders.putAll(otherRegistry.dataLoaders);
209266
return this;
210267
}
211268

269+
public Builder instrumentation(DataLoaderInstrumentation instrumentation) {
270+
this.instrumentation = instrumentation;
271+
return this;
272+
}
273+
212274
/**
213275
* @return the newly built {@link DataLoaderRegistry}
214276
*/

‎src/main/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentation.java‎

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,10 @@
33
import org.dataloader.BatchLoaderEnvironment;
44
import org.dataloader.DataLoader;
55
import org.dataloader.DispatchResult;
6+
import org.dataloader.annotations.PublicApi;
67

78
import java.util.ArrayList;
9+
import java.util.Collection;
810
import java.util.List;
911
import java.util.Objects;
1012
import java.util.function.Function;
@@ -14,6 +16,7 @@
1416
* This {@link DataLoaderInstrumentation} can chain together multiple instrumentations and have them all called in
1517
* the order of the provided list.
1618
*/
19+
@PublicApi
1720
public class ChainedDataLoaderInstrumentation implements DataLoaderInstrumentation {
1821
private final List<DataLoaderInstrumentation> instrumentations;
1922

@@ -25,15 +28,44 @@ public ChainedDataLoaderInstrumentation(List<DataLoaderInstrumentation> instrume
2528
this.instrumentations = List.copyOf(instrumentations);
2629
}
2730

31+
public List<DataLoaderInstrumentation> getInstrumentations() {
32+
return instrumentations;
33+
}
34+
2835
/**
2936
* Adds a new {@link DataLoaderInstrumentation} to the list and creates a new {@link ChainedDataLoaderInstrumentation}
3037
*
3138
* @param instrumentation the one to add
3239
* @return a new ChainedDataLoaderInstrumentation object
3340
*/
3441
public ChainedDataLoaderInstrumentation add(DataLoaderInstrumentation instrumentation) {
35-
ArrayList<DataLoaderInstrumentation> list = new ArrayList<>(instrumentations);
42+
ArrayList<DataLoaderInstrumentation> list = new ArrayList<>(this.instrumentations);
43+
list.add(instrumentation);
44+
return new ChainedDataLoaderInstrumentation(list);
45+
}
46+
47+
/**
48+
* Prepends a new {@link DataLoaderInstrumentation} to the list and creates a new {@link ChainedDataLoaderInstrumentation}
49+
*
50+
* @param instrumentation the one to add
51+
* @return a new ChainedDataLoaderInstrumentation object
52+
*/
53+
public ChainedDataLoaderInstrumentation prepend(DataLoaderInstrumentation instrumentation) {
54+
ArrayList<DataLoaderInstrumentation> list = new ArrayList<>();
3655
list.add(instrumentation);
56+
list.addAll(this.instrumentations);
57+
return new ChainedDataLoaderInstrumentation(list);
58+
}
59+
60+
/**
61+
* Adds a collection of {@link DataLoaderInstrumentation} to the list and creates a new {@link ChainedDataLoaderInstrumentation}
62+
*
63+
* @param instrumentations the new ones to add
64+
* @return a new ChainedDataLoaderInstrumentation object
65+
*/
66+
public ChainedDataLoaderInstrumentation addAll(Collection<DataLoaderInstrumentation> instrumentations) {
67+
ArrayList<DataLoaderInstrumentation> list = new ArrayList<>(this.instrumentations);
68+
list.addAll(instrumentations);
3769
return new ChainedDataLoaderInstrumentation(list);
3870
}
3971

‎src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentation.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,14 @@
33
import org.dataloader.BatchLoaderEnvironment;
44
import org.dataloader.DataLoader;
55
import org.dataloader.DispatchResult;
6+
import org.dataloader.annotations.PublicSpi;
67

78
import java.util.List;
89

910
/**
1011
* This interface is called when certain actions happen inside a data loader
1112
*/
13+
@PublicSpi
1214
public interface DataLoaderInstrumentation {
1315
/**
1416
* This call back is done just before the {@link DataLoader#dispatch()} is invoked,

‎src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentationContext.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
package org.dataloader.instrumentation;
22

3+
import org.dataloader.annotations.PublicSpi;
4+
35
import java.util.concurrent.CompletableFuture;
46

57
/**
@@ -10,6 +12,7 @@
1012
* This pattern of construction of an object then call back is intended to allow "timers" to be created that can instrument what has
1113
* just happened or "loggers" to be called to record what has happened.
1214
*/
15+
@PublicSpi
1316
public interface DataLoaderInstrumentationContext<T> {
1417
/**
1518
* This is invoked when the instrumentation step is initially dispatched. Note this is NOT

‎src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentationHelper.java‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,11 @@
11
package org.dataloader.instrumentation;
22

3+
import org.dataloader.annotations.PublicApi;
4+
5+
@PublicApi
36
public class DataLoaderInstrumentationHelper {
47

8+
@SuppressWarnings("RedundantMethodOverride")
59
private static final DataLoaderInstrumentationContext<?> NOOP_CTX = new DataLoaderInstrumentationContext<>() {
610
@Override
711
public void onDispatched() {
@@ -12,17 +16,26 @@ public void onCompleted(Object result, Throwable t) {
1216
}
1317
};
1418

19+
/**
20+
* Returns a noop {@link DataLoaderInstrumentationContext} of the right type
21+
*
22+
* @param <T> for two
23+
* @return a noop context
24+
*/
1525
public static <T> DataLoaderInstrumentationContext<T> noOpCtx() {
1626
//noinspection unchecked
1727
return (DataLoaderInstrumentationContext<T>) NOOP_CTX;
1828
}
1929

30+
/**
31+
* A well known noop {@link DataLoaderInstrumentation}
32+
*/
2033
public static final DataLoaderInstrumentation NOOP_INSTRUMENTATION = new DataLoaderInstrumentation() {
2134
};
2235

2336
/**
2437
* Check the {@link DataLoaderInstrumentationContext} to see if its null and returns a noop if it is or else the original
25-
* context
38+
* context. This is a bit of a helper method.
2639
*
2740
* @param ic the context in play
2841
* @param <T> for two

‎src/test/java/org/dataloader/instrumentation/CapturingInstrumentation.java‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
package org.dataloader.instrumentation;
2+
3+
import org.dataloader.BatchLoaderEnvironment;
4+
import org.dataloader.DataLoader;
5+
import org.dataloader.DispatchResult;
6+
7+
import java.util.ArrayList;
8+
import java.util.List;
9+
10+
class CapturingInstrumentation implements DataLoaderInstrumentation {
11+
String name;
12+
List<String> methods = new ArrayList<>();
13+
14+
public CapturingInstrumentation(String name) {
15+
this.name = name;
16+
}
17+
18+
@Override
19+
public DataLoaderInstrumentationContext<DispatchResult<?>> beginDispatch(DataLoader<?, ?> dataLoader) {
20+
methods.add(name + "_beginDispatch");
21+
return new DataLoaderInstrumentationContext<>() {
22+
@Override
23+
public void onDispatched() {
24+
methods.add(name + "_beginDispatch_onDispatched");
25+
}
26+
27+
@Override
28+
public void onCompleted(DispatchResult<?> result, Throwable t) {
29+
methods.add(name + "_beginDispatch_onCompleted");
30+
}
31+
};
32+
}
33+
34+
@Override
35+
public DataLoaderInstrumentationContext<List<?>> beginBatchLoader(DataLoader<?, ?> dataLoader, List<?> keys, BatchLoaderEnvironment environment) {
36+
methods.add(name + "_beginBatchLoader");
37+
return new DataLoaderInstrumentationContext<>() {
38+
@Override
39+
public void onDispatched() {
40+
methods.add(name + "_beginBatchLoader_onDispatched");
41+
}
42+
43+
@Override
44+
public void onCompleted(List<?> result, Throwable t) {
45+
methods.add(name + "_beginBatchLoader_onCompleted");
46+
}
47+
};
48+
}
49+
}

‎src/test/java/org/dataloader/instrumentation/CapturingInstrumentationReturnsNull.java‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
package org.dataloader.instrumentation;
2+
3+
import org.dataloader.BatchLoaderEnvironment;
4+
import org.dataloader.DataLoader;
5+
import org.dataloader.DispatchResult;
6+
7+
import java.util.List;
8+
9+
class CapturingInstrumentationReturnsNull extends CapturingInstrumentation {
10+
11+
public CapturingInstrumentationReturnsNull(String name) {
12+
super(name);
13+
}
14+
15+
@Override
16+
public DataLoaderInstrumentationContext<DispatchResult<?>> beginDispatch(DataLoader<?, ?> dataLoader) {
17+
methods.add(name + "_beginDispatch");
18+
return null;
19+
}
20+
21+
@Override
22+
public DataLoaderInstrumentationContext<List<?>> beginBatchLoader(DataLoader<?, ?> dataLoader, List<?> keys, BatchLoaderEnvironment environment) {
23+
methods.add(name + "_beginBatchLoader");
24+
return null;
25+
}
26+
}

0 commit comments

Comments
 (0)