Adding a ScheduledDataLoaderRegistry

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

@@ -6,6 +6,7 @@
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -19,7 +20,7 @@
  */
 @PublicApi
 public class DataLoaderRegistry {
-    private final Map<String, DataLoader<?, ?>> dataLoaders = new ConcurrentHashMap<>();
+    protected final Map<String, DataLoader<?, ?>> dataLoaders = new ConcurrentHashMap<>();
 
     public DataLoaderRegistry() {
     }
@@ -28,6 +29,7 @@ private DataLoaderRegistry(Builder builder) {
         this.dataLoaders.putAll(builder.dataLoaders);
     }
 
+
     /**
      * This will register a new dataloader
      *
@@ -84,6 +86,13 @@ public DataLoaderRegistry combine(DataLoaderRegistry registry) {
         return new ArrayList<>(dataLoaders.values());
     }
 
+    /**
+     * @return the currently registered data loaders as a map
+     */
+    public Map<String, DataLoader<?, ?>> getDataLoadersMap() {
+        return new LinkedHashMap<>(dataLoaders);
+    }
+
     /**
      * This will unregister a new dataloader
      *

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

@@ -0,0 +1,92 @@
+package org.dataloader.registries;
+
+import org.dataloader.DataLoader;
+
+import java.time.Duration;
+import java.util.Objects;
+
+/**
+ * A predicate class used by {@link ScheduledDataLoaderRegistry} to decide whether to dispatch or not
+ */
+@FunctionalInterface
+public interface DispatchPredicate {
+    /**
+     * This predicate tests whether the data loader should be dispatched or not.
+     *
+     * @param dataLoaderKey the key of the data loader when registered
+     * @param dataLoader    the dataloader to dispatch
+     *
+     * @return true if the data loader SHOULD be dispatched
+     */
+    boolean test(String dataLoaderKey, DataLoader<?, ?> dataLoader);
+
+
+    /**
+     * Returns a composed predicate that represents a short-circuiting logical
+     * AND of this predicate and another.
+     *
+     * @param other a predicate that will be logically-ANDed with this
+     *              predicate
+     *
+     * @return a composed predicate that represents the short-circuiting logical
+     * AND of this predicate and the {@code other} predicate
+     */
+    default DispatchPredicate and(DispatchPredicate other) {
+        Objects.requireNonNull(other);
+        return (k, dl) -> test(k, dl) && other.test(k, dl);
+    }
+
+    /**
+     * Returns a predicate that represents the logical negation of this
+     * predicate.
+     *
+     * @return a predicate that represents the logical negation of this
+     * predicate
+     */
+    default DispatchPredicate negate() {
+        return (k, dl) -> !test(k, dl);
+    }
+
+    /**
+     * Returns a composed predicate that represents a short-circuiting logical
+     * OR of this predicate and another.
+     *
+     * @param other a predicate that will be logically-ORed with this
+     *              predicate
+     *
+     * @return a composed predicate that represents the short-circuiting logical
+     * OR of this predicate and the {@code other} predicate
+     */
+    default DispatchPredicate or(DispatchPredicate other) {
+        Objects.requireNonNull(other);
+        return (k, dl) -> test(k, dl) || other.test(k, dl);
+    }
+
+    /**
+     * This predicate will return true if the {@link DataLoader} has not be dispatched
+     * for at least the duration length of time.
+     *
+     * @param duration the length of time to check
+     *
+     * @return true if the data loader has not been dispatched in duration time
+     */
+    static DispatchPredicate dispatchIfLongerThan(Duration duration) {
+        return (dataLoaderKey, dataLoader) -> {
+            int i = dataLoader.getTimeSinceDispatch().compareTo(duration);
+            return i > 0;
+        };
+    }
+
+    /**
+     * This predicate will return true if the {@link DataLoader#dispatchDepth()} is greater than the specified depth.
+     *
+     * This will act as minimum batch size.  There must be more then `depth` items queued for the predicate to return true.
+     *
+     * @param depth the value to be greater than
+     *
+     * @return true if the {@link DataLoader#dispatchDepth()} is greater than the specified depth.
+     */
+    static DispatchPredicate dispatchIfDepthGreaterThan(int depth) {
+        return (dataLoaderKey, dataLoader) -> dataLoader.dispatchDepth() > depth;
+    }
+}

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

@@ -0,0 +1,184 @@
+package org.dataloader.registries;
+
+import org.dataloader.DataLoader;
+import org.dataloader.DataLoaderRegistry;
+import org.dataloader.annotations.PublicApi;
+
+import java.time.Duration;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+
+import static org.dataloader.impl.Assertions.nonNull;
+
+/**
+ * This {@link DataLoaderRegistry} will use a {@link DispatchPredicate} when {@link #dispatchAll()} is called
+ * to test (for each {@link DataLoader} in the registry) if a dispatch should proceed.  If the predicate returns false, then a task is scheduled
+ * to perform that predicate dispatch again via the {@link ScheduledExecutorService}.
+ * <p>
+ * This will continue to loop (test false and reschedule) until such time as the predicate returns true, in which case
+ * no rescheduling will occur and you will need to call dispatch again to restart the process.
+ * <p>
+ * If you wanted to create a ScheduledDataLoaderRegistry that started a rescheduling immediately, just create one and
+ * call {@link #rescheduleNow()}.
+ */
+@PublicApi
+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);
+        this.scheduledExecutorService = builder.scheduledExecutorService;
+        this.dispatchPredicate = builder.dispatchPredicate;
+        this.schedule = builder.schedule;
+        this.closed = false;
+    }
+
+    /**
+     * Once closed this registry will never again reschedule checks
+     */
+    @Override
+    public void close() {
+        closed = true;
+    }
+
+    /**
+     * @return how long the {@link ScheduledExecutorService} task will wait before checking the predicate again
+     */
+    public Duration getScheduleDuration() {
+        return schedule;
+    }
+
+    @Override
+    public void dispatchAll() {
+        dispatchAllWithCount();
+    }
+
+    @Override
+    public int dispatchAllWithCount() {
+        int sum = 0;
+        for (Map.Entry<String, DataLoader<?, ?>> entry : dataLoaders.entrySet()) {
+            DataLoader<?, ?> dataLoader = entry.getValue();
+            String key = entry.getKey();
+            if (dispatchPredicate.test(key, dataLoader)) {
+                sum += dataLoader.dispatchWithCounts().getKeysCount();
+            } else {
+                reschedule(key, dataLoader);
+            }
+        }
+        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
+     */
+    public void rescheduleNow() {
+        dataLoaders.forEach(this::reschedule);
+    }
+
+    private void reschedule(String key, DataLoader<?, ?> dataLoader) {
+        if (!closed) {
+            Runnable runThis = () -> dispatchOrReschedule(key, dataLoader);
+            scheduledExecutorService.schedule(runThis, schedule.toMillis(), TimeUnit.MILLISECONDS);
+        }
+    }
+
+    private void dispatchOrReschedule(String key, DataLoader<?, ?> dataLoader) {
+        if (dispatchPredicate.test(key, dataLoader)) {
+            dataLoader.dispatch();
+        } else {
+            reschedule(key, dataLoader);
+        }
+    }
+
+    /**
+     * By default this will create use a {@link Executors#newSingleThreadScheduledExecutor()}
+     * and a schedule duration of 10 milli seconds.
+     *
+     * @return A builder of {@link ScheduledDataLoaderRegistry}s
+     */
+    public static Builder newScheduledRegistry() {
+        return new Builder();
+    }
+
+    public static class 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);
+            return this;
+        }
+
+        public Builder schedule(Duration schedule) {
+            this.schedule = 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}
+         */
+        public ScheduledDataLoaderRegistry build() {
+            return new ScheduledDataLoaderRegistry(this);
+        }
+    }
+}

src/test/java/ReadmeExamples.java

@@ -10,9 +10,12 @@
 import org.dataloader.fixtures.SecurityCtx;
 import org.dataloader.fixtures.User;
 import org.dataloader.fixtures.UserManager;
+import org.dataloader.registries.DispatchPredicate;
+import org.dataloader.registries.ScheduledDataLoaderRegistry;
 import org.dataloader.stats.Statistics;
 import org.dataloader.stats.ThreadLocalStatisticsCollector;
 
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -269,4 +272,14 @@ private void statsConfigExample() {
         DataLoader<String, User> userDataLoader = DataLoaderFactory.newDataLoader(userBatchLoader, options);
     }
 
+    private void ScheduledDispatche() {
+        DispatchPredicate depthOrTimePredicate = DispatchPredicate.dispatchIfDepthGreaterThan(10)
+                .or(DispatchPredicate.dispatchIfLongerThan(Duration.ofMillis(200)));
+
+        ScheduledDataLoaderRegistry registry = ScheduledDataLoaderRegistry.newScheduledRegistry()
+                .dispatchPredicate(depthOrTimePredicate)
+                .schedule(Duration.ofMillis(10))
+                .register("users", userDataLoader)
+                .build();
+    }
 }

src/test/java/org/dataloader/ClockDataLoader.java

@@ -4,11 +4,11 @@
 
 public class ClockDataLoader<K, V> extends DataLoader<K, V> {
 
-    ClockDataLoader(Object batchLoadFunction, Clock clock) {
+    public ClockDataLoader(Object batchLoadFunction, Clock clock) {
         this(batchLoadFunction, null, clock);
     }
 
-    ClockDataLoader(Object batchLoadFunction, DataLoaderOptions options, Clock clock) {
+    public ClockDataLoader(Object batchLoadFunction, DataLoaderOptions options, Clock clock) {
         super(batchLoadFunction, options, clock);
     }
 

src/test/java/org/dataloader/fixtures/TestKit.java

@@ -31,6 +31,10 @@ public static <K, V> BatchLoader<K, V> keysAsValues(List<List<K>> loadCalls) {
         };
     }
 
+    public static <K, V> DataLoader<K, V> idLoader() {
+        return idLoader(null, new ArrayList<>());
+    }
+
     public static <K, V> DataLoader<K, V> idLoader(List<List<K>> loadCalls) {
         return idLoader(null, loadCalls);
     }