From e55b6431fa9ce01dc6c6017da87b2d6ae95a7855 Mon Sep 17 00:00:00 2001 From: bbaker Date: Tue, 21 Jan 2025 10:11:57 +1100 Subject: [PATCH 1/8] Instrumentatin support for dataloader --- .../org/dataloader/DataLoaderFactory.java | 2 +- .../java/org/dataloader/DataLoaderHelper.java | 38 +++++- .../org/dataloader/DataLoaderOptions.java | 38 +++--- .../DataLoaderInstrumentation.java | 36 ++++++ .../DataLoaderInstrumentationContext.java | 28 +++++ .../DataLoaderInstrumentationHelper.java | 34 ++++++ .../DataLoaderInstrumentationTest.java | 108 ++++++++++++++++++ 7 files changed, 264 insertions(+), 20 deletions(-) create mode 100644 src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentation.java create mode 100644 src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentationContext.java create mode 100644 src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentationHelper.java create mode 100644 src/test/java/org/dataloader/instrumentation/DataLoaderInstrumentationTest.java diff --git a/src/main/java/org/dataloader/DataLoaderFactory.java b/src/main/java/org/dataloader/DataLoaderFactory.java index 0dc029a..a87e4eb 100644 --- a/src/main/java/org/dataloader/DataLoaderFactory.java +++ b/src/main/java/org/dataloader/DataLoaderFactory.java @@ -561,7 +561,7 @@ public Builder options(DataLoaderOptions options) { return this; } - DataLoader build() { + public DataLoader build() { return mkDataLoader(batchLoadFunction, options); } } diff --git a/src/main/java/org/dataloader/DataLoaderHelper.java b/src/main/java/org/dataloader/DataLoaderHelper.java index 29701b5..9b5a59e 100644 --- a/src/main/java/org/dataloader/DataLoaderHelper.java +++ b/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 dispatch() { + DataLoaderInstrumentationContext> instrCtx = ctxOrNoopCtx(instrumentation().beginDispatch(dataLoader)); + boolean batchingEnabled = loaderOptions.batchingEnabled(); final List keys; final List callContexts; @@ -175,7 +180,8 @@ DispatchResult 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 dispatch() { lastDispatchTime.set(now()); } if (!batchingEnabled) { - return emptyDispatchResult(); + instrCtx.onDispatched(); + return endDispatchCtx(instrCtx, emptyDispatchResult()); } final int totalEntriesHandled = keys.size(); // @@ -213,7 +220,15 @@ DispatchResult dispatch() { } else { futureList = dispatchQueueBatch(keys, callContexts, queuedFutures); } - return new DispatchResult<>(futureList, totalEntriesHandled); + instrCtx.onDispatched(); + return endDispatchCtx(instrCtx, new DispatchResult<>(futureList, totalEntriesHandled)); + } + + private DispatchResult endDispatchCtx(DataLoaderInstrumentationContext> instrCtx, DispatchResult dispatchResult) { + // once the CF completes, we can tell the instrumentation + dispatchResult.getPromisedResults() + .whenComplete((result, throwable) -> instrCtx.onCompleted(dispatchResult, throwable)); + return dispatchResult; } private CompletableFuture> sliceIntoBatchesOfBatches(List keys, List> queuedFutures, List callContexts, int maxBatchSize) { @@ -427,11 +442,14 @@ CompletableFuture> invokeLoader(List keys, List keyContexts, } CompletableFuture> invokeLoader(List keys, List keyContexts, List> queuedFutures) { + Object context = loaderOptions.getBatchLoaderContextProvider().getContext(); + BatchLoaderEnvironment environment = BatchLoaderEnvironment.newBatchLoaderEnvironment() + .context(context).keyContexts(keys, keyContexts).build(); + + DataLoaderInstrumentationContext> instrCtx = ctxOrNoopCtx(instrumentation().beginBatchLoader(dataLoader, keys, environment)); + CompletableFuture> 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> invokeLoader(List keys, List 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> invokeListBatchLoader(List keys, BatchLoaderEnvironment environment) { CompletionStage> 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(); diff --git a/src/main/java/org/dataloader/DataLoaderOptions.java b/src/main/java/org/dataloader/DataLoaderOptions.java index b96e785..715bcc5 100644 --- a/src/main/java/org/dataloader/DataLoaderOptions.java +++ b/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. - * + *

* 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 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) { @@ -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; + } } diff --git a/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentation.java b/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentation.java new file mode 100644 index 0000000..f0da788 --- /dev/null +++ b/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> 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> beginBatchLoader(DataLoader dataLoader, List keys, BatchLoaderEnvironment environment) { + return null; + } +} diff --git a/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentationContext.java b/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentationContext.java new file mode 100644 index 0000000..ae0bbc1 --- /dev/null +++ b/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. + *

+ * 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 { + /** + * 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) { + } +} diff --git a/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentationHelper.java b/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentationHelper.java new file mode 100644 index 0000000..5ff545d --- /dev/null +++ b/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 DataLoaderInstrumentationContext noOpCtx() { + //noinspection unchecked + return (DataLoaderInstrumentationContext) 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 for two + * @return a non null context + */ + public static DataLoaderInstrumentationContext ctxOrNoopCtx(DataLoaderInstrumentationContext ic) { + return ic == null ? noOpCtx() : ic; + } +} diff --git a/src/test/java/org/dataloader/instrumentation/DataLoaderInstrumentationTest.java b/src/test/java/org/dataloader/instrumentation/DataLoaderInstrumentationTest.java new file mode 100644 index 0000000..2074aef --- /dev/null +++ b/src/test/java/org/dataloader/instrumentation/DataLoaderInstrumentationTest.java @@ -0,0 +1,108 @@ +package org.dataloader.instrumentation; + +import org.dataloader.BatchLoader; +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.junit.jupiter.api.Test; + +import java.util.List; +import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; + +import static org.awaitility.Awaitility.await; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.greaterThan; +import static org.hamcrest.Matchers.is; + +class DataLoaderInstrumentationTest { + + BatchLoader snoozingBatchLoader = keys -> CompletableFuture.supplyAsync(() -> { + TestKit.snooze(100); + return keys; + }); + + @Test + void canMonitorDispatching() { + AtomicLong timer = new AtomicLong(); + AtomicReference> dlRef = new AtomicReference<>(); + + DataLoaderInstrumentation instrumentation = new DataLoaderInstrumentation() { + + @Override + public DataLoaderInstrumentationContext> beginDispatch(DataLoader dataLoader) { + dlRef.set(dataLoader); + + long then = System.currentTimeMillis(); + return new DataLoaderInstrumentationContext<>() { + @Override + public void onCompleted(DispatchResult result, Throwable t) { + timer.set(System.currentTimeMillis() - then); + } + }; + } + + @Override + public DataLoaderInstrumentationContext> beginBatchLoader(DataLoader dataLoader, List keys, BatchLoaderEnvironment environment) { + return DataLoaderInstrumentationHelper.noOpCtx(); + } + }; + + DataLoaderOptions options = new DataLoaderOptions().setInstrumentation(instrumentation).setMaxBatchSize(5); + + DataLoader dl = DataLoaderFactory.newDataLoader(snoozingBatchLoader, options); + + for (int i = 0; i < 20; i++) { + dl.load("X"+ i); + } + + CompletableFuture> dispatch = dl.dispatch(); + + await().until(dispatch::isDone); + assertThat(timer.get(), greaterThan(150L)); // we must have called batch load 4 times + assertThat(dlRef.get(), is(dl)); + } + + @Test + void canMonitorBatchLoading() { + AtomicLong timer = new AtomicLong(); + AtomicReference beRef = new AtomicReference<>(); + AtomicReference> dlRef = new AtomicReference<>(); + + DataLoaderInstrumentation instrumentation = new DataLoaderInstrumentation() { + + @Override + public DataLoaderInstrumentationContext> beginBatchLoader(DataLoader dataLoader, List keys, BatchLoaderEnvironment environment) { + dlRef.set(dataLoader); + beRef.set(environment); + + long then = System.currentTimeMillis(); + return new DataLoaderInstrumentationContext<>() { + @Override + public void onCompleted(List result, Throwable t) { + timer.set(System.currentTimeMillis() - then); + } + }; + } + }; + + DataLoaderOptions options = new DataLoaderOptions().setInstrumentation(instrumentation); + DataLoader dl = DataLoaderFactory.newDataLoader(snoozingBatchLoader, options); + + dl.load("A"); + dl.load("B"); + + CompletableFuture> dispatch = dl.dispatch(); + + await().until(dispatch::isDone); + assertThat(timer.get(), greaterThan(50L)); + assertThat(dlRef.get(), is(dl)); + assertThat(beRef.get().getKeyContexts().keySet(), equalTo(Set.of("A", "B"))); + } +} \ No newline at end of file From 97364ed880009fe09c5c76a1d1624275dd3ad2f8 Mon Sep 17 00:00:00 2001 From: bbaker Date: Tue, 21 Jan 2025 14:33:11 +1100 Subject: [PATCH 2/8] Instrumentatin support for dataloader - better tests --- .../org/dataloader/fixtures/Stopwatch.java | 57 +++++++++++++++++++ .../DataLoaderInstrumentationTest.java | 36 +++++++----- 2 files changed, 79 insertions(+), 14 deletions(-) create mode 100644 src/test/java/org/dataloader/fixtures/Stopwatch.java diff --git a/src/test/java/org/dataloader/fixtures/Stopwatch.java b/src/test/java/org/dataloader/fixtures/Stopwatch.java new file mode 100644 index 0000000..c815a8b --- /dev/null +++ b/src/test/java/org/dataloader/fixtures/Stopwatch.java @@ -0,0 +1,57 @@ +package org.dataloader.fixtures; + +import java.time.Duration; + +public class Stopwatch { + + public static Stopwatch stopwatchStarted() { + return new Stopwatch().start(); + } + + public static Stopwatch stopwatchUnStarted() { + return new Stopwatch(); + } + + private long started = -1; + private long stopped = -1; + + public Stopwatch start() { + synchronized (this) { + if (started != -1) { + throw new IllegalStateException("You have started it before"); + } + started = System.currentTimeMillis(); + } + return this; + } + + private Stopwatch() { + } + + public long elapsed() { + synchronized (this) { + if (started == -1) { + throw new IllegalStateException("You haven't started it"); + } + if (stopped == -1) { + return System.currentTimeMillis() - started; + } else { + return stopped - started; + } + } + } + + public Duration duration() { + return Duration.ofMillis(elapsed()); + } + + public Duration stop() { + synchronized (this) { + if (started != -1) { + throw new IllegalStateException("You have started it"); + } + stopped = System.currentTimeMillis(); + return duration(); + } + } +} diff --git a/src/test/java/org/dataloader/instrumentation/DataLoaderInstrumentationTest.java b/src/test/java/org/dataloader/instrumentation/DataLoaderInstrumentationTest.java index 2074aef..4f43719 100644 --- a/src/test/java/org/dataloader/instrumentation/DataLoaderInstrumentationTest.java +++ b/src/test/java/org/dataloader/instrumentation/DataLoaderInstrumentationTest.java @@ -6,13 +6,14 @@ import org.dataloader.DataLoaderFactory; import org.dataloader.DataLoaderOptions; import org.dataloader.DispatchResult; +import org.dataloader.fixtures.Stopwatch; import org.dataloader.fixtures.TestKit; import org.junit.jupiter.api.Test; +import java.util.ArrayList; import java.util.List; import java.util.Set; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; import static org.awaitility.Awaitility.await; @@ -30,7 +31,7 @@ class DataLoaderInstrumentationTest { @Test void canMonitorDispatching() { - AtomicLong timer = new AtomicLong(); + Stopwatch stopwatch = Stopwatch.stopwatchUnStarted(); AtomicReference> dlRef = new AtomicReference<>(); DataLoaderInstrumentation instrumentation = new DataLoaderInstrumentation() { @@ -38,12 +39,11 @@ void canMonitorDispatching() { @Override public DataLoaderInstrumentationContext> beginDispatch(DataLoader dataLoader) { dlRef.set(dataLoader); - - long then = System.currentTimeMillis(); + stopwatch.start(); return new DataLoaderInstrumentationContext<>() { @Override public void onCompleted(DispatchResult result, Throwable t) { - timer.set(System.currentTimeMillis() - then); + stopwatch.stop(); } }; } @@ -54,24 +54,32 @@ public DataLoaderInstrumentationContext> beginBatchLoader(DataLoader dl = DataLoaderFactory.newDataLoader(snoozingBatchLoader, options); + List keys = new ArrayList<>(); for (int i = 0; i < 20; i++) { - dl.load("X"+ i); + String key = "X" + i; + keys.add(key); + dl.load(key); } CompletableFuture> dispatch = dl.dispatch(); await().until(dispatch::isDone); - assertThat(timer.get(), greaterThan(150L)); // we must have called batch load 4 times + // we must have called batch load 4 times at 100ms snooze per call + // but its in parallel via supplyAsync + assertThat(stopwatch.elapsed(), greaterThan(75L)); assertThat(dlRef.get(), is(dl)); + assertThat(dispatch.join(), equalTo(keys)); } @Test void canMonitorBatchLoading() { - AtomicLong timer = new AtomicLong(); + Stopwatch stopwatch = Stopwatch.stopwatchUnStarted(); AtomicReference beRef = new AtomicReference<>(); AtomicReference> dlRef = new AtomicReference<>(); @@ -82,11 +90,11 @@ public DataLoaderInstrumentationContext> beginBatchLoader(DataLoader() { @Override public void onCompleted(List result, Throwable t) { - timer.set(System.currentTimeMillis() - then); + stopwatch.stop(); } }; } @@ -95,13 +103,13 @@ public void onCompleted(List result, Throwable t) { DataLoaderOptions options = new DataLoaderOptions().setInstrumentation(instrumentation); DataLoader dl = DataLoaderFactory.newDataLoader(snoozingBatchLoader, options); - dl.load("A"); - dl.load("B"); + dl.load("A", "kcA"); + dl.load("B", "kcB"); CompletableFuture> dispatch = dl.dispatch(); await().until(dispatch::isDone); - assertThat(timer.get(), greaterThan(50L)); + assertThat(stopwatch.elapsed(), greaterThan(50L)); assertThat(dlRef.get(), is(dl)); assertThat(beRef.get().getKeyContexts().keySet(), equalTo(Set.of("A", "B"))); } From 411b9bde10a7fe118d52aa9297cee8e56b82b02c Mon Sep 17 00:00:00 2001 From: bbaker Date: Wed, 22 Jan 2025 11:58:29 +1100 Subject: [PATCH 3/8] Instrumentation support for dataloader - better tests and added ChainedDataLoaderInstrumentation --- .../ChainedDataLoaderInstrumentation.java | 86 +++++++++ .../DataLoaderInstrumentation.java | 4 +- .../DataLoaderInstrumentationContext.java | 6 +- .../parameterized/TestDataLoaderFactory.java | 4 + .../ChainedDataLoaderInstrumentationTest.java | 166 ++++++++++++++++++ 5 files changed, 262 insertions(+), 4 deletions(-) create mode 100644 src/main/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentation.java create mode 100644 src/test/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentationTest.java diff --git a/src/main/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentation.java b/src/main/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentation.java new file mode 100644 index 0000000..c80b510 --- /dev/null +++ b/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 instrumentations; + + public ChainedDataLoaderInstrumentation() { + instrumentations = List.of(); + } + + public ChainedDataLoaderInstrumentation(List 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 list = new ArrayList<>(instrumentations); + list.add(instrumentation); + return new ChainedDataLoaderInstrumentation(list); + } + + @Override + public DataLoaderInstrumentationContext> beginDispatch(DataLoader dataLoader) { + return chainedCtx(it -> it.beginDispatch(dataLoader)); + } + + @Override + public DataLoaderInstrumentationContext> beginBatchLoader(DataLoader dataLoader, List keys, BatchLoaderEnvironment environment) { + return chainedCtx(it -> it.beginBatchLoader(dataLoader, keys, environment)); + } + + private DataLoaderInstrumentationContext chainedCtx(Function> 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 List> dropNullContexts(Function> mapper) { + return instrumentations.stream() + .map(mapper) + .filter(Objects::nonNull) + .collect(Collectors.toList()); + } + + private static class ChainedInstrumentationContext implements DataLoaderInstrumentationContext { + private final List> contexts; + + public ChainedInstrumentationContext(List> 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)); + } + } +} diff --git a/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentation.java b/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentation.java index f0da788..5cce9db 100644 --- a/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentation.java +++ b/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> 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 diff --git a/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentationContext.java b/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentationContext.java index ae0bbc1..d327acd 100644 --- a/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentationContext.java +++ b/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentationContext.java @@ -12,13 +12,15 @@ */ public interface DataLoaderInstrumentationContext { /** - * 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 diff --git a/src/test/java/org/dataloader/fixtures/parameterized/TestDataLoaderFactory.java b/src/test/java/org/dataloader/fixtures/parameterized/TestDataLoaderFactory.java index 97e35a8..8cbe86c 100644 --- a/src/test/java/org/dataloader/fixtures/parameterized/TestDataLoaderFactory.java +++ b/src/test/java/org/dataloader/fixtures/parameterized/TestDataLoaderFactory.java @@ -25,6 +25,10 @@ public interface TestDataLoaderFactory { // Convenience methods + default DataLoader idLoader(DataLoaderOptions options) { + return idLoader(options, new ArrayList<>()); + } + default DataLoader idLoader(List> calls) { return idLoader(null, calls); } diff --git a/src/test/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentationTest.java b/src/test/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentationTest.java new file mode 100644 index 0000000..93c0edc --- /dev/null +++ b/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 methods = new ArrayList<>(); + + public CapturingInstrumentation(String name) { + this.name = name; + } + + @Override + public DataLoaderInstrumentationContext> 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> 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> beginDispatch(DataLoader dataLoader) { + methods.add(name + "_beginDispatch"); + return null; + } + + @Override + public DataLoaderInstrumentationContext> 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 dl = DataLoaderFactory.newDataLoader(TestKit.keysAsValues(), options); + + dl.load("A"); + dl.load("B"); + + CompletableFuture> 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 dl = DataLoaderFactory.newDataLoader(TestKit.keysAsValues(), options); + + dl.load("A"); + dl.load("B"); + + CompletableFuture> 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 dl = factory.idLoader(options); + + dl.load("A"); + dl.load("B"); + + CompletableFuture> 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"))); + } +} \ No newline at end of file From 3a8adcd87402427f68bc6e3e93248d0cc1640c04 Mon Sep 17 00:00:00 2001 From: bbaker Date: Wed, 22 Jan 2025 14:33:08 +1100 Subject: [PATCH 4/8] Instrumentation support for dataloader -added registry instrumentation --- .../org/dataloader/DataLoaderRegistry.java | 82 ++++++- .../ChainedDataLoaderInstrumentation.java | 34 ++- .../DataLoaderInstrumentation.java | 2 + .../DataLoaderInstrumentationContext.java | 3 + .../DataLoaderInstrumentationHelper.java | 15 +- .../CapturingInstrumentation.java | 49 +++++ .../CapturingInstrumentationReturnsNull.java | 26 +++ .../ChainedDataLoaderInstrumentationTest.java | 80 ++----- ...DataLoaderRegistryInstrumentationTest.java | 205 ++++++++++++++++++ 9 files changed, 421 insertions(+), 75 deletions(-) create mode 100644 src/test/java/org/dataloader/instrumentation/CapturingInstrumentation.java create mode 100644 src/test/java/org/dataloader/instrumentation/CapturingInstrumentationReturnsNull.java create mode 100644 src/test/java/org/dataloader/instrumentation/DataLoaderRegistryInstrumentationTest.java diff --git a/src/main/java/org/dataloader/DataLoaderRegistry.java b/src/main/java/org/dataloader/DataLoaderRegistry.java index 5a3f90f..5e00ee1 100644 --- a/src/main/java/org/dataloader/DataLoaderRegistry.java +++ b/src/main/java/org/dataloader/DataLoaderRegistry.java @@ -1,6 +1,9 @@ package org.dataloader; import org.dataloader.annotations.PublicApi; +import org.dataloader.instrumentation.ChainedDataLoaderInstrumentation; +import org.dataloader.instrumentation.DataLoaderInstrumentation; +import org.dataloader.instrumentation.DataLoaderInstrumentationHelper; import org.dataloader.stats.Statistics; import java.util.ArrayList; @@ -21,25 +24,81 @@ @PublicApi public class DataLoaderRegistry { protected final Map> dataLoaders = new ConcurrentHashMap<>(); + protected final DataLoaderInstrumentation instrumentation; + public DataLoaderRegistry() { + instrumentation = null; } private DataLoaderRegistry(Builder builder) { - this.dataLoaders.putAll(builder.dataLoaders); + instrument(builder.instrumentation, builder.dataLoaders); + this.instrumentation = builder.instrumentation; + } + + private void instrument(DataLoaderInstrumentation registryInstrumentation, Map> incomingDataLoaders) { + this.dataLoaders.putAll(incomingDataLoaders); + if (registryInstrumentation != null) { + this.dataLoaders.replaceAll((k, existingDL) -> instrumentDL(registryInstrumentation, existingDL)); + } + } + + /** + * Can be called to tweak a {@link DataLoader} so that it has the registry {@link DataLoaderInstrumentation} added as the first one. + * + * @param registryInstrumentation the common registry {@link DataLoaderInstrumentation} + * @param existingDL the existing data loader + * @return a new {@link DataLoader} or the same one if there is nothing to change + */ + private static DataLoader instrumentDL(DataLoaderInstrumentation registryInstrumentation, DataLoader existingDL) { + if (registryInstrumentation == null) { + return existingDL; + } + DataLoaderOptions options = existingDL.getOptions(); + DataLoaderInstrumentation existingInstrumentation = options.getInstrumentation(); + // if they have any instrumentations then add to it + if (existingInstrumentation != null) { + if (existingInstrumentation == registryInstrumentation) { + // nothing to change + return existingDL; + } + if (existingInstrumentation == DataLoaderInstrumentationHelper.NOOP_INSTRUMENTATION) { + // replace it with the registry one + return mkInstrumentedDataLoader(existingDL, options, registryInstrumentation); + } + if (existingInstrumentation instanceof ChainedDataLoaderInstrumentation) { + // avoids calling a chained inside a chained + DataLoaderInstrumentation newInstrumentation = ((ChainedDataLoaderInstrumentation) existingInstrumentation).prepend(registryInstrumentation); + return mkInstrumentedDataLoader(existingDL, options, newInstrumentation); + } else { + DataLoaderInstrumentation newInstrumentation = new ChainedDataLoaderInstrumentation().add(registryInstrumentation).add(existingInstrumentation); + return mkInstrumentedDataLoader(existingDL, options, newInstrumentation); + } + } else { + return mkInstrumentedDataLoader(existingDL, options, registryInstrumentation); + } + } + + private static DataLoader mkInstrumentedDataLoader(DataLoader existingDL, DataLoaderOptions options, DataLoaderInstrumentation newInstrumentation) { + return existingDL.transform(builder -> { + options.setInstrumentation(newInstrumentation); + builder.options(options); + }); } + public DataLoaderInstrumentation getInstrumentation() { + return instrumentation; + } /** * This will register a new dataloader * * @param key the key to put the data loader under * @param dataLoader the data loader to register - * * @return this registry */ public DataLoaderRegistry register(String key, DataLoader dataLoader) { - dataLoaders.put(key, dataLoader); + dataLoaders.put(key, instrumentDL(instrumentation, dataLoader)); return this; } @@ -54,13 +113,15 @@ public DataLoaderRegistry register(String key, DataLoader dataLoader) { * @param mappingFunction the function to compute a data loader * @param the type of keys * @param the type of values - * * @return a data loader */ @SuppressWarnings("unchecked") public DataLoader computeIfAbsent(final String key, final Function> mappingFunction) { - return (DataLoader) dataLoaders.computeIfAbsent(key, mappingFunction); + return (DataLoader) dataLoaders.computeIfAbsent(key, (k) -> { + DataLoader dl = mappingFunction.apply(k); + return instrumentDL(instrumentation, dl); + }); } /** @@ -68,7 +129,6 @@ public DataLoader computeIfAbsent(final String key, * and return a new combined registry * * @param registry the registry to combine into this registry - * * @return a new combined registry */ public DataLoaderRegistry combine(DataLoaderRegistry registry) { @@ -97,7 +157,6 @@ public DataLoaderRegistry combine(DataLoaderRegistry registry) { * This will unregister a new dataloader * * @param key the key of the data loader to unregister - * * @return this registry */ public DataLoaderRegistry unregister(String key) { @@ -111,7 +170,6 @@ public DataLoaderRegistry unregister(String key) { * @param key the key of the data loader * @param the type of keys * @param the type of values - * * @return a data loader or null if its not present */ @SuppressWarnings("unchecked") @@ -182,13 +240,13 @@ public static Builder newRegistry() { public static class Builder { private final Map> dataLoaders = new HashMap<>(); + private DataLoaderInstrumentation instrumentation; /** * This will register a new dataloader * * @param key the key to put the data loader under * @param dataLoader the data loader to register - * * @return this builder for a fluent pattern */ public Builder register(String key, DataLoader dataLoader) { @@ -201,7 +259,6 @@ public Builder register(String key, DataLoader dataLoader) { * from a previous {@link DataLoaderRegistry} * * @param otherRegistry the previous {@link DataLoaderRegistry} - * * @return this builder for a fluent pattern */ public Builder registerAll(DataLoaderRegistry otherRegistry) { @@ -209,6 +266,11 @@ public Builder registerAll(DataLoaderRegistry otherRegistry) { return this; } + public Builder instrumentation(DataLoaderInstrumentation instrumentation) { + this.instrumentation = instrumentation; + return this; + } + /** * @return the newly built {@link DataLoaderRegistry} */ diff --git a/src/main/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentation.java b/src/main/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentation.java index c80b510..eb9af0a 100644 --- a/src/main/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentation.java +++ b/src/main/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentation.java @@ -3,8 +3,10 @@ import org.dataloader.BatchLoaderEnvironment; import org.dataloader.DataLoader; import org.dataloader.DispatchResult; +import org.dataloader.annotations.PublicApi; import java.util.ArrayList; +import java.util.Collection; import java.util.List; import java.util.Objects; import java.util.function.Function; @@ -14,6 +16,7 @@ * This {@link DataLoaderInstrumentation} can chain together multiple instrumentations and have them all called in * the order of the provided list. */ +@PublicApi public class ChainedDataLoaderInstrumentation implements DataLoaderInstrumentation { private final List instrumentations; @@ -25,6 +28,10 @@ public ChainedDataLoaderInstrumentation(List instrume this.instrumentations = List.copyOf(instrumentations); } + public List getInstrumentations() { + return instrumentations; + } + /** * Adds a new {@link DataLoaderInstrumentation} to the list and creates a new {@link ChainedDataLoaderInstrumentation} * @@ -32,8 +39,33 @@ public ChainedDataLoaderInstrumentation(List instrume * @return a new ChainedDataLoaderInstrumentation object */ public ChainedDataLoaderInstrumentation add(DataLoaderInstrumentation instrumentation) { - ArrayList list = new ArrayList<>(instrumentations); + ArrayList list = new ArrayList<>(this.instrumentations); + list.add(instrumentation); + return new ChainedDataLoaderInstrumentation(list); + } + + /** + * Prepends 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 prepend(DataLoaderInstrumentation instrumentation) { + ArrayList list = new ArrayList<>(); list.add(instrumentation); + list.addAll(this.instrumentations); + return new ChainedDataLoaderInstrumentation(list); + } + + /** + * Adds a collection of {@link DataLoaderInstrumentation} to the list and creates a new {@link ChainedDataLoaderInstrumentation} + * + * @param instrumentations the new ones to add + * @return a new ChainedDataLoaderInstrumentation object + */ + public ChainedDataLoaderInstrumentation addAll(Collection instrumentations) { + ArrayList list = new ArrayList<>(this.instrumentations); + list.addAll(instrumentations); return new ChainedDataLoaderInstrumentation(list); } diff --git a/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentation.java b/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentation.java index 5cce9db..78f2cf5 100644 --- a/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentation.java +++ b/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentation.java @@ -3,12 +3,14 @@ import org.dataloader.BatchLoaderEnvironment; import org.dataloader.DataLoader; import org.dataloader.DispatchResult; +import org.dataloader.annotations.PublicSpi; import java.util.List; /** * This interface is called when certain actions happen inside a data loader */ +@PublicSpi public interface DataLoaderInstrumentation { /** * This call back is done just before the {@link DataLoader#dispatch()} is invoked, diff --git a/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentationContext.java b/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentationContext.java index d327acd..88b08ef 100644 --- a/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentationContext.java +++ b/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentationContext.java @@ -1,5 +1,7 @@ package org.dataloader.instrumentation; +import org.dataloader.annotations.PublicSpi; + import java.util.concurrent.CompletableFuture; /** @@ -10,6 +12,7 @@ * 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. */ +@PublicSpi public interface DataLoaderInstrumentationContext { /** * This is invoked when the instrumentation step is initially dispatched. Note this is NOT diff --git a/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentationHelper.java b/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentationHelper.java index 5ff545d..844ec37 100644 --- a/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentationHelper.java +++ b/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentationHelper.java @@ -1,7 +1,11 @@ package org.dataloader.instrumentation; +import org.dataloader.annotations.PublicApi; + +@PublicApi public class DataLoaderInstrumentationHelper { + @SuppressWarnings("RedundantMethodOverride") private static final DataLoaderInstrumentationContext NOOP_CTX = new DataLoaderInstrumentationContext<>() { @Override public void onDispatched() { @@ -12,17 +16,26 @@ public void onCompleted(Object result, Throwable t) { } }; + /** + * Returns a noop {@link DataLoaderInstrumentationContext} of the right type + * + * @param for two + * @return a noop context + */ public static DataLoaderInstrumentationContext noOpCtx() { //noinspection unchecked return (DataLoaderInstrumentationContext) NOOP_CTX; } + /** + * A well known noop {@link DataLoaderInstrumentation} + */ 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 + * context. This is a bit of a helper method. * * @param ic the context in play * @param for two diff --git a/src/test/java/org/dataloader/instrumentation/CapturingInstrumentation.java b/src/test/java/org/dataloader/instrumentation/CapturingInstrumentation.java new file mode 100644 index 0000000..f5af683 --- /dev/null +++ b/src/test/java/org/dataloader/instrumentation/CapturingInstrumentation.java @@ -0,0 +1,49 @@ +package org.dataloader.instrumentation; + +import org.dataloader.BatchLoaderEnvironment; +import org.dataloader.DataLoader; +import org.dataloader.DispatchResult; + +import java.util.ArrayList; +import java.util.List; + +class CapturingInstrumentation implements DataLoaderInstrumentation { + String name; + List methods = new ArrayList<>(); + + public CapturingInstrumentation(String name) { + this.name = name; + } + + @Override + public DataLoaderInstrumentationContext> 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> 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"); + } + }; + } +} diff --git a/src/test/java/org/dataloader/instrumentation/CapturingInstrumentationReturnsNull.java b/src/test/java/org/dataloader/instrumentation/CapturingInstrumentationReturnsNull.java new file mode 100644 index 0000000..0c16429 --- /dev/null +++ b/src/test/java/org/dataloader/instrumentation/CapturingInstrumentationReturnsNull.java @@ -0,0 +1,26 @@ +package org.dataloader.instrumentation; + +import org.dataloader.BatchLoaderEnvironment; +import org.dataloader.DataLoader; +import org.dataloader.DispatchResult; + +import java.util.List; + +class CapturingInstrumentationReturnsNull extends CapturingInstrumentation { + + public CapturingInstrumentationReturnsNull(String name) { + super(name); + } + + @Override + public DataLoaderInstrumentationContext> beginDispatch(DataLoader dataLoader) { + methods.add(name + "_beginDispatch"); + return null; + } + + @Override + public DataLoaderInstrumentationContext> beginBatchLoader(DataLoader dataLoader, List keys, BatchLoaderEnvironment environment) { + methods.add(name + "_beginBatchLoader"); + return null; + } +} diff --git a/src/test/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentationTest.java b/src/test/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentationTest.java index 93c0edc..3168734 100644 --- a/src/test/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentationTest.java +++ b/src/test/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentationTest.java @@ -1,17 +1,15 @@ 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.BeforeEach; 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; @@ -22,65 +20,16 @@ class ChainedDataLoaderInstrumentationTest { - static class CapturingInstrumentation implements DataLoaderInstrumentation { - String name; - List methods = new ArrayList<>(); - - public CapturingInstrumentation(String name) { - this.name = name; - } - - @Override - public DataLoaderInstrumentationContext> 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> 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 { + CapturingInstrumentation capturingA; + CapturingInstrumentation capturingB; + CapturingInstrumentation capturingButReturnsNull; - public CapturingInstrumentationReturnsNull(String name) { - super(name); - } - @Override - public DataLoaderInstrumentationContext> beginDispatch(DataLoader dataLoader) { - methods.add(name + "_beginDispatch"); - return null; - } - - @Override - public DataLoaderInstrumentationContext> beginBatchLoader(DataLoader dataLoader, List keys, BatchLoaderEnvironment environment) { - methods.add(name + "_beginBatchLoader"); - return null; - } + @BeforeEach + void setUp() { + capturingA = new CapturingInstrumentation("A"); + capturingB = new CapturingInstrumentation("B"); + capturingButReturnsNull = new CapturingInstrumentationReturnsNull("NULL"); } @Test @@ -128,9 +77,6 @@ void canChainTogetherOneInstrumentation() { @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) @@ -163,4 +109,12 @@ public void canChainTogetherManyInstrumentationsWithDifferentBatchLoaders(TestDa // it returned null on all its contexts - nothing to call back on assertThat(capturingButReturnsNull.methods, equalTo(List.of("NULL_beginDispatch", "NULL_beginBatchLoader"))); } + + @Test + void addition_works() { + ChainedDataLoaderInstrumentation chainedItn = new ChainedDataLoaderInstrumentation() + .add(capturingA).prepend(capturingB).addAll(List.of(capturingButReturnsNull)); + + assertThat(chainedItn.getInstrumentations(), equalTo(List.of(capturingB, capturingA, capturingButReturnsNull))); + } } \ No newline at end of file diff --git a/src/test/java/org/dataloader/instrumentation/DataLoaderRegistryInstrumentationTest.java b/src/test/java/org/dataloader/instrumentation/DataLoaderRegistryInstrumentationTest.java new file mode 100644 index 0000000..639ff48 --- /dev/null +++ b/src/test/java/org/dataloader/instrumentation/DataLoaderRegistryInstrumentationTest.java @@ -0,0 +1,205 @@ +package org.dataloader.instrumentation; + +import org.dataloader.DataLoader; +import org.dataloader.DataLoaderRegistry; +import org.dataloader.fixtures.TestKit; +import org.dataloader.fixtures.parameterized.TestDataLoaderFactory; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; + +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; + +import static org.awaitility.Awaitility.await; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.instanceOf; + +class DataLoaderRegistryInstrumentationTest { + DataLoader dlX; + DataLoader dlY; + DataLoader dlZ; + + CapturingInstrumentation instrA; + CapturingInstrumentation instrB; + ChainedDataLoaderInstrumentation chainedInstrA; + ChainedDataLoaderInstrumentation chainedInstrB; + + @BeforeEach + void setUp() { + dlX = TestKit.idLoader(); + dlY = TestKit.idLoader(); + dlZ = TestKit.idLoader(); + instrA = new CapturingInstrumentation("A"); + instrB = new CapturingInstrumentation("B"); + chainedInstrA = new ChainedDataLoaderInstrumentation().add(instrA); + chainedInstrB = new ChainedDataLoaderInstrumentation().add(instrB); + } + + @Test + void canInstrumentRegisteredDLsViaBuilder() { + + assertThat(dlX.getOptions().getInstrumentation(), equalTo(DataLoaderInstrumentationHelper.NOOP_INSTRUMENTATION)); + + DataLoaderRegistry registry = DataLoaderRegistry.newRegistry() + .instrumentation(chainedInstrA) + .register("X", dlX) + .register("Y", dlY) + .register("Z", dlZ) + .build(); + + assertThat(registry.getInstrumentation(), equalTo(chainedInstrA)); + + for (String key : List.of("X", "Y", "Z")) { + DataLoaderInstrumentation instrumentation = registry.getDataLoader(key).getOptions().getInstrumentation(); + assertThat(instrumentation, instanceOf(ChainedDataLoaderInstrumentation.class)); + List instrumentations = ((ChainedDataLoaderInstrumentation) instrumentation).getInstrumentations(); + assertThat(instrumentations, equalTo(List.of(instrA))); + } + } + + @Test + void canInstrumentRegisteredDLsViaBuilderCombined() { + + DataLoaderRegistry registry1 = DataLoaderRegistry.newRegistry() + .register("X", dlX) + .register("Y", dlY) + .build(); + + DataLoaderRegistry registry = DataLoaderRegistry.newRegistry() + .instrumentation(chainedInstrA) + .register("Z", dlZ) + .registerAll(registry1) + .build(); + + for (String key : List.of("X", "Y", "Z")) { + DataLoaderInstrumentation instrumentation = registry.getDataLoader(key).getOptions().getInstrumentation(); + assertThat(instrumentation, instanceOf(ChainedDataLoaderInstrumentation.class)); + List instrumentations = ((ChainedDataLoaderInstrumentation) instrumentation).getInstrumentations(); + assertThat(instrumentations, equalTo(List.of(instrA))); + } + } + + @Test + void canInstrumentViaMutativeRegistration() { + + DataLoaderRegistry registry = DataLoaderRegistry.newRegistry() + .instrumentation(chainedInstrA) + .build(); + + registry.register("X", dlX); + registry.computeIfAbsent("Y", l -> dlY); + registry.computeIfAbsent("Z", l -> dlZ); + + for (String key : List.of("X", "Y", "Z")) { + DataLoaderInstrumentation instrumentation = registry.getDataLoader(key).getOptions().getInstrumentation(); + assertThat(instrumentation, instanceOf(ChainedDataLoaderInstrumentation.class)); + List instrumentations = ((ChainedDataLoaderInstrumentation) instrumentation).getInstrumentations(); + assertThat(instrumentations, equalTo(List.of(instrA))); + } + } + + @Test + void wontDoAnyThingIfThereIsNoRegistryInstrumentation() { + DataLoaderRegistry registry = DataLoaderRegistry.newRegistry() + .register("X", dlX) + .register("Y", dlY) + .register("Z", dlZ) + .build(); + + for (String key : List.of("X", "Y", "Z")) { + DataLoaderInstrumentation instrumentation = registry.getDataLoader(key).getOptions().getInstrumentation(); + assertThat(instrumentation, equalTo(DataLoaderInstrumentationHelper.NOOP_INSTRUMENTATION)); + } + } + + @Test + void wontDoAnyThingIfThereTheyAreTheSameInstrumentationAlready() { + DataLoader newX = dlX.transform(builder -> dlX.getOptions().setInstrumentation(instrA)); + DataLoader newY = dlX.transform(builder -> dlY.getOptions().setInstrumentation(instrA)); + DataLoader newZ = dlX.transform(builder -> dlY.getOptions().setInstrumentation(instrA)); + DataLoaderRegistry registry = DataLoaderRegistry.newRegistry() + .instrumentation(instrA) + .register("X", newX) + .register("Y", newY) + .register("Z", newZ) + .build(); + + Map> dls = Map.of("X", newX, "Y", newY, "Z", newZ); + + assertThat(registry.getInstrumentation(), equalTo(instrA)); + + for (String key : List.of("X", "Y", "Z")) { + DataLoader dataLoader = registry.getDataLoader(key); + DataLoaderInstrumentation instrumentation = dataLoader.getOptions().getInstrumentation(); + assertThat(instrumentation, equalTo(instrA)); + // it's the same DL - it's not changed because it has the same instrumentation + assertThat(dls.get(key), equalTo(dataLoader)); + } + } + + @Test + void ifTheDLHasAInstrumentationThenItsTurnedIntoAChainedOne() { + DataLoader newX = dlX.transform(builder -> dlX.getOptions().setInstrumentation(instrA)); + + DataLoaderRegistry registry = DataLoaderRegistry.newRegistry() + .instrumentation(instrB) + .register("X", newX) + .build(); + + DataLoader dataLoader = registry.getDataLoader("X"); + DataLoaderInstrumentation instrumentation = dataLoader.getOptions().getInstrumentation(); + assertThat(instrumentation, instanceOf(ChainedDataLoaderInstrumentation.class)); + + List instrumentations = ((ChainedDataLoaderInstrumentation) instrumentation).getInstrumentations(); + // it gets turned into a chained one and the registry one goes first + assertThat(instrumentations, equalTo(List.of(instrB, instrA))); + } + + @Test + void chainedInstrumentationsWillBeCombined() { + DataLoader newX = dlX.transform(builder -> builder.options(dlX.getOptions().setInstrumentation(chainedInstrB))); + + DataLoaderRegistry registry = DataLoaderRegistry.newRegistry() + .instrumentation(instrA) + .register("X", newX) + .build(); + + DataLoader dataLoader = registry.getDataLoader("X"); + DataLoaderInstrumentation instrumentation = dataLoader.getOptions().getInstrumentation(); + assertThat(instrumentation, instanceOf(ChainedDataLoaderInstrumentation.class)); + + List instrumentations = ((ChainedDataLoaderInstrumentation) instrumentation).getInstrumentations(); + // it gets turned into a chained one and the registry one goes first + assertThat(instrumentations, equalTo(List.of(instrA, instrB))); + } + + @ParameterizedTest + @MethodSource("org.dataloader.fixtures.parameterized.TestDataLoaderFactories#get") + public void endToEndIntegrationTest(TestDataLoaderFactory factory) { + DataLoader dl = factory.idLoader(); + + DataLoaderRegistry registry = DataLoaderRegistry.newRegistry() + .instrumentation(instrA) + .register("X", dl) + .build(); + + // since the data-loader changed when registered you MUST get the data loader from the registry + // not direct to the old one + DataLoader dataLoader = registry.getDataLoader("X"); + CompletableFuture loadA = dataLoader.load("A"); + + registry.dispatchAll(); + + await().until(loadA::isDone); + assertThat(loadA.join(), equalTo("A")); + + assertThat(instrA.methods, equalTo(List.of("A_beginDispatch", + "A_beginBatchLoader", "A_beginBatchLoader_onDispatched", "A_beginBatchLoader_onCompleted", + "A_beginDispatch_onDispatched", "A_beginDispatch_onCompleted"))); + + } +} \ No newline at end of file From 292b2835b78d745bb755b4222ae98ab4e7c84cd5 Mon Sep 17 00:00:00 2001 From: bbaker Date: Wed, 22 Jan 2025 14:50:45 +1100 Subject: [PATCH 5/8] Instrumentation support for dataloader - updated ScheduledDataLoaderRegistry --- .../org/dataloader/DataLoaderRegistry.java | 33 ++++++++++++++----- .../ScheduledDataLoaderRegistry.java | 11 +++++-- .../ChainedDataLoaderInstrumentationTest.java | 2 +- .../DataLoaderInstrumentationTest.java | 2 +- ...DataLoaderRegistryInstrumentationTest.java | 27 +++++++++++++-- 5 files changed, 61 insertions(+), 14 deletions(-) diff --git a/src/main/java/org/dataloader/DataLoaderRegistry.java b/src/main/java/org/dataloader/DataLoaderRegistry.java index 5e00ee1..d54ee53 100644 --- a/src/main/java/org/dataloader/DataLoaderRegistry.java +++ b/src/main/java/org/dataloader/DataLoaderRegistry.java @@ -19,28 +19,45 @@ /** * This allows data loaders to be registered together into a single place, so * they can be dispatched as one. It also allows you to retrieve data loaders by - * name from a central place + * name from a central place. + *

+ * Notes on {@link DataLoaderInstrumentation} : A {@link DataLoaderRegistry} can have an instrumentation + * associated with it. As each {@link DataLoader} is added to the registry, the {@link DataLoaderInstrumentation} + * of the registry is applied to that {@link DataLoader}. + *

+ * The {@link DataLoader} is changed and hence the object in the registry is not the + * same one as was originally registered. So you MUST get access to the {@link DataLoader} via {@link DataLoaderRegistry#getDataLoader(String)} methods + * and not use the original {@link DataLoader} object. + *

+ * If the {@link DataLoader} has no {@link DataLoaderInstrumentation} then the registry one is added to it. If it does have one already + * then a {@link ChainedDataLoaderInstrumentation} is created with the registry {@link DataLoaderInstrumentation} in it first and then any other + * {@link DataLoaderInstrumentation}s added after that. */ @PublicApi public class DataLoaderRegistry { - protected final Map> dataLoaders = new ConcurrentHashMap<>(); + protected final Map> dataLoaders; protected final DataLoaderInstrumentation instrumentation; public DataLoaderRegistry() { - instrumentation = null; + this(new ConcurrentHashMap<>(),null); } private DataLoaderRegistry(Builder builder) { - instrument(builder.instrumentation, builder.dataLoaders); - this.instrumentation = builder.instrumentation; + this(builder.dataLoaders,builder.instrumentation); } - private void instrument(DataLoaderInstrumentation registryInstrumentation, Map> incomingDataLoaders) { - this.dataLoaders.putAll(incomingDataLoaders); + protected DataLoaderRegistry(Map> dataLoaders, DataLoaderInstrumentation instrumentation ) { + this.dataLoaders = instrumentDLs(dataLoaders, instrumentation); + this.instrumentation = instrumentation; + } + + private Map> instrumentDLs(Map> incomingDataLoaders, DataLoaderInstrumentation registryInstrumentation) { + Map> dataLoaders = new ConcurrentHashMap<>(incomingDataLoaders); if (registryInstrumentation != null) { - this.dataLoaders.replaceAll((k, existingDL) -> instrumentDL(registryInstrumentation, existingDL)); + dataLoaders.replaceAll((k, existingDL) -> instrumentDL(registryInstrumentation, existingDL)); } + return dataLoaders; } /** diff --git a/src/main/java/org/dataloader/registries/ScheduledDataLoaderRegistry.java b/src/main/java/org/dataloader/registries/ScheduledDataLoaderRegistry.java index 6ea9425..b6bc257 100644 --- a/src/main/java/org/dataloader/registries/ScheduledDataLoaderRegistry.java +++ b/src/main/java/org/dataloader/registries/ScheduledDataLoaderRegistry.java @@ -3,6 +3,7 @@ import org.dataloader.DataLoader; import org.dataloader.DataLoaderRegistry; import org.dataloader.annotations.ExperimentalApi; +import org.dataloader.instrumentation.DataLoaderInstrumentation; import java.time.Duration; import java.util.LinkedHashMap; @@ -64,8 +65,7 @@ public class ScheduledDataLoaderRegistry extends DataLoaderRegistry implements A private volatile boolean closed; private ScheduledDataLoaderRegistry(Builder builder) { - super(); - this.dataLoaders.putAll(builder.dataLoaders); + super(builder.dataLoaders, builder.instrumentation); this.scheduledExecutorService = builder.scheduledExecutorService; this.defaultExecutorUsed = builder.defaultExecutorUsed; this.schedule = builder.schedule; @@ -271,6 +271,8 @@ public static class Builder { private boolean defaultExecutorUsed = false; private Duration schedule = Duration.ofMillis(10); private boolean tickerMode = false; + private DataLoaderInstrumentation instrumentation; + /** * If you provide a {@link ScheduledExecutorService} then it will NOT be shutdown when @@ -363,6 +365,11 @@ public Builder tickerMode(boolean tickerMode) { return this; } + public Builder instrumentation(DataLoaderInstrumentation instrumentation) { + this.instrumentation = instrumentation; + return this; + } + /** * @return the newly built {@link ScheduledDataLoaderRegistry} */ diff --git a/src/test/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentationTest.java b/src/test/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentationTest.java index 3168734..b3272ce 100644 --- a/src/test/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentationTest.java +++ b/src/test/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentationTest.java @@ -18,7 +18,7 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.equalTo; -class ChainedDataLoaderInstrumentationTest { +public class ChainedDataLoaderInstrumentationTest { CapturingInstrumentation capturingA; CapturingInstrumentation capturingB; diff --git a/src/test/java/org/dataloader/instrumentation/DataLoaderInstrumentationTest.java b/src/test/java/org/dataloader/instrumentation/DataLoaderInstrumentationTest.java index 4f43719..a35e13a 100644 --- a/src/test/java/org/dataloader/instrumentation/DataLoaderInstrumentationTest.java +++ b/src/test/java/org/dataloader/instrumentation/DataLoaderInstrumentationTest.java @@ -22,7 +22,7 @@ import static org.hamcrest.Matchers.greaterThan; import static org.hamcrest.Matchers.is; -class DataLoaderInstrumentationTest { +public class DataLoaderInstrumentationTest { BatchLoader snoozingBatchLoader = keys -> CompletableFuture.supplyAsync(() -> { TestKit.snooze(100); diff --git a/src/test/java/org/dataloader/instrumentation/DataLoaderRegistryInstrumentationTest.java b/src/test/java/org/dataloader/instrumentation/DataLoaderRegistryInstrumentationTest.java index 639ff48..5690bdc 100644 --- a/src/test/java/org/dataloader/instrumentation/DataLoaderRegistryInstrumentationTest.java +++ b/src/test/java/org/dataloader/instrumentation/DataLoaderRegistryInstrumentationTest.java @@ -4,6 +4,7 @@ import org.dataloader.DataLoaderRegistry; import org.dataloader.fixtures.TestKit; import org.dataloader.fixtures.parameterized.TestDataLoaderFactory; +import org.dataloader.registries.ScheduledDataLoaderRegistry; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; @@ -18,7 +19,7 @@ import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.instanceOf; -class DataLoaderRegistryInstrumentationTest { +public class DataLoaderRegistryInstrumentationTest { DataLoader dlX; DataLoader dlY; DataLoader dlZ; @@ -177,6 +178,29 @@ void chainedInstrumentationsWillBeCombined() { assertThat(instrumentations, equalTo(List.of(instrA, instrB))); } + @SuppressWarnings("resource") + @Test + void canInstrumentScheduledRegistryViaBuilder() { + + assertThat(dlX.getOptions().getInstrumentation(), equalTo(DataLoaderInstrumentationHelper.NOOP_INSTRUMENTATION)); + + ScheduledDataLoaderRegistry registry = ScheduledDataLoaderRegistry.newScheduledRegistry() + .instrumentation(chainedInstrA) + .register("X", dlX) + .register("Y", dlY) + .register("Z", dlZ) + .build(); + + assertThat(registry.getInstrumentation(), equalTo(chainedInstrA)); + + for (String key : List.of("X", "Y", "Z")) { + DataLoaderInstrumentation instrumentation = registry.getDataLoader(key).getOptions().getInstrumentation(); + assertThat(instrumentation, instanceOf(ChainedDataLoaderInstrumentation.class)); + List instrumentations = ((ChainedDataLoaderInstrumentation) instrumentation).getInstrumentations(); + assertThat(instrumentations, equalTo(List.of(instrA))); + } + } + @ParameterizedTest @MethodSource("org.dataloader.fixtures.parameterized.TestDataLoaderFactories#get") public void endToEndIntegrationTest(TestDataLoaderFactory factory) { @@ -200,6 +224,5 @@ public void endToEndIntegrationTest(TestDataLoaderFactory factory) { assertThat(instrA.methods, equalTo(List.of("A_beginDispatch", "A_beginBatchLoader", "A_beginBatchLoader_onDispatched", "A_beginBatchLoader_onCompleted", "A_beginDispatch_onDispatched", "A_beginDispatch_onCompleted"))); - } } \ No newline at end of file From 0b9fba57a316e01c0f6bfb56156e92913edf73bb Mon Sep 17 00:00:00 2001 From: bbaker Date: Wed, 22 Jan 2025 15:39:51 +1100 Subject: [PATCH 6/8] Instrumentation support for dataloader - Adde doco and SimpleDataLoaderInstrumentationContext --- README.md | 59 +++++++++++++++++++ .../DataLoaderInstrumentationHelper.java | 27 +++++++++ ...impleDataLoaderInstrumentationContext.java | 35 +++++++++++ src/test/java/ReadmeExamples.java | 47 +++++++++++++++ ...eDataLoaderInstrumentationContextTest.java | 49 +++++++++++++++ 5 files changed, 217 insertions(+) create mode 100644 src/main/java/org/dataloader/instrumentation/SimpleDataLoaderInstrumentationContext.java create mode 100644 src/test/java/org/dataloader/instrumentation/SimpleDataLoaderInstrumentationContextTest.java diff --git a/README.md b/README.md index 2a0e08f..61180d9 100644 --- a/README.md +++ b/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> 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> 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 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 userDataLoader = DataLoaderFactory.newDataLoader(userBatchLoader); + DataLoader teamsDataLoader = DataLoaderFactory.newDataLoader(teamsBatchLoader); + + DataLoaderRegistry registry = DataLoaderRegistry.newRegistry() + .instrumentation(timingInstrumentation) + .register("users", userDataLoader) + .register("teams", teamsDataLoader) + .build(); + + DataLoader 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) diff --git a/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentationHelper.java b/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentationHelper.java index 844ec37..9e60060 100644 --- a/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentationHelper.java +++ b/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 DataLoaderInstrumentationContext 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 the generic type + * @return an instrumentation context + */ + public static DataLoaderInstrumentationContext 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 the generic type + * @return an instrumentation context + */ + public static DataLoaderInstrumentationContext whenCompleted(BiConsumer 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. diff --git a/src/main/java/org/dataloader/instrumentation/SimpleDataLoaderInstrumentationContext.java b/src/main/java/org/dataloader/instrumentation/SimpleDataLoaderInstrumentationContext.java new file mode 100644 index 0000000..f629a05 --- /dev/null +++ b/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 implements DataLoaderInstrumentationContext { + + private final BiConsumer codeToRunOnComplete; + private final Runnable codeToRunOnDispatch; + + SimpleDataLoaderInstrumentationContext(Runnable codeToRunOnDispatch, BiConsumer 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); + } + } +} diff --git a/src/test/java/ReadmeExamples.java b/src/test/java/ReadmeExamples.java index 9e30c90..3b8f57e 100644 --- a/src/test/java/ReadmeExamples.java +++ b/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 userBatchLoader; + BatchLoader 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> 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> 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 userDataLoader = DataLoaderFactory.newDataLoader(userBatchLoader, options); + } + + private void registryExample() { + DataLoader userDataLoader = DataLoaderFactory.newDataLoader(userBatchLoader); + DataLoader teamsDataLoader = DataLoaderFactory.newDataLoader(teamsBatchLoader); + + DataLoaderRegistry registry = DataLoaderRegistry.newRegistry() + .instrumentation(timingInstrumentation) + .register("users", userDataLoader) + .register("teams", teamsDataLoader) + .build(); + + DataLoader changedUsersDataLoader = registry.getDataLoader("users"); + + } } diff --git a/src/test/java/org/dataloader/instrumentation/SimpleDataLoaderInstrumentationContextTest.java b/src/test/java/org/dataloader/instrumentation/SimpleDataLoaderInstrumentationContextTest.java new file mode 100644 index 0000000..38328eb --- /dev/null +++ b/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 actual = new AtomicReference<>(); + AtomicReference actualErr = new AtomicReference<>(); + + DataLoaderInstrumentationContext 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 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)); + } +} \ No newline at end of file From 610343f5dfeacee7aa6e294a0f903f4048d74e06 Mon Sep 17 00:00:00 2001 From: bbaker Date: Wed, 22 Jan 2025 15:58:31 +1100 Subject: [PATCH 7/8] Instrumentation support for dataloader - Adde more doco --- src/test/java/ReadmeExamples.java | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/src/test/java/ReadmeExamples.java b/src/test/java/ReadmeExamples.java index 3b8f57e..1f718aa 100644 --- a/src/test/java/ReadmeExamples.java +++ b/src/test/java/ReadmeExamples.java @@ -427,4 +427,22 @@ private void registryExample() { DataLoader changedUsersDataLoader = registry.getDataLoader("users"); } + + private void combiningRegistryExample() { + DataLoader userDataLoader = DataLoaderFactory.newDataLoader(userBatchLoader); + DataLoader teamsDataLoader = DataLoaderFactory.newDataLoader(teamsBatchLoader); + + DataLoaderRegistry registry = DataLoaderRegistry.newRegistry() + .register("users", userDataLoader) + .register("teams", teamsDataLoader) + .build(); + + DataLoaderRegistry registryCombined = DataLoaderRegistry.newRegistry() + .instrumentation(timingInstrumentation) + .registerAll(registry) + .build(); + + DataLoader changedUsersDataLoader = registryCombined.getDataLoader("users"); + + } } From 53799826e394aa69340e78d557a88f08ec0f0984 Mon Sep 17 00:00:00 2001 From: bbaker Date: Sat, 1 Mar 2025 09:20:39 +1100 Subject: [PATCH 8/8] Merged in master and tweaked immutable building --- .../org/dataloader/DataLoaderOptions.java | 65 ++++++++++--------- .../org/dataloader/DataLoaderRegistry.java | 21 +++--- .../ChainedDataLoaderInstrumentationTest.java | 8 +-- ...DataLoaderRegistryInstrumentationTest.java | 13 ++-- 4 files changed, 61 insertions(+), 46 deletions(-) diff --git a/src/main/java/org/dataloader/DataLoaderOptions.java b/src/main/java/org/dataloader/DataLoaderOptions.java index 408b32d..e4b286a 100644 --- a/src/main/java/org/dataloader/DataLoaderOptions.java +++ b/src/main/java/org/dataloader/DataLoaderOptions.java @@ -86,6 +86,7 @@ private DataLoaderOptions(Builder builder) { this.environmentProvider = builder.environmentProvider; this.valueCacheOptions = builder.valueCacheOptions; this.batchLoaderScheduler = builder.batchLoaderScheduler; + this.instrumentation = builder.instrumentation; } /** @@ -174,7 +175,7 @@ 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 + * @return a new data loader options instance for fluent coding */ public DataLoaderOptions setBatchingEnabled(boolean batchingEnabled) { return builder().setBatchingEnabled(batchingEnabled).build(); @@ -193,7 +194,7 @@ 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 + * @return a new data loader options instance for fluent coding */ public DataLoaderOptions setCachingEnabled(boolean cachingEnabled) { return builder().setCachingEnabled(cachingEnabled).build(); @@ -217,7 +218,7 @@ 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 + * @return a new data loader options instance for fluent coding */ public DataLoaderOptions setCachingExceptionsEnabled(boolean cachingExceptionsEnabled) { return builder().setCachingExceptionsEnabled(cachingExceptionsEnabled).build(); @@ -238,7 +239,7 @@ public Optional 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 + * @return a new data loader options instance for fluent coding */ public DataLoaderOptions setCacheKeyFunction(CacheKey cacheKeyFunction) { return builder().setCacheKeyFunction(cacheKeyFunction).build(); @@ -259,7 +260,7 @@ 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 + * @return a new data loader options instance for fluent coding */ public DataLoaderOptions setCacheMap(CacheMap cacheMap) { return builder().setCacheMap(cacheMap).build(); @@ -280,7 +281,7 @@ public int maxBatchSize() { * before they are split into multiple class * * @param maxBatchSize the maximum batch size - * @return the data loader options for fluent coding + * @return a new data loader options instance for fluent coding */ public DataLoaderOptions setMaxBatchSize(int maxBatchSize) { return builder().setMaxBatchSize(maxBatchSize).build(); @@ -299,7 +300,7 @@ public StatisticsCollector getStatisticsCollector() { * a common value * * @param statisticsCollector the statistics collector to use - * @return the data loader options for fluent coding + * @return a new data loader options instance for fluent coding */ public DataLoaderOptions setStatisticsCollector(Supplier statisticsCollector) { return builder().setStatisticsCollector(nonNull(statisticsCollector)).build(); @@ -316,7 +317,7 @@ 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 + * @return a new data loader options instance for fluent coding */ public DataLoaderOptions setBatchLoaderContextProvider(BatchLoaderContextProvider contextProvider) { return builder().setBatchLoaderContextProvider(nonNull(contextProvider)).build(); @@ -337,7 +338,7 @@ 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 + * @return a new data loader options instance for fluent coding */ public DataLoaderOptions setValueCache(ValueCache valueCache) { return builder().setValueCache(valueCache).build(); @@ -354,7 +355,7 @@ 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 + * @return a new data loader options instance for fluent coding */ public DataLoaderOptions setValueCacheOptions(ValueCacheOptions valueCacheOptions) { return builder().setValueCacheOptions(nonNull(valueCacheOptions)).build(); @@ -372,12 +373,29 @@ public BatchLoaderScheduler getBatchLoaderScheduler() { * to some future time. * * @param batchLoaderScheduler the scheduler - * @return the data loader options for fluent coding + * @return a new data loader options instance for fluent coding */ public DataLoaderOptions setBatchLoaderScheduler(BatchLoaderScheduler batchLoaderScheduler) { return builder().setBatchLoaderScheduler(batchLoaderScheduler).build(); } + /** + * @return the {@link DataLoaderInstrumentation} to use + */ + public DataLoaderInstrumentation getInstrumentation() { + return instrumentation; + } + + /** + * Sets in a new {@link DataLoaderInstrumentation} + * + * @param instrumentation the new {@link DataLoaderInstrumentation} + * @return a new data loader options instance for fluent coding + */ + public DataLoaderOptions setInstrumentation(DataLoaderInstrumentation instrumentation) { + return builder().setInstrumentation(instrumentation).build(); + } + private Builder builder() { return new Builder(this); } @@ -394,6 +412,7 @@ public static class Builder { private BatchLoaderContextProvider environmentProvider; private ValueCacheOptions valueCacheOptions; private BatchLoaderScheduler batchLoaderScheduler; + private DataLoaderInstrumentation instrumentation; public Builder() { this(new DataLoaderOptions()); // use the defaults of the DataLoaderOptions for this builder @@ -411,6 +430,7 @@ public Builder() { this.environmentProvider = other.environmentProvider; this.valueCacheOptions = other.valueCacheOptions; this.batchLoaderScheduler = other.batchLoaderScheduler; + this.instrumentation = other.instrumentation; } public Builder setBatchingEnabled(boolean batchingEnabled) { @@ -468,27 +488,14 @@ public Builder setBatchLoaderScheduler(BatchLoaderScheduler batchLoaderScheduler return this; } + public Builder setInstrumentation(DataLoaderInstrumentation instrumentation) { + this.instrumentation = nonNull(instrumentation); + return this; + } + public DataLoaderOptions build() { return new DataLoaderOptions(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; - } } diff --git a/src/main/java/org/dataloader/DataLoaderRegistry.java b/src/main/java/org/dataloader/DataLoaderRegistry.java index d54ee53..06c93c4 100644 --- a/src/main/java/org/dataloader/DataLoaderRegistry.java +++ b/src/main/java/org/dataloader/DataLoaderRegistry.java @@ -31,7 +31,8 @@ *

* If the {@link DataLoader} has no {@link DataLoaderInstrumentation} then the registry one is added to it. If it does have one already * then a {@link ChainedDataLoaderInstrumentation} is created with the registry {@link DataLoaderInstrumentation} in it first and then any other - * {@link DataLoaderInstrumentation}s added after that. + * {@link DataLoaderInstrumentation}s added after that. If the registry {@link DataLoaderInstrumentation} instance and {@link DataLoader} {@link DataLoaderInstrumentation} instance + * are the same object, then nothing is changed, since the same instrumentation code is being run. */ @PublicApi public class DataLoaderRegistry { @@ -40,14 +41,14 @@ public class DataLoaderRegistry { public DataLoaderRegistry() { - this(new ConcurrentHashMap<>(),null); + this(new ConcurrentHashMap<>(), null); } private DataLoaderRegistry(Builder builder) { - this(builder.dataLoaders,builder.instrumentation); + this(builder.dataLoaders, builder.instrumentation); } - protected DataLoaderRegistry(Map> dataLoaders, DataLoaderInstrumentation instrumentation ) { + protected DataLoaderRegistry(Map> dataLoaders, DataLoaderInstrumentation instrumentation) { this.dataLoaders = instrumentDLs(dataLoaders, instrumentation); this.instrumentation = instrumentation; } @@ -97,12 +98,16 @@ protected DataLoaderRegistry(Map> dataLoaders, DataLoa } private static DataLoader mkInstrumentedDataLoader(DataLoader existingDL, DataLoaderOptions options, DataLoaderInstrumentation newInstrumentation) { - return existingDL.transform(builder -> { - options.setInstrumentation(newInstrumentation); - builder.options(options); - }); + return existingDL.transform(builder -> builder.options(setInInstrumentation(options, newInstrumentation))); + } + + private static DataLoaderOptions setInInstrumentation(DataLoaderOptions options, DataLoaderInstrumentation newInstrumentation) { + return options.transform(optionsBuilder -> optionsBuilder.setInstrumentation(newInstrumentation)); } + /** + * @return the {@link DataLoaderInstrumentation} associated with this registry which can be null + */ public DataLoaderInstrumentation getInstrumentation() { return instrumentation; } diff --git a/src/test/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentationTest.java b/src/test/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentationTest.java index b3272ce..d791762 100644 --- a/src/test/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentationTest.java +++ b/src/test/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentationTest.java @@ -14,7 +14,7 @@ import java.util.concurrent.CompletableFuture; import static org.awaitility.Awaitility.await; -import static org.dataloader.DataLoaderOptions.newOptions; +import static org.dataloader.DataLoaderOptions.newOptionsBuilder; import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.equalTo; @@ -37,7 +37,7 @@ void canChainTogetherZeroInstrumentation() { // just to prove its useless but harmless ChainedDataLoaderInstrumentation chainedItn = new ChainedDataLoaderInstrumentation(); - DataLoaderOptions options = newOptions().setInstrumentation(chainedItn); + DataLoaderOptions options = newOptionsBuilder().setInstrumentation(chainedItn).build(); DataLoader dl = DataLoaderFactory.newDataLoader(TestKit.keysAsValues(), options); @@ -57,7 +57,7 @@ void canChainTogetherOneInstrumentation() { ChainedDataLoaderInstrumentation chainedItn = new ChainedDataLoaderInstrumentation() .add(capturingA); - DataLoaderOptions options = newOptions().setInstrumentation(chainedItn); + DataLoaderOptions options = newOptionsBuilder().setInstrumentation(chainedItn).build(); DataLoader dl = DataLoaderFactory.newDataLoader(TestKit.keysAsValues(), options); @@ -83,7 +83,7 @@ public void canChainTogetherManyInstrumentationsWithDifferentBatchLoaders(TestDa .add(capturingB) .add(capturingButReturnsNull); - DataLoaderOptions options = newOptions().setInstrumentation(chainedItn); + DataLoaderOptions options = newOptionsBuilder().setInstrumentation(chainedItn).build(); DataLoader dl = factory.idLoader(options); diff --git a/src/test/java/org/dataloader/instrumentation/DataLoaderRegistryInstrumentationTest.java b/src/test/java/org/dataloader/instrumentation/DataLoaderRegistryInstrumentationTest.java index 5690bdc..465aa4d 100644 --- a/src/test/java/org/dataloader/instrumentation/DataLoaderRegistryInstrumentationTest.java +++ b/src/test/java/org/dataloader/instrumentation/DataLoaderRegistryInstrumentationTest.java @@ -1,6 +1,7 @@ package org.dataloader.instrumentation; import org.dataloader.DataLoader; +import org.dataloader.DataLoaderOptions; import org.dataloader.DataLoaderRegistry; import org.dataloader.fixtures.TestKit; import org.dataloader.fixtures.parameterized.TestDataLoaderFactory; @@ -119,9 +120,9 @@ void wontDoAnyThingIfThereIsNoRegistryInstrumentation() { @Test void wontDoAnyThingIfThereTheyAreTheSameInstrumentationAlready() { - DataLoader newX = dlX.transform(builder -> dlX.getOptions().setInstrumentation(instrA)); - DataLoader newY = dlX.transform(builder -> dlY.getOptions().setInstrumentation(instrA)); - DataLoader newZ = dlX.transform(builder -> dlY.getOptions().setInstrumentation(instrA)); + DataLoader newX = dlX.transform(builder -> builder.options(dlX.getOptions().setInstrumentation(instrA))); + DataLoader newY = dlX.transform(builder -> builder.options(dlY.getOptions().setInstrumentation(instrA))); + DataLoader newZ = dlX.transform(builder -> builder.options(dlZ.getOptions().setInstrumentation(instrA))); DataLoaderRegistry registry = DataLoaderRegistry.newRegistry() .instrumentation(instrA) .register("X", newX) @@ -144,7 +145,8 @@ void wontDoAnyThingIfThereTheyAreTheSameInstrumentationAlready() { @Test void ifTheDLHasAInstrumentationThenItsTurnedIntoAChainedOne() { - DataLoader newX = dlX.transform(builder -> dlX.getOptions().setInstrumentation(instrA)); + DataLoaderOptions options = dlX.getOptions().setInstrumentation(instrA); + DataLoader newX = dlX.transform(builder -> builder.options(options)); DataLoaderRegistry registry = DataLoaderRegistry.newRegistry() .instrumentation(instrB) @@ -162,7 +164,8 @@ void ifTheDLHasAInstrumentationThenItsTurnedIntoAChainedOne() { @Test void chainedInstrumentationsWillBeCombined() { - DataLoader newX = dlX.transform(builder -> builder.options(dlX.getOptions().setInstrumentation(chainedInstrB))); + DataLoaderOptions options = dlX.getOptions().setInstrumentation(chainedInstrB); + DataLoader newX = dlX.transform(builder -> builder.options(options)); DataLoaderRegistry registry = DataLoaderRegistry.newRegistry() .instrumentation(instrA)