Repository: kagkarlsson/db-scheduler (java)

Issue:
enableImmediateExecution() does not work for rescheduling in a CompletionHandler
`builder.enableImmediateExecution()` does not work for rescheduling in a CompletionHandler, only when using the `SchedulerClient`. Should be fixed to allow for multi-step jobs 


Discussion:

Patch:
diff --git a/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/ClientEvent.java b/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/ClientEvent.java
--- a/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/ClientEvent.java
+++ b/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/ClientEvent.java
@@ -21,7 +21,7 @@
 
 public class ClientEvent {
 
-    enum EventType {
+    public enum EventType {
         SCHEDULE,
         RESCHEDULE,
         CANCEL
diff --git a/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/ExecutePicked.java b/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/ExecutePicked.java
--- a/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/ExecutePicked.java
+++ b/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/ExecutePicked.java
@@ -35,6 +35,7 @@ class ExecutePicked implements Runnable {
     private static final Logger LOG = LoggerFactory.getLogger(ExecutePicked.class);
     private final Executor executor;
     private final TaskRepository taskRepository;
+    private SchedulerClientEventListener earlyExecutionListener;
     private final SchedulerClient schedulerClient;
     private final StatsRegistry statsRegistry;
     private final TaskResolver taskResolver;
@@ -43,11 +44,12 @@ class ExecutePicked implements Runnable {
     private final Clock clock;
     private final Execution pickedExecution;
 
-    public ExecutePicked(Executor executor, TaskRepository taskRepository, SchedulerClient schedulerClient, StatsRegistry statsRegistry,
+    public ExecutePicked(Executor executor, TaskRepository taskRepository, SchedulerClientEventListener earlyExecutionListener, SchedulerClient schedulerClient, StatsRegistry statsRegistry,
                          TaskResolver taskResolver, SchedulerState schedulerState, ConfigurableLogger failureLogger,
                          Clock clock, Execution pickedExecution) {
         this.executor = executor;
         this.taskRepository = taskRepository;
+        this.earlyExecutionListener = earlyExecutionListener;
         this.schedulerClient = schedulerClient;
         this.statsRegistry = statsRegistry;
         this.taskResolver = taskResolver;
@@ -99,7 +101,7 @@ private void executePickedExecution(Execution execution) {
     private void complete(CompletionHandler completion, Execution execution, Instant executionStarted) {
         ExecutionComplete completeEvent = ExecutionComplete.success(execution, executionStarted, clock.now());
         try {
-            completion.complete(completeEvent, new ExecutionOperations(taskRepository, execution));
+            completion.complete(completeEvent, new ExecutionOperations(taskRepository, earlyExecutionListener, execution));
             statsRegistry.registerSingleCompletedExecution(completeEvent);
         } catch (Throwable e) {
             statsRegistry.register(StatsRegistry.SchedulerStatsEvent.COMPLETIONHANDLER_ERROR);
@@ -115,7 +117,7 @@ private void failure(Task task, Execution execution, Throwable cause, Instant ex
 
         ExecutionComplete completeEvent = ExecutionComplete.failure(execution, executionStarted, clock.now(), cause);
         try {
-            task.getFailureHandler().onFailure(completeEvent, new ExecutionOperations(taskRepository, execution));
+            task.getFailureHandler().onFailure(completeEvent, new ExecutionOperations(taskRepository, earlyExecutionListener, execution));
             statsRegistry.registerSingleCompletedExecution(completeEvent);
         } catch (Throwable e) {
             statsRegistry.register(StatsRegistry.SchedulerStatsEvent.FAILUREHANDLER_ERROR);
diff --git a/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/FetchCandidates.java b/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/FetchCandidates.java
--- a/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/FetchCandidates.java
+++ b/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/FetchCandidates.java
@@ -32,6 +32,7 @@ public class FetchCandidates implements PollStrategy {
     private final Executor executor;
     private final TaskRepository taskRepository;
     private final SchedulerClient schedulerClient;
+    private SchedulerClientEventListener earlyExecutionListener;
     private final StatsRegistry statsRegistry;
     private final SchedulerState schedulerState;
     private final ConfigurableLogger failureLogger;
@@ -44,12 +45,13 @@ public class FetchCandidates implements PollStrategy {
     private final int upperLimit;
 
     public FetchCandidates(Executor executor, TaskRepository taskRepository, SchedulerClient schedulerClient,
-                           int threadpoolSize, StatsRegistry statsRegistry, SchedulerState schedulerState,
+                           SchedulerClientEventListener earlyExecutionListener, int threadpoolSize, StatsRegistry statsRegistry, SchedulerState schedulerState,
                            ConfigurableLogger failureLogger, TaskResolver taskResolver, Clock clock,
                            PollingStrategyConfig pollingStrategyConfig, Runnable triggerCheckForNewExecutions) {
         this.executor = executor;
         this.taskRepository = taskRepository;
         this.schedulerClient = schedulerClient;
+        this.earlyExecutionListener = earlyExecutionListener;
         this.statsRegistry = statsRegistry;
         this.schedulerState = schedulerState;
         this.failureLogger = failureLogger;
@@ -82,7 +84,7 @@ public void run() {
             executor.addToQueue(
                 () -> {
                     final Optional<Execution> candidate = new PickDue(e, newDueBatch).call();
-                    candidate.ifPresent(picked -> new ExecutePicked(executor, taskRepository, schedulerClient, statsRegistry,
+                    candidate.ifPresent(picked -> new ExecutePicked(executor, taskRepository, earlyExecutionListener, schedulerClient, statsRegistry,
                         taskResolver, schedulerState, failureLogger,
                         clock, picked).run());
                 },
diff --git a/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/LockAndFetchCandidates.java b/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/LockAndFetchCandidates.java
--- a/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/LockAndFetchCandidates.java
+++ b/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/LockAndFetchCandidates.java
@@ -30,6 +30,7 @@ public class LockAndFetchCandidates implements PollStrategy {
     private final Executor executor;
     private final TaskRepository taskRepository;
     private final SchedulerClient schedulerClient;
+    private SchedulerClientEventListener earlyExecutionListener;
     private final StatsRegistry statsRegistry;
     private final TaskResolver taskResolver;
     private final SchedulerState schedulerState;
@@ -42,12 +43,13 @@ public class LockAndFetchCandidates implements PollStrategy {
     private AtomicBoolean moreExecutionsInDatabase = new AtomicBoolean(false);
 
     public LockAndFetchCandidates(Executor executor, TaskRepository taskRepository, SchedulerClient schedulerClient,
-                                  int threadpoolSize, StatsRegistry statsRegistry, SchedulerState schedulerState,
+                                  SchedulerClientEventListener earlyExecutionListener, int threadpoolSize, StatsRegistry statsRegistry, SchedulerState schedulerState,
                                   ConfigurableLogger failureLogger, TaskResolver taskResolver, Clock clock,
                                   PollingStrategyConfig pollingStrategyConfig, Runnable triggerCheckForNewExecutions) {
         this.executor = executor;
         this.taskRepository = taskRepository;
         this.schedulerClient = schedulerClient;
+        this.earlyExecutionListener = earlyExecutionListener;
         this.statsRegistry = statsRegistry;
         this.taskResolver = taskResolver;
         this.schedulerState = schedulerState;
@@ -87,7 +89,7 @@ public void run() {
 
         for (Execution picked : pickedExecutions) {
             executor.addToQueue(
-                new ExecutePicked(executor, taskRepository, schedulerClient, statsRegistry,
+                new ExecutePicked(executor, taskRepository, earlyExecutionListener, schedulerClient, statsRegistry,
                     taskResolver, schedulerState, failureLogger,
                     clock, picked),
                 () -> {
diff --git a/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/Scheduler.java b/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/Scheduler.java
--- a/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/Scheduler.java
+++ b/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/Scheduler.java
@@ -39,7 +39,6 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 
 import static com.github.kagkarlsson.scheduler.ExecutorUtils.defaultThreadFactoryWithPrefix;
@@ -69,6 +68,7 @@ public class Scheduler implements SchedulerClient {
     private final Waiter heartbeatWaiter;
     final SettableSchedulerState schedulerState = new SettableSchedulerState();
     final ConfigurableLogger failureLogger;
+    private SchedulerClientEventListener earlyExecutionListener;
 
     protected Scheduler(Clock clock, TaskRepository schedulerTaskRepository, TaskRepository clientTaskRepository, TaskResolver taskResolver, int threadpoolSize, ExecutorService executorService, SchedulerName schedulerName,
                         Waiter executeDueWaiter, Duration heartbeatInterval, boolean enableImmediateExecution, StatsRegistry statsRegistry, PollingStrategyConfig pollingStrategyConfig, Duration deleteUnresolvedAfter, Duration shutdownMaxWait,
@@ -88,15 +88,15 @@ protected Scheduler(Clock clock, TaskRepository schedulerTaskRepository, TaskRep
         this.statsRegistry = statsRegistry;
         this.dueExecutor = Executors.newSingleThreadExecutor(defaultThreadFactoryWithPrefix(THREAD_PREFIX + "-execute-due-"));
         this.housekeeperExecutor = Executors.newScheduledThreadPool(3, defaultThreadFactoryWithPrefix(THREAD_PREFIX + "-housekeeper-"));
-        SchedulerClientEventListener earlyExecutionListener = (enableImmediateExecution ? new TriggerCheckForDueExecutions(schedulerState, clock, executeDueWaiter) : SchedulerClientEventListener.NOOP);
+        earlyExecutionListener = (enableImmediateExecution ? new TriggerCheckForDueExecutions(schedulerState, clock, executeDueWaiter) : SchedulerClientEventListener.NOOP);
         delegate = new StandardSchedulerClient(clientTaskRepository, earlyExecutionListener, clock);
         this.failureLogger = ConfigurableLogger.create(LOG, logLevel, logStackTrace);
 
         if (pollingStrategyConfig.type == PollingStrategyConfig.Type.LOCK_AND_FETCH) {
             schedulerTaskRepository.checkSupportsLockAndFetch();
-            executeDueStrategy = new LockAndFetchCandidates(executor, schedulerTaskRepository, this, threadpoolSize, statsRegistry, schedulerState, failureLogger, taskResolver, clock, pollingStrategyConfig, this::triggerCheckForDueExecutions);
+            executeDueStrategy = new LockAndFetchCandidates(executor, schedulerTaskRepository, this, earlyExecutionListener, threadpoolSize, statsRegistry, schedulerState, failureLogger, taskResolver, clock, pollingStrategyConfig, this::triggerCheckForDueExecutions);
         } else if (pollingStrategyConfig.type == PollingStrategyConfig.Type.FETCH) {
-            executeDueStrategy = new FetchCandidates(executor, schedulerTaskRepository, this, threadpoolSize, statsRegistry, schedulerState, failureLogger, taskResolver, clock, pollingStrategyConfig, this::triggerCheckForDueExecutions);
+            executeDueStrategy = new FetchCandidates(executor, schedulerTaskRepository, this, earlyExecutionListener, threadpoolSize, statsRegistry, schedulerState, failureLogger, taskResolver, clock, pollingStrategyConfig, this::triggerCheckForDueExecutions);
         } else {
             throw new IllegalArgumentException("Unknown polling-strategy type: " + pollingStrategyConfig.type);
         }
@@ -266,7 +266,7 @@ protected void detectDeadExecutions() {
                     Optional<Task> task = taskResolver.resolve(execution.taskInstance.getTaskName());
                     if (task.isPresent()) {
                         statsRegistry.register(SchedulerStatsEvent.DEAD_EXECUTION);
-                        task.get().getDeadExecutionHandler().deadExecution(execution, new ExecutionOperations(schedulerTaskRepository, execution));
+                        task.get().getDeadExecutionHandler().deadExecution(execution, new ExecutionOperations(schedulerTaskRepository, earlyExecutionListener, execution));
                     } else {
                         LOG.error("Failed to find implementation for task with name '{}' for detected dead execution. Either delete the execution from the databaser, or add an implementation for it.", execution.taskInstance.getTaskName());
                     }
diff --git a/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/task/ExecutionOperations.java b/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/task/ExecutionOperations.java
--- a/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/task/ExecutionOperations.java
+++ b/db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/task/ExecutionOperations.java
@@ -15,6 +15,10 @@
  */
 package com.github.kagkarlsson.scheduler.task;
 
+import com.github.kagkarlsson.scheduler.ClientEvent;
+import com.github.kagkarlsson.scheduler.ClientEvent.ClientEventContext;
+import com.github.kagkarlsson.scheduler.ClientEvent.EventType;
+import com.github.kagkarlsson.scheduler.SchedulerClientEventListener;
 import com.github.kagkarlsson.scheduler.TaskRepository;
 
 import java.time.Instant;
@@ -22,10 +26,12 @@
 public class ExecutionOperations<T> {
 
     private final TaskRepository taskRepository;
+    private final SchedulerClientEventListener earlyExecutionListener;
     private final Execution execution;
 
-    public ExecutionOperations(TaskRepository taskRepository, Execution execution) {
+    public ExecutionOperations(TaskRepository taskRepository, SchedulerClientEventListener earlyExecutionListener, Execution execution) {
         this.taskRepository = taskRepository;
+        this.earlyExecutionListener = earlyExecutionListener;
         this.execution = execution;
     }
 
@@ -39,7 +45,7 @@ public void reschedule(ExecutionComplete completed, Instant nextExecutionTime) {
         } else {
             taskRepository.reschedule(execution, nextExecutionTime, execution.lastSuccess, completed.getTimeDone(), execution.consecutiveFailures + 1);
         }
-
+        hintExecutionScheduled(completed, nextExecutionTime);
     }
 
     public void reschedule(ExecutionComplete completed, Instant nextExecutionTime, T newData) {
@@ -48,6 +54,13 @@ public void reschedule(ExecutionComplete completed, Instant nextExecutionTime, T
         } else {
             taskRepository.reschedule(execution, nextExecutionTime, newData, execution.lastSuccess, completed.getTimeDone(), execution.consecutiveFailures + 1);
         }
+        hintExecutionScheduled(completed, nextExecutionTime);
+    }
+
+    private void hintExecutionScheduled(ExecutionComplete completed, Instant nextExecutionTime) {
+        // Hint that a new execution was scheduled in-case we want to go check for it immediately
+        earlyExecutionListener.newEvent(new ClientEvent(
+            new ClientEventContext(EventType.RESCHEDULE, completed.getExecution().taskInstance, nextExecutionTime)));
     }
 
 }


Test patch:
diff --git a/db-scheduler/src/test/java/com/github/kagkarlsson/scheduler/TestSchedules.java b/db-scheduler/src/test/java/com/github/kagkarlsson/scheduler/TestSchedules.java
new file mode 100644
--- /dev/null
+++ b/db-scheduler/src/test/java/com/github/kagkarlsson/scheduler/TestSchedules.java
@@ -0,0 +1,23 @@
+package com.github.kagkarlsson.scheduler;
+
+import com.github.kagkarlsson.scheduler.task.ExecutionComplete;
+import com.github.kagkarlsson.scheduler.task.schedule.Schedule;
+
+import java.time.Instant;
+
+public class TestSchedules {
+
+    public static Schedule now() {
+        return new Schedule() {
+            @Override
+            public Instant getNextExecutionTime(ExecutionComplete executionComplete) {
+                return executionComplete.getTimeDone();
+            }
+
+            @Override
+            public boolean isDeterministic() {
+                return false;
+            }
+        };
+    }
+}
diff --git a/db-scheduler/src/test/java/com/github/kagkarlsson/scheduler/functional/ImmediateExecutionTest.java b/db-scheduler/src/test/java/com/github/kagkarlsson/scheduler/functional/ImmediateExecutionTest.java
--- a/db-scheduler/src/test/java/com/github/kagkarlsson/scheduler/functional/ImmediateExecutionTest.java
+++ b/db-scheduler/src/test/java/com/github/kagkarlsson/scheduler/functional/ImmediateExecutionTest.java
@@ -1,15 +1,13 @@
 package com.github.kagkarlsson.scheduler.functional;
 
 import co.unruly.matchers.TimeMatchers;
-import com.github.kagkarlsson.scheduler.DbUtils;
-import com.github.kagkarlsson.scheduler.EmbeddedPostgresqlExtension;
-import com.github.kagkarlsson.scheduler.Scheduler;
-import com.github.kagkarlsson.scheduler.SchedulerName;
-import com.github.kagkarlsson.scheduler.StopSchedulerExtension;
-import com.github.kagkarlsson.scheduler.TestTasks;
+import com.github.kagkarlsson.scheduler.*;
 import com.github.kagkarlsson.scheduler.helper.TestableRegistry;
-import com.github.kagkarlsson.scheduler.task.ExecutionComplete;
+import com.github.kagkarlsson.scheduler.task.*;
+import com.github.kagkarlsson.scheduler.task.helper.CustomTask;
 import com.github.kagkarlsson.scheduler.task.helper.OneTimeTask;
+import com.github.kagkarlsson.scheduler.task.helper.Tasks;
+import com.github.kagkarlsson.scheduler.task.schedule.Schedule;
 import com.github.kagkarlsson.scheduler.testhelper.SettableClock;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
@@ -20,10 +18,13 @@
 import java.time.Duration;
 import java.time.Instant;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.core.Is.is;
+import static org.junit.Assert.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
 
 
 public class ImmediateExecutionTest {
@@ -40,9 +41,9 @@ public void setUp() {
         clock = new SettableClock();
     }
 
-    @RepeatedTest(10)
+    @Test
     public void test_immediate_execution() {
-        Assertions.assertTimeoutPreemptively(Duration.ofSeconds(10), () -> {
+        Assertions.assertTimeoutPreemptively(Duration.ofSeconds(2), () -> {
 
             Instant now = Instant.now();
             OneTimeTask<Void> task = TestTasks.oneTime("onetime-a", Void.class, TestTasks.DO_NOTHING);
@@ -51,15 +52,7 @@ public void test_immediate_execution() {
 
             TestableRegistry registry = TestableRegistry.create().waitConditions(executeDueCondition, completedCondition).build();
 
-            Scheduler scheduler = Scheduler.create(postgres.getDataSource(), task)
-                .pollingInterval(Duration.ofMinutes(1))
-                .enableImmediateExecution()
-                .schedulerName(new SchedulerName.Fixed("test"))
-                .statsRegistry(registry)
-                .build();
-            stopScheduler.register(scheduler);
-
-            scheduler.start();
+            Scheduler scheduler = createAndStartScheduler(task, registry);
             executeDueCondition.waitFor();
 
             scheduler.schedule(task.instance("1"), clock.now());
@@ -67,13 +60,63 @@ public void test_immediate_execution() {
 
             List<ExecutionComplete> completed = registry.getCompleted();
             assertThat(completed, hasSize(1));
-            completed.stream().forEach(e -> {
+            completed.forEach(e -> {
+                assertThat(e.getResult(), is(ExecutionComplete.Result.OK));
+                Duration durationUntilExecuted = Duration.between(now, e.getTimeDone());
+                assertThat(durationUntilExecuted, TimeMatchers.shorterThan(Duration.ofSeconds(1)));
+            });
+            registry.assertNoFailures();
+        });
+    }
+
+    @Test
+    public void test_immediate_execution_by_completion_handler() {
+        Assertions.assertTimeoutPreemptively(Duration.ofSeconds(2), () -> {
+
+            Instant now = Instant.now();
+            CustomTask<Integer> task = Tasks.custom("two-step", Integer.class)
+                .execute((taskInstance, executionContext) -> {
+                    if (taskInstance.getData() >= 2) {
+                        return new CompletionHandler.OnCompleteRemove<>();
+                    } else {
+                        return new CompletionHandler.OnCompleteReschedule<>(TestSchedules.now(), taskInstance.getData() + 1);
+                    }
+                });
+
+            TestableRegistry.Condition completedCondition = TestableRegistry.Conditions.completed(2);
+            TestableRegistry.Condition executeDueCondition = TestableRegistry.Conditions.ranExecuteDue(1);
+
+            TestableRegistry registry = TestableRegistry.create().waitConditions(executeDueCondition, completedCondition).build();
+
+            Scheduler scheduler = createAndStartScheduler(task, registry);
+            executeDueCondition.waitFor();
+
+            scheduler.schedule(task.instance("id1", 1), clock.now());
+            completedCondition.waitFor();
+
+            List<ExecutionComplete> completed = registry.getCompleted();
+            assertThat(completed, hasSize(2));
+            completed.forEach(e -> {
                 assertThat(e.getResult(), is(ExecutionComplete.Result.OK));
                 Duration durationUntilExecuted = Duration.between(now, e.getTimeDone());
                 assertThat(durationUntilExecuted, TimeMatchers.shorterThan(Duration.ofSeconds(1)));
             });
+            assertEquals(scheduler.getScheduledExecutions().size(), 0);
             registry.assertNoFailures();
         });
     }
 
+    private Scheduler createAndStartScheduler(Task task, TestableRegistry registry) {
+        Scheduler scheduler = Scheduler.create(postgres.getDataSource(), task)
+            .pollingInterval(Duration.ofMinutes(1))
+            .enableImmediateExecution()
+            .schedulerName(new SchedulerName.Fixed("test"))
+            .statsRegistry(registry)
+            .build();
+        stopScheduler.register(scheduler);
+
+        scheduler.start();
+        return scheduler;
+    }
+
 }