Update dispatch strategy and impl bfs strategy

Author: zsmoore

src/main/java/org/dataloader/DataLoader.java

@@ -229,9 +229,7 @@ public Optional<CompletableFuture<V>> getIfCompleted(K key) {
      * @return the future of the value
      */
     public CompletableFuture<V> load(@NonNull K key, @Nullable Object keyContext) {
-        CompletableFuture<V> result = loadImpl(key, keyContext);
-        options.getDispatchStrategy().loadCalled(this);
-        return result;
+        return loadImpl(key, keyContext);
     }
 
     private CompletableFuture<V> loadImpl(@NonNull K key, @Nullable Object keyContext) {
@@ -283,7 +281,6 @@ public CompletableFuture<List<V>> loadMany(List<K> keys, List<Object> keyContext
             }
             collect.add(loadImpl(key, keyContext));
         }
-        options.getDispatchStrategy().loadCalled(this);
         return CompletableFutureKit.allOf(collect);
     }
 
@@ -311,7 +308,6 @@ public CompletableFuture<Map<K, V>> loadMany(Map<K, ?> keysAndContexts) {
             Object keyContext = entry.getValue();
             collect.put(key, loadImpl(key, keyContext));
         }
-        options.getDispatchStrategy().loadCalled(this);
         return CompletableFutureKit.allOf(collect);
     }
 

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

@@ -158,7 +158,9 @@ CompletableFuture<V> load(K key, Object loadContext) {
                     // We already have a promise for this key, no need to check value cache or queue up load
                     stats.incrementCacheHitCount(new IncrementCacheHitCountStatisticsContext<>(key, loadContext));
                     ctx.onDispatched();
+                    loaderOptions.getDispatchStrategy().loadCalled();
                     cachedFuture.whenComplete(ctx::onCompleted);
+                    cachedFuture.whenComplete((result, error) -> loaderOptions.getDispatchStrategy().loadCompleted());
                     return cachedFuture;
                 }
             } catch (Exception ignored) {
@@ -173,7 +175,9 @@ CompletableFuture<V> load(K key, Object loadContext) {
                     // another thread was faster and created a matching CF ... hence this is really a cachehit and we are done
                     stats.incrementCacheHitCount(new IncrementCacheHitCountStatisticsContext<>(key, loadContext));
                     ctx.onDispatched();
+                    loaderOptions.getDispatchStrategy().loadCalled();
                     cachedFuture.whenComplete(ctx::onCompleted);
+                    cachedFuture.whenComplete((result, error) -> loaderOptions.getDispatchStrategy().loadCompleted());
                     return cachedFuture;
                 }
             }
@@ -190,14 +194,18 @@ CompletableFuture<V> load(K key, Object loadContext) {
                     // meaning this is a cache hit and we are done
                     stats.incrementCacheHitCount(new IncrementCacheHitCountStatisticsContext<>(key, loadContext));
                     ctx.onDispatched();
+                    loaderOptions.getDispatchStrategy().loadCalled();
                     cachedFuture.whenComplete(ctx::onCompleted);
+                    cachedFuture.whenComplete((result, error) -> loaderOptions.getDispatchStrategy().loadCompleted());
                     return cachedFuture;
                 }
             }
         }
 
         ctx.onDispatched();
+        loaderOptions.getDispatchStrategy().loadCalled();
         loadCallFuture.whenComplete(ctx::onCompleted);
+        loadCallFuture.whenComplete((result, error) -> loaderOptions.getDispatchStrategy().loadCompleted());
         return loadCallFuture;
     }
 

src/main/java/org/dataloader/DataLoaderRegistry.java

@@ -58,14 +58,15 @@ private DataLoaderRegistry(Builder builder) {
     }
 
     protected DataLoaderRegistry(Map<String, DataLoader<?, ?>> dataLoaders, @Nullable DataLoaderInstrumentation instrumentation, DispatchStrategy dispatchStrategy) {
-        this.dataLoaders = instrumentDLs(dataLoaders, instrumentation);
+        this.dataLoaders = instrumentDLs(dataLoaders, instrumentation, dispatchStrategy);
         this.instrumentation = instrumentation;
         this.dispatchStrategy = dispatchStrategy;
+        dispatchStrategy.onRegistryCreation(this);
     }
 
-    private Map<String, DataLoader<?, ?>> instrumentDLs(Map<String, DataLoader<?, ?>> incomingDataLoaders, @Nullable DataLoaderInstrumentation registryInstrumentation) {
+    private Map<String, DataLoader<?, ?>> instrumentDLs(Map<String, DataLoader<?, ?>> incomingDataLoaders, @Nullable DataLoaderInstrumentation registryInstrumentation, DispatchStrategy dispatchStrategy) {
         Map<String, DataLoader<?, ?>> dataLoaders = new ConcurrentHashMap<>(incomingDataLoaders);
-        if (registryInstrumentation != null) {
+        if (registryInstrumentation != null || dispatchStrategy != DispatchStrategy.NO_OP) {
             dataLoaders.replaceAll((k, existingDL) -> nameAndInstrumentDL(k, registryInstrumentation, existingDL, dispatchStrategy));
         }
         return dataLoaders;
@@ -83,7 +84,7 @@ protected DataLoaderRegistry(Map<String, DataLoader<?, ?>> dataLoaders, @Nullabl
     private static DataLoader<?, ?> nameAndInstrumentDL(String key, @Nullable DataLoaderInstrumentation registryInstrumentation, DataLoader<?, ?> existingDL, DispatchStrategy dispatchStrategy) {
         existingDL = checkAndSetName(key, existingDL);
 
-        if (registryInstrumentation == null) {
+        if (registryInstrumentation == null && dispatchStrategy == DispatchStrategy.NO_OP) {
             return existingDL;
         }
         DataLoaderOptions options = existingDL.getOptions();
@@ -120,12 +121,17 @@ protected DataLoaderRegistry(Map<String, DataLoader<?, ?>> dataLoaders, @Nullabl
         return dataLoader;
     }
 
-    private static DataLoader<?, ?> mkInstrumentedDataLoader(DataLoader<?, ?> existingDL, DataLoaderOptions options, DataLoaderInstrumentation newInstrumentation, DispatchStrategy dispatchStrategy) {
+    private static DataLoader<?, ?> mkInstrumentedDataLoader(DataLoader<?, ?> existingDL, DataLoaderOptions options, @Nullable DataLoaderInstrumentation newInstrumentation, DispatchStrategy dispatchStrategy) {
         return existingDL.transform(builder -> builder.options(setInInstrumentation(options, newInstrumentation, dispatchStrategy)));
     }
 
-    private static DataLoaderOptions setInInstrumentation(DataLoaderOptions options, DataLoaderInstrumentation newInstrumentation, DispatchStrategy dispatchStrategy) {
-        return options.transform(optionsBuilder -> optionsBuilder.setInstrumentation(newInstrumentation).setDispatchStrategy(dispatchStrategy));
+    private static DataLoaderOptions setInInstrumentation(DataLoaderOptions options, @Nullable DataLoaderInstrumentation newInstrumentation, DispatchStrategy dispatchStrategy) {
+        return options.transform(optionsBuilder -> {
+            optionsBuilder.setDispatchStrategy(dispatchStrategy);
+            if (newInstrumentation != null) {
+                optionsBuilder.setInstrumentation(newInstrumentation);
+            }
+        });
     }
 
     /**

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

@@ -3,14 +3,24 @@
 import org.dataloader.annotations.PublicApi;
 import org.jspecify.annotations.NullMarked;
 
+import java.util.concurrent.atomic.AtomicBoolean;
+
 @NullMarked
 @PublicApi
 public interface DispatchStrategy {
 
     DispatchStrategy NO_OP = new DispatchStrategy() {
     };
 
-    default void loadCalled(DataLoader<?, ?> dataLoader) {
+    default void onRegistryCreation(DataLoaderRegistry registry) {
+
+    }
+
+    default void loadCalled() {
+
+    }
+
+    default void loadCompleted() {
 
     }
 }

src/main/java/org/dataloader/NotBusyDispatchStrategy.java

@@ -2,6 +2,7 @@
 
 import org.dataloader.DataLoader;
 import org.dataloader.DataLoaderRegistry;
+import org.dataloader.DispatchStrategy;
 import org.dataloader.annotations.ExperimentalApi;
 import org.dataloader.impl.Assertions;
 import org.dataloader.instrumentation.DataLoaderInstrumentation;
@@ -69,7 +70,7 @@ public class ScheduledDataLoaderRegistry extends DataLoaderRegistry implements A
     private volatile boolean closed;
 
     private ScheduledDataLoaderRegistry(Builder builder) {
-        super(builder.dataLoaders, builder.instrumentation);
+        super(builder.dataLoaders, builder.instrumentation, DispatchStrategy.NO_OP);
         this.scheduledExecutorService = Assertions.nonNull(builder.scheduledExecutorService);
         this.defaultExecutorUsed = builder.defaultExecutorUsed;
         this.schedule = builder.schedule;