Adds a predicate to DataLoaderRegistry and a per dataloader map of predicates is also possible

Author: bbakerman

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

@@ -1,6 +1,7 @@
 package org.dataloader;
 
 import org.dataloader.annotations.PublicApi;
+import org.dataloader.registries.DispatchPredicate;
 import org.dataloader.stats.Statistics;
 
 import java.util.ArrayList;
@@ -21,12 +22,17 @@
 @PublicApi
 public class DataLoaderRegistry {
     protected final Map<String, DataLoader<?, ?>> dataLoaders = new ConcurrentHashMap<>();
+    protected final Map<DataLoader<?, ?>, DispatchPredicate> dataLoaderPredicates = new ConcurrentHashMap<>();
+    protected final DispatchPredicate dispatchPredicate;
+
 
     public DataLoaderRegistry() {
+        this.dispatchPredicate = DispatchPredicate.DISPATCH_ALWAYS;
     }
 
-    private DataLoaderRegistry(Builder builder) {
+    protected DataLoaderRegistry(Builder<?> builder) {
         this.dataLoaders.putAll(builder.dataLoaders);
+        this.dispatchPredicate = builder.dispatchPredicate;
     }
 
 
@@ -43,6 +49,21 @@ public DataLoaderRegistry register(String key, DataLoader<?, ?> dataLoader) {
         return this;
     }
 
+    /**
+     * This will register a new dataloader and dispatch predicate associated with that data loader
+     *
+     * @param key               the key to put the data loader under
+     * @param dataLoader        the data loader to register
+     * @param dispatchPredicate the dispatch predicate to associate with this data loader
+     *
+     * @return this registry
+     */
+    public DataLoaderRegistry register(String key, DataLoader<?, ?> dataLoader, DispatchPredicate dispatchPredicate) {
+        dataLoaders.put(key, dataLoader);
+        dataLoaderPredicates.put(dataLoader, dispatchPredicate);
+        return this;
+    }
+
     /**
      * Computes a data loader if absent or return it if it was
      * already registered at that key.
@@ -76,6 +97,8 @@ public DataLoaderRegistry combine(DataLoaderRegistry registry) {
 
         this.dataLoaders.forEach(combined::register);
         registry.dataLoaders.forEach(combined::register);
+        combined.dataLoaderPredicates.putAll(this.dataLoaderPredicates);
+        combined.dataLoaderPredicates.putAll(registry.dataLoaderPredicates);
         return combined;
     }
 
@@ -101,7 +124,10 @@ public DataLoaderRegistry combine(DataLoaderRegistry registry) {
      * @return this registry
      */
     public DataLoaderRegistry unregister(String key) {
-        dataLoaders.remove(key);
+        DataLoader<?, ?> dataLoader = dataLoaders.remove(key);
+        if (dataLoader != null) {
+            dataLoaderPredicates.remove(dataLoader);
+        }
         return this;
     }
 
@@ -131,7 +157,7 @@ public Set<String> getKeys() {
      * {@link org.dataloader.DataLoader}s
      */
     public void dispatchAll() {
-        getDataLoaders().forEach(DataLoader::dispatch);
+        dispatchAllWithCount();
     }
 
     /**
@@ -142,8 +168,12 @@ public void dispatchAll() {
      */
     public int dispatchAllWithCount() {
         int sum = 0;
-        for (DataLoader<?, ?> dataLoader : getDataLoaders()) {
-            sum += dataLoader.dispatchWithCounts().getKeysCount();
+        for (Map.Entry<String, DataLoader<?, ?>> entry : dataLoaders.entrySet()) {
+            DataLoader<?, ?> dataLoader = entry.getValue();
+            String key = entry.getKey();
+            if (shouldDispatch(key, dataLoader)) {
+                sum += dataLoader.dispatchWithCounts().getKeysCount();
+            }
         }
         return sum;
     }
@@ -154,12 +184,59 @@ public int dispatchAllWithCount() {
      */
     public int dispatchDepth() {
         int totalDispatchDepth = 0;
-        for (DataLoader<?, ?> dataLoader : getDataLoaders()) {
-            totalDispatchDepth += dataLoader.dispatchDepth();
+        for (Map.Entry<String, DataLoader<?, ?>> entry : dataLoaders.entrySet()) {
+            DataLoader<?, ?> dataLoader = entry.getValue();
+            String key = entry.getKey();
+            if (shouldDispatch(key, dataLoader)) {
+                totalDispatchDepth += dataLoader.dispatchDepth();
+            }
         }
         return totalDispatchDepth;
     }
 
+    /**
+     * This will immediately dispatch the {@link DataLoader}s in the registry
+     * without testing the predicate
+     */
+    public void dispatchAllImmediately() {
+        dispatchAllWithCountImmediately();
+    }
+
+    /**
+     * This will immediately dispatch the {@link DataLoader}s in the registry
+     * without testing the predicate
+     *
+     * @return total number of entries that were dispatched from registered {@link org.dataloader.DataLoader}s.
+     */
+    public int dispatchAllWithCountImmediately() {
+        int sum = 0;
+        for (Map.Entry<String, DataLoader<?, ?>> entry : dataLoaders.entrySet()) {
+            DataLoader<?, ?> dataLoader = entry.getValue();
+            sum += dataLoader.dispatchWithCounts().getKeysCount();
+        }
+        return sum;
+    }
+
+
+    /**
+     * Returns true if the dataloader has a predicate which returned true, OR the overall
+     * registry predicate returned true.
+     *
+     * @param dataLoaderKey the key in the dataloader map
+     * @param dataLoader    the dataloader
+     *
+     * @return true if it should dispatch
+     */
+    protected boolean shouldDispatch(String dataLoaderKey, DataLoader<?, ?> dataLoader) {
+        DispatchPredicate dispatchPredicate = dataLoaderPredicates.get(dataLoader);
+        if (dispatchPredicate != null) {
+            if (dispatchPredicate.test(dataLoaderKey, dataLoader)) {
+                return true;
+            }
+        }
+        return this.dispatchPredicate.test(dataLoaderKey, dataLoader);
+    }
+
     /**
      * @return a combined set of statistics for all data loaders in this registry presented
      * as the sum of all their statistics
@@ -175,13 +252,22 @@ public Statistics getStatistics() {
     /**
      * @return A builder of {@link DataLoaderRegistry}s
      */
-    public static Builder newRegistry() {
+    public static Builder<?> newRegistry() {
+        //noinspection rawtypes
         return new Builder();
     }
 
-    public static class Builder {
+    public static class Builder<B extends Builder<B>> {
 
         private final Map<String, DataLoader<?, ?>> dataLoaders = new HashMap<>();
+        private final Map<DataLoader<?, ?>, DispatchPredicate> dataLoaderPredicates = new ConcurrentHashMap<>();
+
+        private DispatchPredicate dispatchPredicate = DispatchPredicate.DISPATCH_ALWAYS;
+
+        private B self() {
+            //noinspection unchecked
+            return (B) this;
+        }
 
         /**
          * This will register a new dataloader
@@ -191,22 +277,51 @@ public static class Builder {
          *
          * @return this builder for a fluent pattern
          */
-        public Builder register(String key, DataLoader<?, ?> dataLoader) {
+        public B register(String key, DataLoader<?, ?> dataLoader) {
             dataLoaders.put(key, dataLoader);
-            return this;
+            return self();
         }
 
         /**
-         * This will combine together the data loaders in this builder with the ones
+         * This will register a new dataloader with a specific {@link DispatchPredicate}
+         *
+         * @param key               the key to put the data loader under
+         * @param dataLoader        the data loader to register
+         * @param dispatchPredicate the dispatch predicate
+         *
+         * @return this builder for a fluent pattern
+         */
+        public B register(String key, DataLoader<?, ?> dataLoader, DispatchPredicate dispatchPredicate) {
+            register(key, dataLoader);
+            dataLoaderPredicates.put(dataLoader, dispatchPredicate);
+            return self();
+        }
+
+        /**
+         * This will combine the data loaders in this builder with the ones
          * from a previous {@link DataLoaderRegistry}
          *
          * @param otherRegistry the previous {@link DataLoaderRegistry}
          *
          * @return this builder for a fluent pattern
          */
-        public Builder registerAll(DataLoaderRegistry otherRegistry) {
+        public B registerAll(DataLoaderRegistry otherRegistry) {
             dataLoaders.putAll(otherRegistry.dataLoaders);
-            return this;
+            dataLoaderPredicates.putAll(otherRegistry.dataLoaderPredicates);
+            return self();
+        }
+
+        /**
+         * This sets a predicate on the {@link DataLoaderRegistry} that will control
+         * whether all {@link DataLoader}s in the {@link DataLoaderRegistry }should be dispatched.
+         *
+         * @param dispatchPredicate the predicate
+         *
+         * @return this builder for a fluent pattern
+         */
+        public B dispatchPredicate(DispatchPredicate dispatchPredicate) {
+            this.dispatchPredicate = dispatchPredicate;
+            return self();
         }
 
         /**

src/main/java/org/dataloader/registries/DispatchPredicate.java

@@ -10,6 +10,16 @@
  */
 @FunctionalInterface
 public interface DispatchPredicate {
+
+    /**
+     * A predicate that always returns true
+     */
+    DispatchPredicate DISPATCH_ALWAYS = (dataLoaderKey, dataLoader) -> true;
+    /**
+     * A predicate that always returns false
+     */
+    DispatchPredicate DISPATCH_NEVER = (dataLoaderKey, dataLoader) -> true;
+
     /**
      * This predicate tests whether the data loader should be dispatched or not.
      *

src/main/java/org/dataloader/registries/ScheduledDataLoaderRegistry.java

@@ -5,7 +5,6 @@
 import org.dataloader.annotations.ExperimentalApi;
 
 import java.time.Duration;
-import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
@@ -30,14 +29,12 @@
 public class ScheduledDataLoaderRegistry extends DataLoaderRegistry implements AutoCloseable {
 
     private final ScheduledExecutorService scheduledExecutorService;
-    private final DispatchPredicate dispatchPredicate;
     private final Duration schedule;
     private volatile boolean closed;
 
     private ScheduledDataLoaderRegistry(Builder builder) {
-        this.dataLoaders.putAll(builder.dataLoaders);
+        super(builder);
         this.scheduledExecutorService = builder.scheduledExecutorService;
-        this.dispatchPredicate = builder.dispatchPredicate;
         this.schedule = builder.schedule;
         this.closed = false;
     }
@@ -77,24 +74,6 @@ public int dispatchAllWithCount() {
         return sum;
     }
 
-    /**
-     * This will immediately dispatch the {@link DataLoader}s in the registry
-     * without testing the predicate
-     */
-    public void dispatchAllImmediately() {
-        super.dispatchAll();
-    }
-
-    /**
-     * This will immediately dispatch the {@link DataLoader}s in the registry
-     * without testing the predicate
-     *
-     * @return total number of entries that were dispatched from registered {@link org.dataloader.DataLoader}s.
-     */
-    public int dispatchAllWithCountImmediately() {
-        return super.dispatchAllWithCount();
-    }
-
     /**
      * This will schedule a task to check the predicate and dispatch if true right now.  It will not do
      * a pre check of the preodicate like {@link #dispatchAll()} would
@@ -111,7 +90,7 @@ private void reschedule(String key, DataLoader<?, ?> dataLoader) {
     }
 
     private void dispatchOrReschedule(String key, DataLoader<?, ?> dataLoader) {
-        if (dispatchPredicate.test(key, dataLoader)) {
+        if (shouldDispatch(key, dataLoader)) {
             dataLoader.dispatch();
         } else {
             reschedule(key, dataLoader);
@@ -128,12 +107,10 @@ public static Builder newScheduledRegistry() {
         return new Builder();
     }
 
-    public static class Builder {
+    public static class Builder extends DataLoaderRegistry.Builder<ScheduledDataLoaderRegistry.Builder> {
 
         private ScheduledExecutorService scheduledExecutorService = Executors.newSingleThreadScheduledExecutor();
-        private DispatchPredicate dispatchPredicate = (key, dl) -> true;
         private Duration schedule = Duration.ofMillis(10);
-        private final Map<String, DataLoader<?, ?>> dataLoaders = new HashMap<>();
 
         public Builder scheduledExecutorService(ScheduledExecutorService executorService) {
             this.scheduledExecutorService = nonNull(executorService);
@@ -145,37 +122,6 @@ public Builder schedule(Duration schedule) {
             return this;
         }
 
-        public Builder dispatchPredicate(DispatchPredicate dispatchPredicate) {
-            this.dispatchPredicate = nonNull(dispatchPredicate);
-            return this;
-        }
-
-        /**
-         * 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) {
-            dataLoaders.put(key, dataLoader);
-            return this;
-        }
-
-        /**
-         * This will combine together the data loaders in this builder with the ones
-         * from a previous {@link DataLoaderRegistry}
-         *
-         * @param otherRegistry the previous {@link DataLoaderRegistry}
-         *
-         * @return this builder for a fluent pattern
-         */
-        public Builder registerAll(DataLoaderRegistry otherRegistry) {
-            dataLoaders.putAll(otherRegistry.getDataLoadersMap());
-            return this;
-        }
-
         /**
          * @return the newly built {@link ScheduledDataLoaderRegistry}
          */