Instrumentation support for dataloader - better tests and added ChainedDataLoaderInstrumentation

Author: bbakerman

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

@@ -0,0 +1,86 @@
+package org.dataloader.instrumentation;
+
+import org.dataloader.BatchLoaderEnvironment;
+import org.dataloader.DataLoader;
+import org.dataloader.DispatchResult;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+/**
+ * This {@link DataLoaderInstrumentation} can chain together multiple instrumentations and have them all called in
+ * the order of the provided list.
+ */
+public class ChainedDataLoaderInstrumentation implements DataLoaderInstrumentation {
+ private final List<DataLoaderInstrumentation> instrumentations;
+
+ public ChainedDataLoaderInstrumentation() {
+ instrumentations = List.of();
+ }
+
+ public ChainedDataLoaderInstrumentation(List<DataLoaderInstrumentation> instrumentations) {
+ this.instrumentations = List.copyOf(instrumentations);
+ }
+
+ /**
+ * Adds a new {@link DataLoaderInstrumentation} to the list and creates a new {@link ChainedDataLoaderInstrumentation}
+ *
+ * @param instrumentation the one to add
+ * @return a new ChainedDataLoaderInstrumentation object
+ */
+ public ChainedDataLoaderInstrumentation add(DataLoaderInstrumentation instrumentation) {
+ ArrayList<DataLoaderInstrumentation> list = new ArrayList<>(instrumentations);
+ list.add(instrumentation);
+ return new ChainedDataLoaderInstrumentation(list);
+ }
+
+ @Override
+ public DataLoaderInstrumentationContext<DispatchResult<?>> beginDispatch(DataLoader<?, ?> dataLoader) {
+ return chainedCtx(it -> it.beginDispatch(dataLoader));
+ }
+
+ @Override
+ public DataLoaderInstrumentationContext<List<?>> beginBatchLoader(DataLoader<?, ?> dataLoader, List<?> keys, BatchLoaderEnvironment environment) {
+ return chainedCtx(it -> it.beginBatchLoader(dataLoader, keys, environment));
+ }
+
+ private <T> DataLoaderInstrumentationContext<T> chainedCtx(Function<DataLoaderInstrumentation, DataLoaderInstrumentationContext<T>> mapper) {
+ // if we have zero or 1 instrumentations (and 1 is the most common), then we can avoid an object allocation
+ // of the ChainedInstrumentationContext since it won't be needed
+ if (instrumentations.isEmpty()) {
+ return DataLoaderInstrumentationHelper.noOpCtx();
+ }
+ if (instrumentations.size() == 1) {
+ return mapper.apply(instrumentations.get(0));
+ }
+ return new ChainedInstrumentationContext<>(dropNullContexts(mapper));
+ }
+
+ private <T> List<DataLoaderInstrumentationContext<T>> dropNullContexts(Function<DataLoaderInstrumentation, DataLoaderInstrumentationContext<T>> mapper) {
+ return instrumentations.stream()
+ .map(mapper)
+ .filter(Objects::nonNull)
+ .collect(Collectors.toList());
+ }
+
+ private static class ChainedInstrumentationContext<T> implements DataLoaderInstrumentationContext<T> {
+ private final List<DataLoaderInstrumentationContext<T>> contexts;
+
+ public ChainedInstrumentationContext(List<DataLoaderInstrumentationContext<T>> contexts) {
+ this.contexts = contexts;
+ }
+
+ @Override
+ public void onDispatched() {
+ contexts.forEach(DataLoaderInstrumentationContext::onDispatched);
+ }
+
+ @Override
+ public void onCompleted(T result, Throwable t) {
+ contexts.forEach(it -> it.onCompleted(result, t));
+ }
+ }
+}

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

@@ -11,7 +11,7 @@
  */
 public interface DataLoaderInstrumentation {
  /**
- * This call back is done just before the {@link DataLoader#dispatch()} is invoked
+ * This call back is done just before the {@link DataLoader#dispatch()} is invoked,
  * and it completes when the dispatch call promise is done.
  *
  * @param dataLoader the {@link DataLoader} in question
@@ -22,7 +22,7 @@ default DataLoaderInstrumentationContext<DispatchResult<?>> beginDispatch(DataLo
  }
 
  /**
- * This call back is done just before the batch loader of a {@link DataLoader} is invoked. Remember a batch loader
+ * This call back is done just before the `batch loader` of a {@link DataLoader} is invoked. Remember a batch loader
  * could be called multiple times during a dispatch event (because of max batch sizes)
  *
  * @param dataLoader the {@link DataLoader} in question

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

@@ -12,13 +12,15 @@
  */
 public interface DataLoaderInstrumentationContext<T> {
  /**
- * This is invoked when the instrumentation step is initially dispatched
+ * This is invoked when the instrumentation step is initially dispatched. Note this is NOT
+ * the same time as the {@link DataLoaderInstrumentation}`beginXXX()` starts, but rather after all the inner
+ * work has been done.
  */
  default void onDispatched() {
  }
 
  /**
- * This is invoked when the instrumentation step is fully completed
+ * This is invoked when the instrumentation step is fully completed.
  *
  * @param result the result of the step (which may be null)
  * @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

@@ -25,6 +25,10 @@ public interface TestDataLoaderFactory {
 
  // Convenience methods
 
+ default <K> DataLoader<K, K> idLoader(DataLoaderOptions options) {
+ return idLoader(options, new ArrayList<>());
+ }
+
  default <K> DataLoader<K, K> idLoader(List<Collection<K>> calls) {
  return idLoader(null, calls);
  }

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

@@ -0,0 +1,166 @@
+package org.dataloader.instrumentation;
+
+import org.dataloader.BatchLoaderEnvironment;
+import org.dataloader.DataLoader;
+import org.dataloader.DataLoaderFactory;
+import org.dataloader.DataLoaderOptions;
+import org.dataloader.DispatchResult;
+import org.dataloader.fixtures.TestKit;
+import org.dataloader.fixtures.parameterized.TestDataLoaderFactory;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.MethodSource;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+
+import static org.awaitility.Awaitility.await;
+import static org.dataloader.DataLoaderOptions.newOptions;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
+
+class ChainedDataLoaderInstrumentationTest {
+
+ static class CapturingInstrumentation implements DataLoaderInstrumentation {
+ String name;
+ List<String> methods = new ArrayList<>();
+
+ public CapturingInstrumentation(String name) {
+ this.name = name;
+ }
+
+ @Override
+ public DataLoaderInstrumentationContext<DispatchResult<?>> beginDispatch(DataLoader<?, ?> dataLoader) {
+ methods.add(name + "_beginDispatch");
+ return new DataLoaderInstrumentationContext<>() {
+ @Override
+ public void onDispatched() {
+ methods.add(name + "_beginDispatch_onDispatched");
+ }
+
+ @Override
+ public void onCompleted(DispatchResult<?> result, Throwable t) {
+ methods.add(name + "_beginDispatch_onCompleted");
+ }
+ };
+ }
+
+ @Override
+ public DataLoaderInstrumentationContext<List<?>> beginBatchLoader(DataLoader<?, ?> dataLoader, List<?> keys, BatchLoaderEnvironment environment) {
+ methods.add(name + "_beginBatchLoader");
+ return new DataLoaderInstrumentationContext<>() {
+ @Override
+ public void onDispatched() {
+ methods.add(name + "_beginBatchLoader_onDispatched");
+ }
+
+ @Override
+ public void onCompleted(List<?> result, Throwable t) {
+ methods.add(name + "_beginBatchLoader_onCompleted");
+ }
+ };
+ }
+ }
+
+
+ static class CapturingInstrumentationReturnsNull extends CapturingInstrumentation {
+
+ public CapturingInstrumentationReturnsNull(String name) {
+ super(name);
+ }
+
+ @Override
+ public DataLoaderInstrumentationContext<DispatchResult<?>> beginDispatch(DataLoader<?, ?> dataLoader) {
+ methods.add(name + "_beginDispatch");
+ return null;
+ }
+
+ @Override
+ public DataLoaderInstrumentationContext<List<?>> beginBatchLoader(DataLoader<?, ?> dataLoader, List<?> keys, BatchLoaderEnvironment environment) {
+ methods.add(name + "_beginBatchLoader");
+ return null;
+ }
+ }
+
+ @Test
+ void canChainTogetherZeroInstrumentation() {
+ // just to prove its useless but harmless
+ ChainedDataLoaderInstrumentation chainedItn = new ChainedDataLoaderInstrumentation();
+
+ DataLoaderOptions options = newOptions().setInstrumentation(chainedItn);
+
+ DataLoader<String, String> dl = DataLoaderFactory.newDataLoader(TestKit.keysAsValues(), options);
+
+ dl.load("A");
+ dl.load("B");
+
+ CompletableFuture<List<String>> dispatch = dl.dispatch();
+
+ await().until(dispatch::isDone);
+ assertThat(dispatch.join(), equalTo(List.of("A", "B")));
+ }
+
+ @Test
+ void canChainTogetherOneInstrumentation() {
+ CapturingInstrumentation capturingA = new CapturingInstrumentation("A");
+
+ ChainedDataLoaderInstrumentation chainedItn = new ChainedDataLoaderInstrumentation()
+ .add(capturingA);
+
+ DataLoaderOptions options = newOptions().setInstrumentation(chainedItn);
+
+ DataLoader<String, String> dl = DataLoaderFactory.newDataLoader(TestKit.keysAsValues(), options);
+
+ dl.load("A");
+ dl.load("B");
+
+ CompletableFuture<List<String>> dispatch = dl.dispatch();
+
+ await().until(dispatch::isDone);
+
+ assertThat(capturingA.methods, equalTo(List.of("A_beginDispatch",
+ "A_beginBatchLoader", "A_beginBatchLoader_onDispatched", "A_beginBatchLoader_onCompleted",
+ "A_beginDispatch_onDispatched", "A_beginDispatch_onCompleted")));
+ }
+
+
+ @ParameterizedTest
+ @MethodSource("org.dataloader.fixtures.parameterized.TestDataLoaderFactories#get")
+ public void canChainTogetherManyInstrumentationsWithDifferentBatchLoaders(TestDataLoaderFactory factory) {
+ CapturingInstrumentation capturingA = new CapturingInstrumentation("A");
+ CapturingInstrumentation capturingB = new CapturingInstrumentation("B");
+ CapturingInstrumentation capturingButReturnsNull = new CapturingInstrumentationReturnsNull("NULL");
+
+ ChainedDataLoaderInstrumentation chainedItn = new ChainedDataLoaderInstrumentation()
+ .add(capturingA)
+ .add(capturingB)
+ .add(capturingButReturnsNull);
+
+ DataLoaderOptions options = newOptions().setInstrumentation(chainedItn);
+
+ DataLoader<String, String> dl = factory.idLoader(options);
+
+ dl.load("A");
+ dl.load("B");
+
+ CompletableFuture<List<String>> dispatch = dl.dispatch();
+
+ await().until(dispatch::isDone);
+
+ //
+ // A_beginBatchLoader happens before A_beginDispatch_onDispatched because these are sync
+ // and no async - a batch scheduler or async batch loader would change that
+ //
+ assertThat(capturingA.methods, equalTo(List.of("A_beginDispatch",
+ "A_beginBatchLoader", "A_beginBatchLoader_onDispatched", "A_beginBatchLoader_onCompleted",
+ "A_beginDispatch_onDispatched", "A_beginDispatch_onCompleted")));
+
+ assertThat(capturingB.methods, equalTo(List.of("B_beginDispatch",
+ "B_beginBatchLoader", "B_beginBatchLoader_onDispatched", "B_beginBatchLoader_onCompleted",
+ "B_beginDispatch_onDispatched", "B_beginDispatch_onCompleted")));
+
+ // it returned null on all its contexts - nothing to call back on
+ assertThat(capturingButReturnsNull.methods, equalTo(List.of("NULL_beginDispatch", "NULL_beginBatchLoader")));
+ }
+}