This adds a ticker mode to ScheduledDataLoaderRegistry

Author: bbakerman

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

@@ -15,15 +15,34 @@
 
 /**
  * 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}.
+ * 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.
+ * In the default mode, when {@link #tickerMode} is false, the registry 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>
+ * However, when {@link #tickerMode} is true, the registry will always reschedule continuously after the first ever call to {@link #dispatchAll()}.
+ * <p>
+ * This will allow you to chain together {@link DataLoader} load calls like this :
+ * <pre>{@code
+ *   CompletableFuture<String> future = dataLoaderA.load("A")
+ *                                          .thenCompose(value -> dataLoaderB.load(value));
+ * }</pre>
+ * <p>
+ * However, it may mean your batching will not be as efficient as it might be. In environments
+ * like graphql this might mean you are too eager in fetching.  The {@link DispatchPredicate} still runs to decide if
+ * dispatch should happen however in ticker mode it will be continuously rescheduled.
+ * <p>
+ * When {@link #tickerMode} is true, you really SHOULD close the registry say at the end of a request otherwise you will leave a job
+ * on the {@link ScheduledExecutorService} that is continuously dispatching.
  * <p>
  * If you wanted to create a ScheduledDataLoaderRegistry that started a rescheduling immediately, just create one and
  * call {@link #rescheduleNow()}.
  * <p>
+ * By default, it uses a {@link Executors#newSingleThreadScheduledExecutor()}} to schedule the tasks.  However, if you
+ * are creating a {@link ScheduledDataLoaderRegistry} per request you will want to look at sharing this {@link ScheduledExecutorService}
+ * to avoid creating a new thread per registry created.
+ * <p>
  * This code is currently marked as {@link ExperimentalApi}
  */
 @ExperimentalApi
@@ -32,13 +51,15 @@ public class ScheduledDataLoaderRegistry extends DataLoaderRegistry implements A
     private final ScheduledExecutorService scheduledExecutorService;
     private final DispatchPredicate dispatchPredicate;
     private final Duration schedule;
+    private final boolean tickerMode;
     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.tickerMode = builder.tickerMode;
         this.closed = false;
     }
 
@@ -57,6 +78,13 @@ public Duration getScheduleDuration() {
         return schedule;
     }
 
+    /**
+     * @return true of the registry is in ticker mode or false otherwise
+     */
+    public boolean isTickerMode() {
+        return tickerMode;
+    }
+
     @Override
     public void dispatchAll() {
         dispatchAllWithCount();
@@ -68,11 +96,7 @@ public int dispatchAllWithCount() {
         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);
-            }
+            dispatchOrReschedule(key, dataLoader);
         }
         return sum;
     }
@@ -111,9 +135,11 @@ private void reschedule(String key, DataLoader<?, ?> dataLoader) {
     }
 
     private void dispatchOrReschedule(String key, DataLoader<?, ?> dataLoader) {
-        if (dispatchPredicate.test(key, dataLoader)) {
+        boolean shouldDispatch = dispatchPredicate.test(key, dataLoader);
+        if (shouldDispatch) {
             dataLoader.dispatch();
-        } else {
+        }
+        if (tickerMode || !shouldDispatch) {
             reschedule(key, dataLoader);
         }
     }
@@ -134,6 +160,7 @@ public static class Builder {
         private DispatchPredicate dispatchPredicate = (key, dl) -> true;
         private Duration schedule = Duration.ofMillis(10);
         private final Map<String, DataLoader<?, ?>> dataLoaders = new HashMap<>();
+        private boolean tickerMode = false;
 
         public Builder scheduledExecutorService(ScheduledExecutorService executorService) {
             this.scheduledExecutorService = nonNull(executorService);
@@ -176,6 +203,20 @@ public Builder registerAll(DataLoaderRegistry otherRegistry) {
             return this;
         }
 
+        /**
+         * This sets ticker mode on the registry.  When ticker mode is true the registry will
+         * continuously reschedule the data loaders for possible dispatching after the first call
+         * to dispatchAll.
+         *
+         * @param tickerMode true or false
+         *
+         * @return this builder for a fluent pattern
+         */
+        public Builder tickerMode(boolean tickerMode) {
+            this.tickerMode = tickerMode;
+            return this;
+        }
+
         /**
          * @return the newly built {@link ScheduledDataLoaderRegistry}
          */

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

@@ -5,6 +5,7 @@
 import org.dataloader.DataLoaderFactory;
 import org.dataloader.DataLoaderOptions;
 
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
@@ -31,6 +32,23 @@ public static <K, V> BatchLoader<K, V> keysAsValues(List<List<K>> loadCalls) {
         };
     }
 
+    public static <K, V> BatchLoader<K, V> keysAsValuesAsync(Duration delay) {
+        return keysAsValuesAsync(new ArrayList<>(), delay);
+    }
+
+    public static <K, V> BatchLoader<K, V> keysAsValuesAsync(List<List<K>> loadCalls, Duration delay) {
+        return keys -> CompletableFuture.supplyAsync(() -> {
+            snooze(delay.toMillis());
+            List<K> ks = new ArrayList<>(keys);
+            loadCalls.add(ks);
+            @SuppressWarnings("unchecked")
+            List<V> values = keys.stream()
+                    .map(k -> (V) k)
+                    .collect(toList());
+            return values;
+        });
+    }
+
     public static <K, V> DataLoader<K, V> idLoader() {
         return idLoader(null, new ArrayList<>());
     }
@@ -43,6 +61,14 @@ public static <K, V> DataLoader<K, V> idLoader(DataLoaderOptions options, List<L
         return DataLoaderFactory.newDataLoader(keysAsValues(loadCalls), options);
     }
 
+    public static <K, V> DataLoader<K, V> idLoaderAsync(Duration delay) {
+        return idLoaderAsync(null, new ArrayList<>(), delay);
+    }
+
+    public static <K, V> DataLoader<K, V> idLoaderAsync(DataLoaderOptions options, List<List<K>> loadCalls, Duration delay) {
+        return DataLoaderFactory.newDataLoader(keysAsValuesAsync(loadCalls, delay), options);
+    }
+
     public static Collection<Integer> listFrom(int i, int max) {
         List<Integer> ints = new ArrayList<>();
         for (int j = i; j < max; j++) {
@@ -55,7 +81,7 @@ public static <V> CompletableFuture<V> futureError() {
         return failedFuture(new IllegalStateException("Error"));
     }
 
-    public static void snooze(int millis) {
+    public static void snooze(long millis) {
         try {
             Thread.sleep(millis);
         } catch (InterruptedException e) {

src/test/java/org/dataloader/registries/ScheduledDataLoaderRegistryTest.java

@@ -1,6 +1,7 @@
 package org.dataloader.registries;
 
 import junit.framework.TestCase;
+import org.awaitility.core.ConditionTimeoutException;
 import org.dataloader.DataLoader;
 import org.dataloader.DataLoaderFactory;
 import org.dataloader.DataLoaderRegistry;
@@ -11,13 +12,17 @@
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executors;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import static java.util.Arrays.asList;
 import static java.util.Collections.singletonList;
+import static org.awaitility.Awaitility.await;
+import static org.awaitility.Duration.TWO_SECONDS;
 import static org.dataloader.fixtures.TestKit.keysAsValues;
 import static org.dataloader.fixtures.TestKit.snooze;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.is;
 import static org.junit.Assert.assertThat;
 
 public class ScheduledDataLoaderRegistryTest extends TestCase {
@@ -257,4 +262,63 @@ public void test_close_is_a_one_way_door() {
         snooze(200);
         assertEquals(counter.get(), countThen + 1);
     }
+
+    public void test_can_tick_after_first_dispatch_for_chain_data_loaders() {
+
+        // delays much bigger than the tick rate will mean multiple calls to dispatch
+        DataLoader<String, String> dlA = TestKit.idLoaderAsync(Duration.ofMillis(100));
+        DataLoader<String, String> dlB = TestKit.idLoaderAsync(Duration.ofMillis(200));
+
+        CompletableFuture<String> chainedCF = dlA.load("AK1").thenCompose(dlB::load);
+
+        AtomicBoolean done = new AtomicBoolean();
+        chainedCF.whenComplete((v, t) -> done.set(true));
+
+        ScheduledDataLoaderRegistry registry = ScheduledDataLoaderRegistry.newScheduledRegistry()
+                .register("a", dlA)
+                .register("b", dlB)
+                .dispatchPredicate(alwaysDispatch)
+                .schedule(Duration.ofMillis(10))
+                .tickerMode(true)
+                .build();
+
+        assertThat(registry.isTickerMode(), equalTo(true));
+
+        registry.dispatchAll();
+
+        await().atMost(TWO_SECONDS).untilAtomic(done, is(true));
+
+        registry.close();
+    }
+
+    public void test_chain_data_loaders_will_hang_if_not_in_ticker_mode() {
+
+        // delays much bigger than the tick rate will mean multiple calls to dispatch
+        DataLoader<String, String> dlA = TestKit.idLoaderAsync(Duration.ofMillis(100));
+        DataLoader<String, String> dlB = TestKit.idLoaderAsync(Duration.ofMillis(200));
+
+        CompletableFuture<String> chainedCF = dlA.load("AK1").thenCompose(dlB::load);
+
+        AtomicBoolean done = new AtomicBoolean();
+        chainedCF.whenComplete((v, t) -> done.set(true));
+
+        ScheduledDataLoaderRegistry registry = ScheduledDataLoaderRegistry.newScheduledRegistry()
+                .register("a", dlA)
+                .register("b", dlB)
+                .dispatchPredicate(alwaysDispatch)
+                .schedule(Duration.ofMillis(10))
+                .tickerMode(false)
+                .build();
+
+        assertThat(registry.isTickerMode(), equalTo(false));
+
+        registry.dispatchAll();
+
+        try {
+            await().atMost(TWO_SECONDS).untilAtomic(done, is(true));
+            fail("This should not have completed but rather timed out");
+        } catch (ConditionTimeoutException expected) {
+        }
+        registry.close();
+    }
 }