expand Strategy arguments, make sure strategy is called earlier

Author: andimarek

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

@@ -188,18 +188,18 @@ CompletableFuture<V> load(K key, Object loadContext) {
         }
 
         ctx.onDispatched();
-        loaderOptions.getDispatchStrategy().loadCalled();
+        loaderOptions.getDispatchStrategy().loadCalled(dataLoader, key, loadContext, loadCallFuture);
+        loadCallFuture = loadCallFuture.whenComplete((result, error) -> loaderOptions.getDispatchStrategy().loadCompleted(dataLoader, result, error));
         loadCallFuture.whenComplete(ctx::onCompleted);
-        loadCallFuture.whenComplete((result, error) -> loaderOptions.getDispatchStrategy().loadCompleted());
         return loadCallFuture;
     }
 
     private CompletableFuture<V> incrementCacheHitAndReturnCF(DataLoaderInstrumentationContext<Object> ctx, K key, Object loadContext, CompletableFuture<V> cachedFuture) {
         stats.incrementCacheHitCount(new IncrementCacheHitCountStatisticsContext<>(key, loadContext));
         ctx.onDispatched();
-        loaderOptions.getDispatchStrategy().loadCalled();
+        loaderOptions.getDispatchStrategy().loadCalled(dataLoader, key, loadContext, cachedFuture);
+        cachedFuture = cachedFuture.whenComplete((result, error) -> loaderOptions.getDispatchStrategy().loadCompleted(dataLoader, result, error));
         cachedFuture.whenComplete(ctx::onCompleted);
-        cachedFuture.whenComplete((result, error) -> loaderOptions.getDispatchStrategy().loadCompleted());
         return cachedFuture;
     }
 

src/main/java/org/dataloader/DispatchStrategy.java

@@ -2,6 +2,9 @@
 
 import org.dataloader.annotations.PublicApi;
 import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
+
+import java.util.concurrent.CompletableFuture;
 
 
 /**
@@ -28,14 +31,14 @@ default void onRegistryCreation(DataLoaderRegistry registry) {
     /**
      * Called when a {@link DataLoader#load(Object)} is called on a dataloader
      */
-    default void loadCalled() {
+    default void loadCalled(DataLoader<?, ?> dataLoader, Object key, @Nullable Object loadContext, CompletableFuture<?> result) {
 
     }
 
     /**
      * Called when a {@link DataLoader#load(Object)} is executed and completed on a dataloader
      */
-    default void loadCompleted() {
+    default void loadCompleted(DataLoader<?, ?> dataLoader, @Nullable Object result, @Nullable Throwable error) {
 
     }
 }

src/main/java/org/dataloader/strategy/GreedyLevelByLevelChainedDispatchStrategy.java

@@ -1,12 +1,15 @@
 package org.dataloader.strategy;
 
+import org.dataloader.DataLoader;
 import org.dataloader.DataLoaderRegistry;
 import org.dataloader.DispatchStrategy;
 import org.dataloader.annotations.VisibleForTesting;
 import org.dataloader.impl.Assertions;
+import org.jspecify.annotations.NullMarked;
 import org.jspecify.annotations.Nullable;
 
 import java.time.Duration;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
@@ -25,6 +28,7 @@
  * In practice this will greedily fill up DataLoader keys while walking the chain to provide a nice balance of
  * batching/dedupe/caching while not needing to worry about manually dispatching the tree.
  */
+@NullMarked
 public class GreedyLevelByLevelChainedDispatchStrategy implements DispatchStrategy {
 
     private static final Duration DEFAULT_FALLBACK_TIMEOUT = Duration.ofMillis(30);
@@ -35,7 +39,9 @@ public class GreedyLevelByLevelChainedDispatchStrategy implements DispatchStrate
     private final Object dispatchLock = new Object();
 
     // only used for tests
-    private Runnable onIteration;
+    Runnable onIteration = () -> {
+
+    };
 
     private final Duration fallbackTimeout;
     @Nullable private ScheduledFuture<?> fallbackDispatchFuture = null;
@@ -53,7 +59,7 @@ public void onRegistryCreation(DataLoaderRegistry registry) {
     }
 
     @Override
-    public void loadCalled() {
+    public void loadCalled(DataLoader<?, ?> dataLoader, Object key, @Nullable Object loaderContext, CompletableFuture<?> result) {
         // initial load called
         pendingLoadCount.incrementAndGet();
         int previousTotal = totalWorkCount.getAndIncrement();
@@ -63,7 +69,7 @@ public void loadCalled() {
     }
 
     @Override
-    public void loadCompleted() {
+    public void loadCompleted(DataLoader<?, ?> dataLoader, @Nullable Object result, @Nullable Throwable error) {
         pendingLoadCount.decrementAndGet();
     }