Keep one pending ticker task per ScheduledDataLoaderRegistry

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

@@ -15,6 +15,7 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import static org.dataloader.impl.Assertions.nonNull;
 
@@ -33,6 +34,8 @@
  * 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()}.
+ * A single task per registry re-runs the dispatch of all data loaders every {@link #getScheduleDuration()}, however
+ * often {@link #dispatchAll()} or {@link #rescheduleNow()} is called.
  * <p>
  * This will allow you to chain together {@link DataLoader} load calls like this :
  * <pre>{@code
@@ -67,6 +70,7 @@ public class ScheduledDataLoaderRegistry extends DataLoaderRegistry implements A
     private final Duration schedule;
     private final boolean tickerMode;
     private volatile boolean closed;
+    private final AtomicBoolean tickScheduled = new AtomicBoolean();
 
     private ScheduledDataLoaderRegistry(Builder builder) {
         super(builder.dataLoaders, builder.instrumentation);
@@ -214,7 +218,11 @@ public int dispatchAllWithCountImmediately() {
      * a pre-check of the predicate like {@link #dispatchAll()} would
      */
     public void rescheduleNow() {
-        dataLoaders.forEach(this::reschedule);
+        if (tickerMode) {
+            scheduleTick();
+        } else {
+            dataLoaders.forEach(this::reschedule);
+        }
     }
 
     /**
@@ -246,12 +254,32 @@ private int dispatchOrReschedule(String key, DataLoader<?, ?> dataLoader) {
         if (shouldDispatch) {
             sum = dataLoader.dispatchWithCounts().getKeysCount();
         }
-        if (tickerMode || !shouldDispatch) {
+        if (tickerMode) {
+            scheduleTick();
+        } else if (!shouldDispatch) {
             reschedule(key, dataLoader);
         }
         return sum;
     }
 
+    /**
+     * In ticker mode a single task per registry re-runs {@link #dispatchAllWithCount()} every {@link #schedule},
+     * however often {@link #dispatchAll()} is called. Scheduling a task per data loader on every
+     * {@link #dispatchAll()} call would grow the executor queue by the number of registered data loaders on each call.
+     */
+    private void scheduleTick() {
+        if (!closed && tickScheduled.compareAndSet(false, true)) {
+            scheduledExecutorService.schedule(this::tick, schedule.toMillis(), TimeUnit.MILLISECONDS);
+        }
+    }
+
+    private void tick() {
+        tickScheduled.set(false);
+        // re-arm before dispatching so that a data loader throwing on dispatch cannot stop the ticker
+        scheduleTick();
+        dispatchAllWithCount();
+    }
+
     /**
      * By default, this will create use a {@link Executors#newSingleThreadScheduledExecutor()}
      * and a schedule duration of 10 milliseconds.
@@ -349,7 +377,7 @@ public Builder dispatchPredicate(DispatchPredicate dispatchPredicate) {
         /**
          * 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.
+         * to dispatchAll, using a single scheduled task per registry.
          *
          * @param tickerMode true or false
          * @return this builder for a fluent pattern

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

@@ -15,13 +15,16 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
 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.idLoader;
 import static org.dataloader.fixtures.TestKit.snooze;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.equalTo;
@@ -350,6 +353,87 @@ public void chain_data_loaders_will_hang_if_not_in_ticker_mode(TestDataLoaderFac
         registry.close();
     }
 
+    @Test
+    public void ticker_mode_keeps_one_pending_task_however_often_dispatch_all_is_called() {
+        ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor(1);
+        // a tick that never fires during the test, so the queue is only inspected between ticks
+        ScheduledDataLoaderRegistry registry = tickerRegistry(executor, Duration.ofHours(1), 20, alwaysDispatch);
+
+        for (int i = 0; i < 50; i++) {
+            registry.dispatchAll();
+        }
+
+        assertThat(executor.getQueue().size(), equalTo(1));
+        registry.close();
+        executor.shutdownNow();
+    }
+
+    @Test
+    public void ticker_mode_reschedule_now_keeps_one_pending_task() {
+        ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor(1);
+        ScheduledDataLoaderRegistry registry = tickerRegistry(executor, Duration.ofHours(1), 20, alwaysDispatch);
+        registry.dispatchAll();
+
+        for (int i = 0; i < 50; i++) {
+            registry.rescheduleNow();
+        }
+
+        assertThat(executor.getQueue().size(), equalTo(1));
+        registry.close();
+        executor.shutdownNow();
+    }
+
+    @Test
+    public void ticker_mode_keeps_ticking_after_a_dispatch_predicate_throws() throws Exception {
+        ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor(1);
+        AtomicBoolean throwOnNextCheck = new AtomicBoolean();
+        DispatchPredicate throwingOnce = (key, dl) -> {
+            if (throwOnNextCheck.getAndSet(false)) {
+                throw new IllegalStateException("predicate failure");
+            }
+            return true;
+        };
+        ScheduledDataLoaderRegistry registry = tickerRegistry(executor, Duration.ofMillis(10), 1, throwingOnce);
+        registry.dispatchAll();
+        throwOnNextCheck.set(true);
+
+        DataLoader<String, String> loader = registry.getDataLoader("loader-0");
+        CompletableFuture<String> late = loader.load("late");
+
+        assertThat(late.get(2, TimeUnit.SECONDS), equalTo("late"));
+        registry.close();
+        executor.shutdownNow();
+    }
+
+    @Test
+    public void close_stops_the_ticker_but_dispatch_all_still_dispatches() throws Exception {
+        ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor(1);
+        ScheduledDataLoaderRegistry registry = tickerRegistry(executor, Duration.ofMillis(10), 1, alwaysDispatch);
+        registry.dispatchAll();
+        registry.close();
+        await().atMost(TWO_SECONDS).until(() -> executor.getQueue().isEmpty());
+
+        DataLoader<String, String> loader = registry.getDataLoader("loader-0");
+        CompletableFuture<String> afterClose = loader.load("after-close");
+        registry.dispatchAll();
+
+        assertThat(afterClose.getNow(null), equalTo("after-close"));
+        assertThat(executor.getQueue().size(), equalTo(0));
+        executor.shutdownNow();
+    }
+
+    private ScheduledDataLoaderRegistry tickerRegistry(ScheduledExecutorService executor, Duration schedule, int loaderCount, DispatchPredicate loaderPredicate) {
+        ScheduledDataLoaderRegistry.Builder builder = ScheduledDataLoaderRegistry.newScheduledRegistry()
+                .scheduledExecutorService(executor)
+                .tickerMode(true)
+                .schedule(schedule)
+                .dispatchPredicate(neverDispatch);
+        for (int i = 0; i < loaderCount; i++) {
+            builder.register("loader-" + i, idLoader("loader-" + i), loaderPredicate);
+        }
+        return builder.build();
+    }
+
     @Test
     public void executors_are_shutdown() {
         ScheduledDataLoaderRegistry registry = ScheduledDataLoaderRegistry.newScheduledRegistry().build();