Instrumentation support for dataloader - Adde doco and SimpleDataLoaderInstrumentationContext

Author: bbakerman

README.md

@@ -750,6 +750,65 @@ When ticker mode is **true** the `ScheduledDataLoaderRegistry` algorithm is as f
         * If it returns **true**, then `dataLoader.dispatch()` is called **and** a task is scheduled to re-evaluate this specific dataloader in the near future
 * The re-evaluation tasks are run periodically according to the `registry.getScheduleDuration()`
 
+## Instrumenting the data loader code
+
+A `DataLoader` can have a `DataLoaderInstrumentation` associated with it.  This callback interface is intended to provide
+insight into working of the `DataLoader` such as how long it takes to run or to allow for logging of key events.
+
+You set the `DataLoaderInstrumentation` into the `DataLoaderOptions` at build time.
+
+```java
+
+
+        DataLoaderInstrumentation timingInstrumentation = new DataLoaderInstrumentation() {
+            @Override
+            public DataLoaderInstrumentationContext<DispatchResult<?>> beginDispatch(DataLoader<?, ?> dataLoader) {
+                long then = System.currentTimeMillis();
+                return DataLoaderInstrumentationHelper.whenCompleted((result, err) -> {
+                    long ms = System.currentTimeMillis() - then;
+                    System.out.println(format("dispatch time: %d ms", ms));
+                });
+            }
+
+            @Override
+            public DataLoaderInstrumentationContext<List<?>> beginBatchLoader(DataLoader<?, ?> dataLoader, List<?> keys, BatchLoaderEnvironment environment) {
+                long then = System.currentTimeMillis();
+                return DataLoaderInstrumentationHelper.whenCompleted((result, err) -> {
+                    long ms = System.currentTimeMillis() - then;
+                    System.out.println(format("batch loader time: %d ms", ms));
+                });
+            }
+        };
+        DataLoaderOptions options = DataLoaderOptions.newOptions().setInstrumentation(timingInstrumentation);
+        DataLoader<String, User> userDataLoader = DataLoaderFactory.newDataLoader(userBatchLoader, options);
+        
+```
+
+The example shows how long the overall `DataLoader` dispatch takes or how long the batch loader takes to run.
+
+### Instrumenting the DataLoaderRegistry
+
+You can also associate a `DataLoaderInstrumentation` with a `DataLoaderRegistry`.  Every `DataLoader` registered will be changed so that the registry
+`DataLoaderInstrumentation` is associated with it.  This allows you to set just the one `DataLoaderInstrumentation` in place and it applies to all 
+data loaders.
+
+```java
+    DataLoader<String, User> userDataLoader = DataLoaderFactory.newDataLoader(userBatchLoader);
+    DataLoader<String, User> teamsDataLoader = DataLoaderFactory.newDataLoader(teamsBatchLoader);
+    
+    DataLoaderRegistry registry = DataLoaderRegistry.newRegistry()
+            .instrumentation(timingInstrumentation)
+            .register("users", userDataLoader)
+            .register("teams", teamsDataLoader)
+            .build();
+    
+    DataLoader<String, User> changedUsersDataLoader = registry.getDataLoader("users");
+```
+
+The `timingInstrumentation` here will be associated with the `DataLoader` under the key  `users` and the key `teams`.  Note that since
+DataLoader is immutable, a new changed object is created so you must use the registry to get the `DataLoader`.  
+
+
 ## Other information sources
 
 - [Facebook DataLoader Github repo](https://github.com/facebook/dataloader)

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

@@ -2,6 +2,8 @@
 
 import org.dataloader.annotations.PublicApi;
 
+import java.util.function.BiConsumer;
+
 @PublicApi
 public class DataLoaderInstrumentationHelper {
 
@@ -33,6 +35,31 @@ public static <T> DataLoaderInstrumentationContext<T> noOpCtx() {
     public static final DataLoaderInstrumentation NOOP_INSTRUMENTATION = new DataLoaderInstrumentation() {
     };
 
+    /**
+     * Allows for the more fluent away to return an instrumentation context that runs the specified
+     * code on instrumentation step dispatch.
+     *
+     * @param codeToRun the code to run on dispatch
+     * @param <U>       the generic type
+     * @return an instrumentation context
+     */
+    public static <U> DataLoaderInstrumentationContext<U> whenDispatched(Runnable codeToRun) {
+        return new SimpleDataLoaderInstrumentationContext<>(codeToRun, null);
+    }
+
+    /**
+     * Allows for the more fluent away to return an instrumentation context that runs the specified
+     * code on instrumentation step completion.
+     *
+     * @param codeToRun the code to run on completion
+     * @param <U>       the generic type
+     * @return an instrumentation context
+     */
+    public static <U> DataLoaderInstrumentationContext<U> whenCompleted(BiConsumer<U, Throwable> codeToRun) {
+        return new SimpleDataLoaderInstrumentationContext<>(null, codeToRun);
+    }
+
+
     /**
      * Check the {@link DataLoaderInstrumentationContext} to see if its null and returns a noop if it is or else the original
      * context.  This is a bit of a helper method.

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

@@ -0,0 +1,35 @@
+package org.dataloader.instrumentation;
+
+
+import org.dataloader.annotations.Internal;
+
+import java.util.function.BiConsumer;
+
+/**
+ * A simple implementation of {@link DataLoaderInstrumentationContext}
+ */
+@Internal
+class SimpleDataLoaderInstrumentationContext<T> implements DataLoaderInstrumentationContext<T> {
+
+    private final BiConsumer<T, Throwable> codeToRunOnComplete;
+    private final Runnable codeToRunOnDispatch;
+
+    SimpleDataLoaderInstrumentationContext(Runnable codeToRunOnDispatch, BiConsumer<T, Throwable> codeToRunOnComplete) {
+        this.codeToRunOnComplete = codeToRunOnComplete;
+        this.codeToRunOnDispatch = codeToRunOnDispatch;
+    }
+
+    @Override
+    public void onDispatched() {
+        if (codeToRunOnDispatch != null) {
+            codeToRunOnDispatch.run();
+        }
+    }
+
+    @Override
+    public void onCompleted(T result, Throwable t) {
+        if (codeToRunOnComplete != null) {
+            codeToRunOnComplete.accept(result, t);
+        }
+    }
+}

src/test/java/ReadmeExamples.java

@@ -6,12 +6,17 @@
 import org.dataloader.DataLoader;
 import org.dataloader.DataLoaderFactory;
 import org.dataloader.DataLoaderOptions;
+import org.dataloader.DataLoaderRegistry;
+import org.dataloader.DispatchResult;
 import org.dataloader.MappedBatchLoaderWithContext;
 import org.dataloader.MappedBatchPublisher;
 import org.dataloader.Try;
 import org.dataloader.fixtures.SecurityCtx;
 import org.dataloader.fixtures.User;
 import org.dataloader.fixtures.UserManager;
+import org.dataloader.instrumentation.DataLoaderInstrumentation;
+import org.dataloader.instrumentation.DataLoaderInstrumentationContext;
+import org.dataloader.instrumentation.DataLoaderInstrumentationHelper;
 import org.dataloader.registries.DispatchPredicate;
 import org.dataloader.registries.ScheduledDataLoaderRegistry;
 import org.dataloader.scheduler.BatchLoaderScheduler;
@@ -228,6 +233,7 @@ private void clearCacheOnError() {
     }
 
     BatchLoader<String, User> userBatchLoader;
+    BatchLoader<String, User> teamsBatchLoader;
 
     private void disableCache() {
         DataLoaderFactory.newDataLoader(userBatchLoader, DataLoaderOptions.newOptions().setCachingEnabled(false));
@@ -380,4 +386,45 @@ private void ScheduledDispatcherChained() {
                 .build();
 
     }
+
+    private DataLoaderInstrumentation timingInstrumentation = DataLoaderInstrumentationHelper.NOOP_INSTRUMENTATION;
+
+    private void instrumentationExample() {
+
+        DataLoaderInstrumentation timingInstrumentation = new DataLoaderInstrumentation() {
+            @Override
+            public DataLoaderInstrumentationContext<DispatchResult<?>> beginDispatch(DataLoader<?, ?> dataLoader) {
+                long then = System.currentTimeMillis();
+                return DataLoaderInstrumentationHelper.whenCompleted((result, err) -> {
+                    long ms = System.currentTimeMillis() - then;
+                    System.out.println(format("dispatch time: %d ms", ms));
+                });
+            }
+
+            @Override
+            public DataLoaderInstrumentationContext<List<?>> beginBatchLoader(DataLoader<?, ?> dataLoader, List<?> keys, BatchLoaderEnvironment environment) {
+                long then = System.currentTimeMillis();
+                return DataLoaderInstrumentationHelper.whenCompleted((result, err) -> {
+                    long ms = System.currentTimeMillis() - then;
+                    System.out.println(format("batch loader time: %d ms", ms));
+                });
+            }
+        };
+        DataLoaderOptions options = DataLoaderOptions.newOptions().setInstrumentation(timingInstrumentation);
+        DataLoader<String, User> userDataLoader = DataLoaderFactory.newDataLoader(userBatchLoader, options);
+    }
+
+    private void registryExample() {
+        DataLoader<String, User> userDataLoader = DataLoaderFactory.newDataLoader(userBatchLoader);
+        DataLoader<String, User> teamsDataLoader = DataLoaderFactory.newDataLoader(teamsBatchLoader);
+
+        DataLoaderRegistry registry = DataLoaderRegistry.newRegistry()
+                .instrumentation(timingInstrumentation)
+                .register("users", userDataLoader)
+                .register("teams", teamsDataLoader)
+                .build();
+
+        DataLoader<String, User> changedUsersDataLoader = registry.getDataLoader("users");
+
+    }
 }

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

@@ -0,0 +1,49 @@
+package org.dataloader.instrumentation;
+
+import org.hamcrest.Matchers;
+import org.junit.jupiter.api.Test;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.nullValue;
+
+public class SimpleDataLoaderInstrumentationContextTest {
+
+    @Test
+    void canRunCompletedCodeAsExpected() {
+        AtomicReference<Object> actual = new AtomicReference<>();
+        AtomicReference<Object> actualErr = new AtomicReference<>();
+
+        DataLoaderInstrumentationContext<Object> ctx = DataLoaderInstrumentationHelper.whenCompleted((r, err) -> {
+            actualErr.set(err);
+            actual.set(r);
+        });
+
+        ctx.onDispatched(); // nothing happens
+        assertThat(actual.get(), nullValue());
+        assertThat(actualErr.get(), nullValue());
+
+        ctx.onCompleted("X", null);
+        assertThat(actual.get(), Matchers.equalTo("X"));
+        assertThat(actualErr.get(), nullValue());
+
+        ctx.onCompleted(null, new RuntimeException());
+        assertThat(actual.get(), nullValue());
+        assertThat(actualErr.get(), Matchers.instanceOf(RuntimeException.class));
+    }
+
+    @Test
+    void canRunOnDispatchCodeAsExpected() {
+        AtomicBoolean dispatchedCalled = new AtomicBoolean();
+
+        DataLoaderInstrumentationContext<Object> ctx = DataLoaderInstrumentationHelper.whenDispatched(() -> dispatchedCalled.set(true));
+
+        ctx.onCompleted("X", null); // nothing happens
+        assertThat(dispatchedCalled.get(), Matchers.equalTo(false));
+
+        ctx.onDispatched();
+        assertThat(dispatchedCalled.get(), Matchers.equalTo(true));
+    }
+}