Merge remote-tracking branch 'origin/master' into delegating-dataloader-support

bbakerman · bbakerman · commit 5452d738ccb1 · 2025-04-02T12:38:11.000+11:00
diff --git a/src/main/java/org/dataloader/DataLoaderHelper.java b/src/main/java/org/dataloader/DataLoaderHelper.java
@@ -148,12 +148,16 @@ CompletableFuture<V> load(K key, Object loadContext) {
             boolean cachingEnabled = loaderOptions.cachingEnabled();
 
             stats.incrementLoadCount(new IncrementLoadCountStatisticsContext<>(key, loadContext));
-
+            DataLoaderInstrumentationContext<Object> ctx = ctxOrNoopCtx(instrumentation().beginLoad(dataLoader, key,loadContext));
+            CompletableFuture<V> cf;
             if (cachingEnabled) {
-                return loadFromCache(key, loadContext, batchingEnabled);
+                cf = loadFromCache(key, loadContext, batchingEnabled);
             } else {
-                return queueOrInvokeLoader(key, loadContext, batchingEnabled, false);
+                cf = queueOrInvokeLoader(key, loadContext, batchingEnabled, false);
             }
+            ctx.onDispatched();
+            cf.whenComplete(ctx::onCompleted);
+            return cf;
         }
     }
 
diff --git a/src/main/java/org/dataloader/DataLoaderOptions.java b/src/main/java/org/dataloader/DataLoaderOptions.java
@@ -139,7 +139,7 @@ public static DataLoaderOptions.Builder newDataLoaderOptions(DataLoaderOptions o
      * @return a new {@link DataLoaderOptions} object
      */
     public DataLoaderOptions transform(Consumer<Builder> builderConsumer) {
-        Builder builder = newOptionsBuilder();
+        Builder builder = newDataLoaderOptions(this);
         builderConsumer.accept(builder);
         return builder.build();
     }
diff --git a/src/main/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentation.java b/src/main/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentation.java
@@ -69,6 +69,12 @@ public ChainedDataLoaderInstrumentation addAll(Collection<DataLoaderInstrumentat
         return new ChainedDataLoaderInstrumentation(list);
     }
 
+
+    @Override
+    public DataLoaderInstrumentationContext<Object> beginLoad(DataLoader<?, ?> dataLoader, Object key, Object loadContext) {
+        return chainedCtx(it -> it.beginLoad(dataLoader, key, loadContext));
+    }
+
     @Override
     public DataLoaderInstrumentationContext<DispatchResult<?>> beginDispatch(DataLoader<?, ?> dataLoader) {
         return chainedCtx(it -> it.beginDispatch(dataLoader));
diff --git a/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentation.java b/src/main/java/org/dataloader/instrumentation/DataLoaderInstrumentation.java
@@ -12,6 +12,21 @@
  */
 @PublicSpi
 public interface DataLoaderInstrumentation {
+    /**
+     * This call back is done just before the {@link DataLoader#load(Object)} methods are invoked,
+     * and it completes when the load promise is completed.  If the value is a cached {@link java.util.concurrent.CompletableFuture}
+     * then it might return almost immediately, otherwise it will return
+     * when the batch load function is invoked and values get returned
+     *
+     * @param dataLoader  the {@link DataLoader} in question
+     * @param key         the key used during the {@link DataLoader#load(Object)} call
+     * @param loadContext the load context used during the {@link DataLoader#load(Object, Object)} call
+     * @return a DataLoaderInstrumentationContext or null to be more performant
+     */
+    default DataLoaderInstrumentationContext<Object> beginLoad(DataLoader<?, ?> dataLoader, Object key, Object loadContext) {
+        return null;
+    }
+
     /**
      * This call back is done just before the {@link DataLoader#dispatch()} is invoked,
      * and it completes when the dispatch call promise is done.
diff --git a/src/test/java/org/dataloader/DataLoaderBuilderTest.java b/src/test/java/org/dataloader/DataLoaderBuilderTest.java
@@ -46,19 +46,31 @@ void canBuildNewDataLoaders() {
 
     @Test
     void theDataLoaderCanTransform() {
-        DataLoader<String, Object> dataLoader1 = DataLoaderFactory.newDataLoader(batchLoader1, defaultOptions);
-        assertThat(dataLoader1.getOptions(), equalTo(defaultOptions));
-        assertThat(dataLoader1.getBatchLoadFunction(), equalTo(batchLoader1));
+        DataLoader<String, Object> dataLoaderOrig = DataLoaderFactory.newDataLoader(batchLoader1, defaultOptions);
+        assertThat(dataLoaderOrig.getOptions(), equalTo(defaultOptions));
+        assertThat(dataLoaderOrig.getBatchLoadFunction(), equalTo(batchLoader1));
         //
         // we can transform the data loader
         //
-        DataLoader<String, Object> dataLoader2 = dataLoader1.transform(it -> {
+        DataLoader<String, Object> dataLoaderTransformed = dataLoaderOrig.transform(it -> {
             it.options(differentOptions);
             it.batchLoadFunction(batchLoader2);
         });
 
-        assertThat(dataLoader2, not(equalTo(dataLoader1)));
-        assertThat(dataLoader2.getOptions(), equalTo(differentOptions));
-        assertThat(dataLoader2.getBatchLoadFunction(), equalTo(batchLoader2));
+        assertThat(dataLoaderTransformed, not(equalTo(dataLoaderOrig)));
+        assertThat(dataLoaderTransformed.getOptions(), equalTo(differentOptions));
+        assertThat(dataLoaderTransformed.getBatchLoadFunction(), equalTo(batchLoader2));
+
+        // can copy values
+        dataLoaderOrig = DataLoaderFactory.newDataLoader(batchLoader1, defaultOptions);
+
+        dataLoaderTransformed = dataLoaderOrig.transform(it -> {
+            it.batchLoadFunction(batchLoader2);
+        });
+
+        assertThat(dataLoaderTransformed, not(equalTo(dataLoaderOrig)));
+        assertThat(dataLoaderTransformed.getOptions(), equalTo(defaultOptions));
+        assertThat(dataLoaderTransformed.getBatchLoadFunction(), equalTo(batchLoader2));
+
     }
 }
diff --git a/src/test/java/org/dataloader/DataLoaderCacheMapTest.java b/src/test/java/org/dataloader/DataLoaderCacheMapTest.java
@@ -43,7 +43,7 @@ public void should_access_to_future_dependants() {
         Collection<CompletableFuture<Integer>> futures = dataLoader.getCacheMap().getAll();
 
         List<CompletableFuture<Integer>> futuresList = new ArrayList<>(futures);
-        assertThat(futuresList.get(0).getNumberOfDependents(), equalTo(2));
-        assertThat(futuresList.get(1).getNumberOfDependents(), equalTo(1));
+        assertThat(futuresList.get(0).getNumberOfDependents(), equalTo(4)); // instrumentation is depending on the CF completing
+        assertThat(futuresList.get(1).getNumberOfDependents(), equalTo(2));
     }
 }
diff --git a/src/test/java/org/dataloader/DataLoaderOptionsTest.java b/src/test/java/org/dataloader/DataLoaderOptionsTest.java
@@ -2,9 +2,11 @@
 
 import org.dataloader.impl.DefaultCacheMap;
 import org.dataloader.impl.NoOpValueCache;
+import org.dataloader.instrumentation.DataLoaderInstrumentation;
 import org.dataloader.scheduler.BatchLoaderScheduler;
 import org.dataloader.stats.NoOpStatisticsCollector;
 import org.dataloader.stats.StatisticsCollector;
+import org.hamcrest.CoreMatchers;
 import org.junit.jupiter.api.Test;
 
 import java.util.List;
@@ -184,4 +186,45 @@ void canBuildViaBuilderOk() {
         assertThat(builtOptions.getStatisticsCollector(),
                 equalTo(testStatisticsCollectorSupplier.get()));
     }
+
+    @Test
+    void canCopyExistingOptionValuesOnTransform() {
+
+        DataLoaderInstrumentation instrumentation1 = new DataLoaderInstrumentation() {
+        };
+        DataLoaderInstrumentation instrumentation2 = new DataLoaderInstrumentation() {
+        };
+        BatchLoaderContextProvider contextProvider1 = () -> null;
+
+        DataLoaderOptions startingOptions = DataLoaderOptions.newOptionsBuilder().setBatchingEnabled(false)
+                .setCachingEnabled(false)
+                .setInstrumentation(instrumentation1)
+                .setBatchLoaderContextProvider(contextProvider1)
+                .build();
+
+        assertThat(startingOptions.batchingEnabled(), equalTo(false));
+        assertThat(startingOptions.cachingEnabled(), equalTo(false));
+        assertThat(startingOptions.getInstrumentation(), equalTo(instrumentation1));
+        assertThat(startingOptions.getBatchLoaderContextProvider(), equalTo(contextProvider1));
+
+        DataLoaderOptions newOptions = startingOptions.transform(builder ->
+                builder.setBatchingEnabled(true).setInstrumentation(instrumentation2));
+
+
+        // immutable
+        assertThat(newOptions, CoreMatchers.not(startingOptions));
+        assertThat(startingOptions.batchingEnabled(), equalTo(false));
+        assertThat(startingOptions.cachingEnabled(), equalTo(false));
+        assertThat(startingOptions.getInstrumentation(), equalTo(instrumentation1));
+        assertThat(startingOptions.getBatchLoaderContextProvider(), equalTo(contextProvider1));
+
+        // stayed the same
+        assertThat(newOptions.cachingEnabled(), equalTo(false));
+        assertThat(newOptions.getBatchLoaderContextProvider(), equalTo(contextProvider1));
+
+        // was changed
+        assertThat(newOptions.batchingEnabled(), equalTo(true));
+        assertThat(newOptions.getInstrumentation(), equalTo(instrumentation2));
+
+    }
 }
diff --git a/src/test/java/org/dataloader/instrumentation/CapturingInstrumentation.java b/src/test/java/org/dataloader/instrumentation/CapturingInstrumentation.java
@@ -6,15 +6,49 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.stream.Collectors;
 
 class CapturingInstrumentation implements DataLoaderInstrumentation {
-    String name;
-    List<String> methods = new ArrayList<>();
+    protected String name;
+    protected List<String> methods = new ArrayList<>();
 
     public CapturingInstrumentation(String name) {
         this.name = name;
     }
 
+    public String getName() {
+        return name;
+    }
+
+    public List<String> methods() {
+        return methods;
+    }
+
+    public List<String> notLoads() {
+        return methods.stream().filter(method -> !method.contains("beginLoad")).collect(Collectors.toList());
+    }
+
+    public List<String> onlyLoads() {
+        return methods.stream().filter(method -> method.contains("beginLoad")).collect(Collectors.toList());
+    }
+
+
+    @Override
+    public DataLoaderInstrumentationContext<Object> beginLoad(DataLoader<?, ?> dataLoader, Object key, Object loadContext) {
+        methods.add(name + "_beginLoad" +"_k:" + key);
+        return new DataLoaderInstrumentationContext<>() {
+            @Override
+            public void onDispatched() {
+                methods.add(name + "_beginLoad_onDispatched"+"_k:" + key);
+            }
+
+            @Override
+            public void onCompleted(Object result, Throwable t) {
+                methods.add(name + "_beginLoad_onCompleted"+"_k:" + key);
+            }
+        };
+    }
+
     @Override
     public DataLoaderInstrumentationContext<DispatchResult<?>> beginDispatch(DataLoader<?, ?> dataLoader) {
         methods.add(name + "_beginDispatch");
diff --git a/src/test/java/org/dataloader/instrumentation/CapturingInstrumentationReturnsNull.java b/src/test/java/org/dataloader/instrumentation/CapturingInstrumentationReturnsNull.java
@@ -12,6 +12,12 @@ public CapturingInstrumentationReturnsNull(String name) {
         super(name);
     }
 
+    @Override
+    public DataLoaderInstrumentationContext<Object> beginLoad(DataLoader<?, ?> dataLoader, Object key, Object loadContext) {
+        methods.add(name + "_beginLoad" +"_k:" + key);
+        return null;
+    }
+
     @Override
     public DataLoaderInstrumentationContext<DispatchResult<?>> beginDispatch(DataLoader<?, ?> dataLoader) {
         methods.add(name + "_beginDispatch");
diff --git a/src/test/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentationTest.java b/src/test/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentationTest.java
@@ -61,16 +61,21 @@ void canChainTogetherOneInstrumentation() {
 
         DataLoader<String, String> dl = DataLoaderFactory.newDataLoader(TestKit.keysAsValues(), options);
 
-        dl.load("A");
-        dl.load("B");
+        dl.load("X");
+        dl.load("Y");
 
         CompletableFuture<List<String>> dispatch = dl.dispatch();
 
         await().until(dispatch::isDone);
 
-        assertThat(capturingA.methods, equalTo(List.of("A_beginDispatch",
+        assertThat(capturingA.notLoads(), equalTo(List.of("A_beginDispatch",
                 "A_beginBatchLoader", "A_beginBatchLoader_onDispatched", "A_beginBatchLoader_onCompleted",
                 "A_beginDispatch_onDispatched", "A_beginDispatch_onCompleted")));
+
+        assertThat(capturingA.onlyLoads(), equalTo(List.of(
+                "A_beginLoad_k:X", "A_beginLoad_onDispatched_k:X", "A_beginLoad_k:Y", "A_beginLoad_onDispatched_k:Y",
+                "A_beginLoad_onCompleted_k:X", "A_beginLoad_onCompleted_k:Y"
+        )));
     }
 
 
@@ -87,8 +92,8 @@ public void canChainTogetherManyInstrumentationsWithDifferentBatchLoaders(TestDa
 
         DataLoader<String, String> dl = factory.idLoader(options);
 
-        dl.load("A");
-        dl.load("B");
+        dl.load("X");
+        dl.load("Y");
 
         CompletableFuture<List<String>> dispatch = dl.dispatch();
 
@@ -98,16 +103,21 @@ public void canChainTogetherManyInstrumentationsWithDifferentBatchLoaders(TestDa
         // 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",
+        assertThat(capturingA.notLoads(), 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",
+        assertThat(capturingA.onlyLoads(), equalTo(List.of(
+                "A_beginLoad_k:X", "A_beginLoad_onDispatched_k:X", "A_beginLoad_k:Y", "A_beginLoad_onDispatched_k:Y",
+                "A_beginLoad_onCompleted_k:X", "A_beginLoad_onCompleted_k:Y"
+        )));
+
+        assertThat(capturingB.notLoads(), 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")));
+        assertThat(capturingButReturnsNull.notLoads(), equalTo(List.of("NULL_beginDispatch", "NULL_beginBatchLoader")));
     }
 
     @Test
diff --git a/src/test/java/org/dataloader/instrumentation/DataLoaderInstrumentationTest.java b/src/test/java/org/dataloader/instrumentation/DataLoaderInstrumentationTest.java
@@ -29,6 +29,61 @@ public class DataLoaderInstrumentationTest {
         return keys;
     });
 
+    @Test
+    void canMonitorLoading() {
+        AtomicReference<DataLoader<?, ?>> dlRef = new AtomicReference<>();
+
+        CapturingInstrumentation instrumentation = new CapturingInstrumentation("x") {
+
+            @Override
+            public DataLoaderInstrumentationContext<Object> beginLoad(DataLoader<?, ?> dataLoader, Object key, Object loadContext) {
+                DataLoaderInstrumentationContext<Object> superCtx = super.beginLoad(dataLoader, key, loadContext);
+                dlRef.set(dataLoader);
+                return superCtx;
+            }
+
+            @Override
+            public DataLoaderInstrumentationContext<List<?>> beginBatchLoader(DataLoader<?, ?> dataLoader, List<?> keys, BatchLoaderEnvironment environment) {
+                return DataLoaderInstrumentationHelper.noOpCtx();
+            }
+        };
+
+        DataLoaderOptions options = new DataLoaderOptions()
+                .setInstrumentation(instrumentation)
+                .setMaxBatchSize(5);
+
+        DataLoader<String, String> dl = DataLoaderFactory.newDataLoader(snoozingBatchLoader, options);
+
+        List<String> keys = new ArrayList<>();
+        for (int i = 0; i < 3; i++) {
+            String key = "X" + i;
+            keys.add(key);
+            dl.load(key);
+        }
+
+        // load a key that is cached
+        dl.load("X0");
+
+        CompletableFuture<List<String>> dispatch = dl.dispatch();
+
+        await().until(dispatch::isDone);
+        assertThat(dlRef.get(), is(dl));
+        assertThat(dispatch.join(), equalTo(keys));
+
+        // the batch loading means they start and are instrumentation dispatched before they all end up completing
+        assertThat(instrumentation.onlyLoads(),
+                equalTo(List.of(
+                        "x_beginLoad_k:X0", "x_beginLoad_onDispatched_k:X0",
+                        "x_beginLoad_k:X1", "x_beginLoad_onDispatched_k:X1",
+                        "x_beginLoad_k:X2", "x_beginLoad_onDispatched_k:X2",
+                        "x_beginLoad_k:X0", "x_beginLoad_onDispatched_k:X0", // second cached call counts
+                        "x_beginLoad_onCompleted_k:X0",
+                        "x_beginLoad_onCompleted_k:X0", // each load call counts
+                        "x_beginLoad_onCompleted_k:X1", "x_beginLoad_onCompleted_k:X2")));
+
+    }
+
+
     @Test
     void canMonitorDispatching() {
         Stopwatch stopwatch = Stopwatch.stopwatchUnStarted();
diff --git a/src/test/java/org/dataloader/instrumentation/DataLoaderRegistryInstrumentationTest.java b/src/test/java/org/dataloader/instrumentation/DataLoaderRegistryInstrumentationTest.java
@@ -224,7 +224,7 @@ public void endToEndIntegrationTest(TestDataLoaderFactory factory) {
         await().until(loadA::isDone);
         assertThat(loadA.join(), equalTo("A"));
 
-        assertThat(instrA.methods, equalTo(List.of("A_beginDispatch",
+        assertThat(instrA.notLoads(), equalTo(List.of("A_beginDispatch",
                 "A_beginBatchLoader", "A_beginBatchLoader_onDispatched", "A_beginBatchLoader_onCompleted",
                 "A_beginDispatch_onDispatched", "A_beginDispatch_onCompleted")));
     }

Original file line number	Diff line number	Diff line change
`@@ -148,12 +148,16 @@ CompletableFuture<V> load(K key, Object loadContext) {`
`148`	`148`	`boolean cachingEnabled = loaderOptions.cachingEnabled();`
`149`	`149`
`150`	`150`	`stats.incrementLoadCount(new IncrementLoadCountStatisticsContext<>(key, loadContext));`
`151`		`-`
	`151`	`+ DataLoaderInstrumentationContext<Object> ctx = ctxOrNoopCtx(instrumentation().beginLoad(dataLoader, key,loadContext));`
	`152`	`+ CompletableFuture<V> cf;`
`152`	`153`	`if (cachingEnabled) {`
`153`		`- return loadFromCache(key, loadContext, batchingEnabled);`
	`154`	`+ cf = loadFromCache(key, loadContext, batchingEnabled);`
`154`	`155`	`} else {`
`155`		`- return queueOrInvokeLoader(key, loadContext, batchingEnabled, false);`
	`156`	`+ cf = queueOrInvokeLoader(key, loadContext, batchingEnabled, false);`
`156`	`157`	`}`
	`158`	`+ ctx.onDispatched();`
	`159`	`+ cf.whenComplete(ctx::onCompleted);`
	`160`	`+ return cf;`
`157`	`161`	`}`
`158`	`162`	`}`
`159`	`163`
Original file line number	Diff line number	Diff line change
`@@ -139,7 +139,7 @@ public static DataLoaderOptions.Builder newDataLoaderOptions(DataLoaderOptions o`
`139`	`139`	`* @return a new {@link DataLoaderOptions} object`
`140`	`140`	`*/`
`141`	`141`	`public DataLoaderOptions transform(Consumer<Builder> builderConsumer) {`
`142`		`- Builder builder = newOptionsBuilder();`
	`142`	`+ Builder builder = newDataLoaderOptions(this);`
`143`	`143`	`builderConsumer.accept(builder);`
`144`	`144`	`return builder.build();`
`145`	`145`	`}`
Original file line number	Diff line number	Diff line change
`@@ -43,7 +43,7 @@ public void should_access_to_future_dependants() {`
`43`	`43`	`Collection<CompletableFuture<Integer>> futures = dataLoader.getCacheMap().getAll();`
`44`	`44`
`45`	`45`	`List<CompletableFuture<Integer>> futuresList = new ArrayList<>(futures);`
`46`		`- assertThat(futuresList.get(0).getNumberOfDependents(), equalTo(2));`
`47`		`- assertThat(futuresList.get(1).getNumberOfDependents(), equalTo(1));`
	`46`	`+ assertThat(futuresList.get(0).getNumberOfDependents(), equalTo(4)); // instrumentation is depending on the CF completing`
	`47`	`+ assertThat(futuresList.get(1).getNumberOfDependents(), equalTo(2));`
`48`	`48`	`}`
`49`	`49`	`}`