Skip to content

Commit 411b9bd

Browse files
bbakermanbbakerman
committed
Instrumentation support for dataloader - better tests and added ChainedDataLoaderInstrumentation
1 parent 97364ed commit 411b9bd

File tree

5 files changed

+262
-4
lines changed

5 files changed

+262
-4
lines changed

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

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
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+
import java.util.Objects;
10+
import java.util.function.Function;
11+
import java.util.stream.Collectors;
12+
13+
/**
14+
* This {@link DataLoaderInstrumentation} can chain together multiple instrumentations and have them all called in
15+
* the order of the provided list.
16+
*/
17+
public class ChainedDataLoaderInstrumentation implements DataLoaderInstrumentation {
18+
private final List<DataLoaderInstrumentation> instrumentations;
19+
20+
public ChainedDataLoaderInstrumentation() {
21+
instrumentations = List.of();
22+
}
23+
24+
public ChainedDataLoaderInstrumentation(List<DataLoaderInstrumentation> instrumentations) {
25+
this.instrumentations = List.copyOf(instrumentations);
26+
}
27+
28+
/**
29+
* Adds a new {@link DataLoaderInstrumentation} to the list and creates a new {@link ChainedDataLoaderInstrumentation}
30+
*
31+
* @param instrumentation the one to add
32+
* @return a new ChainedDataLoaderInstrumentation object
33+
*/
34+
public ChainedDataLoaderInstrumentation add(DataLoaderInstrumentation instrumentation) {
35+
ArrayList<DataLoaderInstrumentation> list = new ArrayList<>(instrumentations);
36+
list.add(instrumentation);
37+
return new ChainedDataLoaderInstrumentation(list);
38+
}
39+
40+
@Override
41+
public DataLoaderInstrumentationContext<DispatchResult<?>> beginDispatch(DataLoader<?, ?> dataLoader) {
42+
return chainedCtx(it -> it.beginDispatch(dataLoader));
43+
}
44+
45+
@Override
46+
public DataLoaderInstrumentationContext<List<?>> beginBatchLoader(DataLoader<?, ?> dataLoader, List<?> keys, BatchLoaderEnvironment environment) {
47+
return chainedCtx(it -> it.beginBatchLoader(dataLoader, keys, environment));
48+
}
49+
50+
private <T> DataLoaderInstrumentationContext<T> chainedCtx(Function<DataLoaderInstrumentation, DataLoaderInstrumentationContext<T>> mapper) {
51+
// if we have zero or 1 instrumentations (and 1 is the most common), then we can avoid an object allocation
52+
// of the ChainedInstrumentationContext since it won't be needed
53+
if (instrumentations.isEmpty()) {
54+
return DataLoaderInstrumentationHelper.noOpCtx();
55+
}
56+
if (instrumentations.size() == 1) {
57+
return mapper.apply(instrumentations.get(0));
58+
}
59+
return new ChainedInstrumentationContext<>(dropNullContexts(mapper));
60+
}
61+
62+
private <T> List<DataLoaderInstrumentationContext<T>> dropNullContexts(Function<DataLoaderInstrumentation, DataLoaderInstrumentationContext<T>> mapper) {
63+
return instrumentations.stream()
64+
.map(mapper)
65+
.filter(Objects::nonNull)
66+
.collect(Collectors.toList());
67+
}
68+
69+
private static class ChainedInstrumentationContext<T> implements DataLoaderInstrumentationContext<T> {
70+
private final List<DataLoaderInstrumentationContext<T>> contexts;
71+
72+
public ChainedInstrumentationContext(List<DataLoaderInstrumentationContext<T>> contexts) {
73+
this.contexts = contexts;
74+
}
75+
76+
@Override
77+
public void onDispatched() {
78+
contexts.forEach(DataLoaderInstrumentationContext::onDispatched);
79+
}
80+
81+
@Override
82+
public void onCompleted(T result, Throwable t) {
83+
contexts.forEach(it -> it.onCompleted(result, t));
84+
}
85+
}
86+
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
*/
1212
public interface DataLoaderInstrumentation {
1313
/**
14-
* This call back is done just before the {@link DataLoader#dispatch()} is invoked
14+
* This call back is done just before the {@link DataLoader#dispatch()} is invoked,
1515
* and it completes when the dispatch call promise is done.
1616
*
1717
* @param dataLoader the {@link DataLoader} in question
@@ -22,7 +22,7 @@ default DataLoaderInstrumentationContext<DispatchResult<?>> beginDispatch(DataLo
2222
}
2323

2424
/**
25-
* This call back is done just before the batch loader of a {@link DataLoader} is invoked. Remember a batch loader
25+
* This call back is done just before the `batch loader` of a {@link DataLoader} is invoked. Remember a batch loader
2626
* could be called multiple times during a dispatch event (because of max batch sizes)
2727
*
2828
* @param dataLoader the {@link DataLoader} in question

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,13 +12,15 @@
1212
*/
1313
public interface DataLoaderInstrumentationContext<T> {
1414
/**
15-
* This is invoked when the instrumentation step is initially dispatched
15+
* This is invoked when the instrumentation step is initially dispatched. Note this is NOT
16+
* the same time as the {@link DataLoaderInstrumentation}`beginXXX()` starts, but rather after all the inner
17+
* work has been done.
1618
*/
1719
default void onDispatched() {
1820
}
1921

2022
/**
21-
* This is invoked when the instrumentation step is fully completed
23+
* This is invoked when the instrumentation step is fully completed.
2224
*
2325
* @param result the result of the step (which may be null)
2426
* @param t this exception will be non-null if an exception was thrown during the step

‎src/test/java/org/dataloader/fixtures/parameterized/TestDataLoaderFactory.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,10 @@ public interface TestDataLoaderFactory {
2525

2626
// Convenience methods
2727

28+
default <K> DataLoader<K, K> idLoader(DataLoaderOptions options) {
29+
return idLoader(options, new ArrayList<>());
30+
}
31+
2832
default <K> DataLoader<K, K> idLoader(List<Collection<K>> calls) {
2933
return idLoader(null, calls);
3034
}

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

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
package org.dataloader.instrumentation;
2+
3+
import org.dataloader.BatchLoaderEnvironment;
4+
import org.dataloader.DataLoader;
5+
import org.dataloader.DataLoaderFactory;
6+
import org.dataloader.DataLoaderOptions;
7+
import org.dataloader.DispatchResult;
8+
import org.dataloader.fixtures.TestKit;
9+
import org.dataloader.fixtures.parameterized.TestDataLoaderFactory;
10+
import org.junit.jupiter.api.Test;
11+
import org.junit.jupiter.params.ParameterizedTest;
12+
import org.junit.jupiter.params.provider.MethodSource;
13+
14+
import java.util.ArrayList;
15+
import java.util.List;
16+
import java.util.concurrent.CompletableFuture;
17+
18+
import static org.awaitility.Awaitility.await;
19+
import static org.dataloader.DataLoaderOptions.newOptions;
20+
import static org.hamcrest.MatcherAssert.assertThat;
21+
import static org.hamcrest.Matchers.equalTo;
22+
23+
class ChainedDataLoaderInstrumentationTest {
24+
25+
static class CapturingInstrumentation implements DataLoaderInstrumentation {
26+
String name;
27+
List<String> methods = new ArrayList<>();
28+
29+
public CapturingInstrumentation(String name) {
30+
this.name = name;
31+
}
32+
33+
@Override
34+
public DataLoaderInstrumentationContext<DispatchResult<?>> beginDispatch(DataLoader<?, ?> dataLoader) {
35+
methods.add(name + "_beginDispatch");
36+
return new DataLoaderInstrumentationContext<>() {
37+
@Override
38+
public void onDispatched() {
39+
methods.add(name + "_beginDispatch_onDispatched");
40+
}
41+
42+
@Override
43+
public void onCompleted(DispatchResult<?> result, Throwable t) {
44+
methods.add(name + "_beginDispatch_onCompleted");
45+
}
46+
};
47+
}
48+
49+
@Override
50+
public DataLoaderInstrumentationContext<List<?>> beginBatchLoader(DataLoader<?, ?> dataLoader, List<?> keys, BatchLoaderEnvironment environment) {
51+
methods.add(name + "_beginBatchLoader");
52+
return new DataLoaderInstrumentationContext<>() {
53+
@Override
54+
public void onDispatched() {
55+
methods.add(name + "_beginBatchLoader_onDispatched");
56+
}
57+
58+
@Override
59+
public void onCompleted(List<?> result, Throwable t) {
60+
methods.add(name + "_beginBatchLoader_onCompleted");
61+
}
62+
};
63+
}
64+
}
65+
66+
67+
static class CapturingInstrumentationReturnsNull extends CapturingInstrumentation {
68+
69+
public CapturingInstrumentationReturnsNull(String name) {
70+
super(name);
71+
}
72+
73+
@Override
74+
public DataLoaderInstrumentationContext<DispatchResult<?>> beginDispatch(DataLoader<?, ?> dataLoader) {
75+
methods.add(name + "_beginDispatch");
76+
return null;
77+
}
78+
79+
@Override
80+
public DataLoaderInstrumentationContext<List<?>> beginBatchLoader(DataLoader<?, ?> dataLoader, List<?> keys, BatchLoaderEnvironment environment) {
81+
methods.add(name + "_beginBatchLoader");
82+
return null;
83+
}
84+
}
85+
86+
@Test
87+
void canChainTogetherZeroInstrumentation() {
88+
// just to prove its useless but harmless
89+
ChainedDataLoaderInstrumentation chainedItn = new ChainedDataLoaderInstrumentation();
90+
91+
DataLoaderOptions options = newOptions().setInstrumentation(chainedItn);
92+
93+
DataLoader<String, String> dl = DataLoaderFactory.newDataLoader(TestKit.keysAsValues(), options);
94+
95+
dl.load("A");
96+
dl.load("B");
97+
98+
CompletableFuture<List<String>> dispatch = dl.dispatch();
99+
100+
await().until(dispatch::isDone);
101+
assertThat(dispatch.join(), equalTo(List.of("A", "B")));
102+
}
103+
104+
@Test
105+
void canChainTogetherOneInstrumentation() {
106+
CapturingInstrumentation capturingA = new CapturingInstrumentation("A");
107+
108+
ChainedDataLoaderInstrumentation chainedItn = new ChainedDataLoaderInstrumentation()
109+
.add(capturingA);
110+
111+
DataLoaderOptions options = newOptions().setInstrumentation(chainedItn);
112+
113+
DataLoader<String, String> dl = DataLoaderFactory.newDataLoader(TestKit.keysAsValues(), options);
114+
115+
dl.load("A");
116+
dl.load("B");
117+
118+
CompletableFuture<List<String>> dispatch = dl.dispatch();
119+
120+
await().until(dispatch::isDone);
121+
122+
assertThat(capturingA.methods, equalTo(List.of("A_beginDispatch",
123+
"A_beginBatchLoader", "A_beginBatchLoader_onDispatched", "A_beginBatchLoader_onCompleted",
124+
"A_beginDispatch_onDispatched", "A_beginDispatch_onCompleted")));
125+
}
126+
127+
128+
@ParameterizedTest
129+
@MethodSource("org.dataloader.fixtures.parameterized.TestDataLoaderFactories#get")
130+
public void canChainTogetherManyInstrumentationsWithDifferentBatchLoaders(TestDataLoaderFactory factory) {
131+
CapturingInstrumentation capturingA = new CapturingInstrumentation("A");
132+
CapturingInstrumentation capturingB = new CapturingInstrumentation("B");
133+
CapturingInstrumentation capturingButReturnsNull = new CapturingInstrumentationReturnsNull("NULL");
134+
135+
ChainedDataLoaderInstrumentation chainedItn = new ChainedDataLoaderInstrumentation()
136+
.add(capturingA)
137+
.add(capturingB)
138+
.add(capturingButReturnsNull);
139+
140+
DataLoaderOptions options = newOptions().setInstrumentation(chainedItn);
141+
142+
DataLoader<String, String> dl = factory.idLoader(options);
143+
144+
dl.load("A");
145+
dl.load("B");
146+
147+
CompletableFuture<List<String>> dispatch = dl.dispatch();
148+
149+
await().until(dispatch::isDone);
150+
151+
//
152+
// A_beginBatchLoader happens before A_beginDispatch_onDispatched because these are sync
153+
// and no async - a batch scheduler or async batch loader would change that
154+
//
155+
assertThat(capturingA.methods, equalTo(List.of("A_beginDispatch",
156+
"A_beginBatchLoader", "A_beginBatchLoader_onDispatched", "A_beginBatchLoader_onCompleted",
157+
"A_beginDispatch_onDispatched", "A_beginDispatch_onCompleted")));
158+
159+
assertThat(capturingB.methods, equalTo(List.of("B_beginDispatch",
160+
"B_beginBatchLoader", "B_beginBatchLoader_onDispatched", "B_beginBatchLoader_onCompleted",
161+
"B_beginDispatch_onDispatched", "B_beginDispatch_onCompleted")));
162+
163+
// it returned null on all its contexts - nothing to call back on
164+
assertThat(capturingButReturnsNull.methods, equalTo(List.of("NULL_beginDispatch", "NULL_beginBatchLoader")));
165+
}
166+
}

0 commit comments

Comments
 (0)