Instrumentatin support for dataloader

Author: bbakerman

src/main/java/org/dataloader/DataLoaderFactory.java

@@ -561,7 +561,7 @@ public Builder<K, V> options(DataLoaderOptions options) {
             return this;
         }
 
-        DataLoader<K, V> build() {
+        public DataLoader<K, V> build() {
             return mkDataLoader(batchLoadFunction, options);
         }
     }

src/main/java/org/dataloader/DataLoaderHelper.java

@@ -3,6 +3,8 @@
 import org.dataloader.annotations.GuardedBy;
 import org.dataloader.annotations.Internal;
 import org.dataloader.impl.CompletableFutureKit;
+import org.dataloader.instrumentation.DataLoaderInstrumentation;
+import org.dataloader.instrumentation.DataLoaderInstrumentationContext;
 import org.dataloader.reactive.ReactiveSupport;
 import org.dataloader.scheduler.BatchLoaderScheduler;
 import org.dataloader.stats.StatisticsCollector;
@@ -34,6 +36,7 @@
 import static java.util.stream.Collectors.toList;
 import static org.dataloader.impl.Assertions.assertState;
 import static org.dataloader.impl.Assertions.nonNull;
+import static org.dataloader.instrumentation.DataLoaderInstrumentationHelper.ctxOrNoopCtx;
 
 /**
  * This helps break up the large DataLoader class functionality, and it contains the logic to dispatch the
@@ -167,6 +170,8 @@ Object getCacheKeyWithContext(K key, Object context) {
     }
 
     DispatchResult<V> dispatch() {
+        DataLoaderInstrumentationContext<DispatchResult<?>> instrCtx = ctxOrNoopCtx(instrumentation().beginDispatch(dataLoader));
+
         boolean batchingEnabled = loaderOptions.batchingEnabled();
         final List<K> keys;
         final List<Object> callContexts;
@@ -175,7 +180,8 @@ DispatchResult<V> dispatch() {
             int queueSize = loaderQueue.size();
             if (queueSize == 0) {
                 lastDispatchTime.set(now());
-                return emptyDispatchResult();
+                instrCtx.onDispatched();
+                return endDispatchCtx(instrCtx, emptyDispatchResult());
             }
 
             // we copy the pre-loaded set of futures ready for dispatch
@@ -192,7 +198,8 @@ DispatchResult<V> dispatch() {
             lastDispatchTime.set(now());
         }
         if (!batchingEnabled) {
-            return emptyDispatchResult();
+            instrCtx.onDispatched();
+            return endDispatchCtx(instrCtx, emptyDispatchResult());
         }
         final int totalEntriesHandled = keys.size();
         //
@@ -213,7 +220,15 @@ DispatchResult<V> dispatch() {
         } else {
             futureList = dispatchQueueBatch(keys, callContexts, queuedFutures);
         }
-        return new DispatchResult<>(futureList, totalEntriesHandled);
+        instrCtx.onDispatched();
+        return endDispatchCtx(instrCtx, new DispatchResult<>(futureList, totalEntriesHandled));
+    }
+
+    private DispatchResult<V> endDispatchCtx(DataLoaderInstrumentationContext<DispatchResult<?>> instrCtx, DispatchResult<V> dispatchResult) {
+        // once the CF completes, we can tell the instrumentation
+        dispatchResult.getPromisedResults()
+                .whenComplete((result, throwable) -> instrCtx.onCompleted(dispatchResult, throwable));
+        return dispatchResult;
     }
 
     private CompletableFuture<List<V>> sliceIntoBatchesOfBatches(List<K> keys, List<CompletableFuture<V>> queuedFutures, List<Object> callContexts, int maxBatchSize) {
@@ -427,11 +442,14 @@ CompletableFuture<List<V>> invokeLoader(List<K> keys, List<Object> keyContexts,
     }
 
     CompletableFuture<List<V>> invokeLoader(List<K> keys, List<Object> keyContexts, List<CompletableFuture<V>> queuedFutures) {
+        Object context = loaderOptions.getBatchLoaderContextProvider().getContext();
+        BatchLoaderEnvironment environment = BatchLoaderEnvironment.newBatchLoaderEnvironment()
+                .context(context).keyContexts(keys, keyContexts).build();
+
+        DataLoaderInstrumentationContext<List<?>> instrCtx = ctxOrNoopCtx(instrumentation().beginBatchLoader(dataLoader, keys, environment));
+
         CompletableFuture<List<V>> batchLoad;
         try {
-            Object context = loaderOptions.getBatchLoaderContextProvider().getContext();
-            BatchLoaderEnvironment environment = BatchLoaderEnvironment.newBatchLoaderEnvironment()
-                    .context(context).keyContexts(keys, keyContexts).build();
             if (isMapLoader()) {
                 batchLoad = invokeMapBatchLoader(keys, environment);
             } else if (isPublisher()) {
@@ -441,12 +459,16 @@ CompletableFuture<List<V>> invokeLoader(List<K> keys, List<Object> keyContexts,
             } else {
                 batchLoad = invokeListBatchLoader(keys, environment);
             }
+            instrCtx.onDispatched();
         } catch (Exception e) {
+            instrCtx.onDispatched();
             batchLoad = CompletableFutureKit.failedFuture(e);
         }
+        batchLoad.whenComplete(instrCtx::onCompleted);
         return batchLoad;
     }
 
+
     @SuppressWarnings("unchecked")
     private CompletableFuture<List<V>> invokeListBatchLoader(List<K> keys, BatchLoaderEnvironment environment) {
         CompletionStage<List<V>> loadResult;
@@ -575,6 +597,10 @@ private boolean isMappedPublisher() {
         return batchLoadFunction instanceof MappedBatchPublisher;
     }
 
+    private DataLoaderInstrumentation instrumentation() {
+        return loaderOptions.getInstrumentation();
+    }
+
     int dispatchDepth() {
         synchronized (dataLoader) {
             return loaderQueue.size();

src/main/java/org/dataloader/DataLoaderOptions.java

@@ -18,6 +18,8 @@
 
 import org.dataloader.annotations.PublicApi;
 import org.dataloader.impl.Assertions;
+import org.dataloader.instrumentation.DataLoaderInstrumentation;
+import org.dataloader.instrumentation.DataLoaderInstrumentationHelper;
 import org.dataloader.scheduler.BatchLoaderScheduler;
 import org.dataloader.stats.NoOpStatisticsCollector;
 import org.dataloader.stats.StatisticsCollector;
@@ -48,6 +50,7 @@ public class DataLoaderOptions {
     private BatchLoaderContextProvider environmentProvider;
     private ValueCacheOptions valueCacheOptions;
     private BatchLoaderScheduler batchLoaderScheduler;
+    private DataLoaderInstrumentation instrumentation;
 
     /**
      * Creates a new data loader options with default settings.
@@ -61,6 +64,7 @@ public DataLoaderOptions() {
         environmentProvider = NULL_PROVIDER;
         valueCacheOptions = ValueCacheOptions.newOptions();
         batchLoaderScheduler = null;
+        instrumentation = DataLoaderInstrumentationHelper.NOOP_INSTRUMENTATION;
     }
 
     /**
@@ -80,7 +84,8 @@ public DataLoaderOptions(DataLoaderOptions other) {
         this.statisticsCollector = other.statisticsCollector;
         this.environmentProvider = other.environmentProvider;
         this.valueCacheOptions = other.valueCacheOptions;
-        batchLoaderScheduler = other.batchLoaderScheduler;
+        this.batchLoaderScheduler = other.batchLoaderScheduler;
+        this.instrumentation = other.instrumentation;
     }
 
     /**
@@ -103,7 +108,6 @@ public boolean batchingEnabled() {
      * Sets the option that determines whether batch loading is enabled.
      *
      * @param batchingEnabled {@code true} to enable batch loading, {@code false} otherwise
-     *
      * @return the data loader options for fluent coding
      */
     public DataLoaderOptions setBatchingEnabled(boolean batchingEnabled) {
@@ -124,7 +128,6 @@ public boolean cachingEnabled() {
      * Sets the option that determines whether caching is enabled.
      *
      * @param cachingEnabled {@code true} to enable caching, {@code false} otherwise
-     *
      * @return the data loader options for fluent coding
      */
     public DataLoaderOptions setCachingEnabled(boolean cachingEnabled) {
@@ -134,7 +137,7 @@ public DataLoaderOptions setCachingEnabled(boolean cachingEnabled) {
 
     /**
      * Option that determines whether to cache exceptional values (the default), or not.
-     *
+     * <p>
      * For short-lived caches (that is request caches) it makes sense to cache exceptions since
      * it's likely the key is still poisoned.  However, if you have long-lived caches, then it may make
      * sense to set this to false since the downstream system may have recovered from its failure
@@ -150,7 +153,6 @@ public boolean cachingExceptionsEnabled() {
      * Sets the option that determines whether exceptional values are cache enabled.
      *
      * @param cachingExceptionsEnabled {@code true} to enable caching exceptional values, {@code false} otherwise
-     *
      * @return the data loader options for fluent coding
      */
     public DataLoaderOptions setCachingExceptionsEnabled(boolean cachingExceptionsEnabled) {
@@ -173,7 +175,6 @@ public Optional<CacheKey> cacheKeyFunction() {
      * Sets the function to use for creating the cache key, if caching is enabled.
      *
      * @param cacheKeyFunction the cache key function to use
-     *
      * @return the data loader options for fluent coding
      */
     public DataLoaderOptions setCacheKeyFunction(CacheKey<?> cacheKeyFunction) {
@@ -196,7 +197,6 @@ public DataLoaderOptions setCacheKeyFunction(CacheKey<?> cacheKeyFunction) {
      * Sets the cache map implementation to use for caching, if caching is enabled.
      *
      * @param cacheMap the cache map instance
-     *
      * @return the data loader options for fluent coding
      */
     public DataLoaderOptions setCacheMap(CacheMap<?, ?> cacheMap) {
@@ -219,7 +219,6 @@ public int maxBatchSize() {
      * before they are split into multiple class
      *
      * @param maxBatchSize the maximum batch size
-     *
      * @return the data loader options for fluent coding
      */
     public DataLoaderOptions setMaxBatchSize(int maxBatchSize) {
@@ -240,7 +239,6 @@ public StatisticsCollector getStatisticsCollector() {
      * a common value
      *
      * @param statisticsCollector the statistics collector to use
-     *
      * @return the data loader options for fluent coding
      */
     public DataLoaderOptions setStatisticsCollector(Supplier<StatisticsCollector> statisticsCollector) {
@@ -259,7 +257,6 @@ public BatchLoaderContextProvider getBatchLoaderContextProvider() {
      * Sets the batch loader environment provider that will be used to give context to batch load functions
      *
      * @param contextProvider the batch loader context provider
-     *
      * @return the data loader options for fluent coding
      */
     public DataLoaderOptions setBatchLoaderContextProvider(BatchLoaderContextProvider contextProvider) {
@@ -282,7 +279,6 @@ public DataLoaderOptions setBatchLoaderContextProvider(BatchLoaderContextProvide
      * Sets the value cache implementation to use for caching values, if caching is enabled.
      *
      * @param valueCache the value cache instance
-     *
      * @return the data loader options for fluent coding
      */
     public DataLoaderOptions setValueCache(ValueCache<?, ?> valueCache) {
@@ -301,7 +297,6 @@ public ValueCacheOptions getValueCacheOptions() {
      * Sets the {@link ValueCacheOptions} that control how the {@link ValueCache} will be used
      *
      * @param valueCacheOptions the value cache options
-     *
      * @return the data loader options for fluent coding
      */
     public DataLoaderOptions setValueCacheOptions(ValueCacheOptions valueCacheOptions) {
@@ -321,11 +316,28 @@ public BatchLoaderScheduler getBatchLoaderScheduler() {
      * to some future time.
      *
      * @param batchLoaderScheduler the scheduler
-     *
      * @return the data loader options for fluent coding
      */
     public DataLoaderOptions setBatchLoaderScheduler(BatchLoaderScheduler batchLoaderScheduler) {
         this.batchLoaderScheduler = batchLoaderScheduler;
         return this;
     }
+
+    /**
+     * @return the {@link DataLoaderInstrumentation} to use
+     */
+    public DataLoaderInstrumentation getInstrumentation() {
+        return instrumentation;
+    }
+
+    /**
+     * Sets in a new {@link DataLoaderInstrumentation}
+     *
+     * @param instrumentation the new {@link DataLoaderInstrumentation}
+     * @return the data loader options for fluent coding
+     */
+    public DataLoaderOptions setInstrumentation(DataLoaderInstrumentation instrumentation) {
+        this.instrumentation = nonNull(instrumentation);
+        return this;
+    }
 }

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

@@ -0,0 +1,36 @@
+package org.dataloader.instrumentation;
+
+import org.dataloader.BatchLoaderEnvironment;
+import org.dataloader.DataLoader;
+import org.dataloader.DispatchResult;
+
+import java.util.List;
+
+/**
+ * This interface is called when certain actions happen inside a data loader
+ */
+public interface DataLoaderInstrumentation {
+    /**
+     * 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
+     * @return a DataLoaderInstrumentationContext or null to be more performant
+     */
+    default DataLoaderInstrumentationContext<DispatchResult<?>> beginDispatch(DataLoader<?, ?> dataLoader) {
+        return null;
+    }
+
+    /**
+     * 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
+     * @param keys        the set of keys being fetched
+     * @param environment the {@link BatchLoaderEnvironment}
+     * @return a DataLoaderInstrumentationContext or null to be more performant
+     */
+    default DataLoaderInstrumentationContext<List<?>> beginBatchLoader(DataLoader<?, ?> dataLoader, List<?> keys, BatchLoaderEnvironment environment) {
+        return null;
+    }
+}

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

@@ -0,0 +1,28 @@
+package org.dataloader.instrumentation;
+
+import java.util.concurrent.CompletableFuture;
+
+/**
+ * When a {@link DataLoaderInstrumentation}.'beginXXX()' method is called then it must return a {@link DataLoaderInstrumentationContext}
+ * that will be invoked when the step is first dispatched and then when it completes.  Sometimes this is effectively the same time
+ * whereas at other times it's when an asynchronous {@link CompletableFuture} completes.
+ * <p>
+ * This pattern of construction of an object then call back is intended to allow "timers" to be created that can instrument what has
+ * just happened or "loggers" to be called to record what has happened.
+ */
+public interface DataLoaderInstrumentationContext<T> {
+    /**
+     * This is invoked when the instrumentation step is initially dispatched
+     */
+    default void onDispatched() {
+    }
+
+    /**
+     * 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
+     */
+    default void onCompleted(T result, Throwable t) {
+    }
+}

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

@@ -0,0 +1,34 @@
+package org.dataloader.instrumentation;
+
+public class DataLoaderInstrumentationHelper {
+
+    private static final DataLoaderInstrumentationContext<?> NOOP_CTX = new DataLoaderInstrumentationContext<>() {
+        @Override
+        public void onDispatched() {
+        }
+
+        @Override
+        public void onCompleted(Object result, Throwable t) {
+        }
+    };
+
+    public static <T> DataLoaderInstrumentationContext<T> noOpCtx() {
+        //noinspection unchecked
+        return (DataLoaderInstrumentationContext<T>) NOOP_CTX;
+    }
+
+    public static final DataLoaderInstrumentation NOOP_INSTRUMENTATION = new DataLoaderInstrumentation() {
+    };
+
+    /**
+     * Check the {@link DataLoaderInstrumentationContext} to see if its null and returns a noop if it is or else the original
+     * context
+     *
+     * @param ic  the context in play
+     * @param <T> for two
+     * @return a non null context
+     */
+    public static <T> DataLoaderInstrumentationContext<T> ctxOrNoopCtx(DataLoaderInstrumentationContext<T> ic) {
+        return ic == null ? noOpCtx() : ic;
+    }
+}