pause scheduler for tests (#337)

Improve test reliability by:

* Add pause support for `SchedulerService` tests.
* Fixing issue where we were scheduling already scheduled
WorkflowSchedules

fixes #335

Author: devhawk

transact/src/main/java/dev/dbos/transact/execution/QueueService.java

@@ -38,11 +38,11 @@ public void setSpeedupForTest() {
     speedup = 0.01;
   }
 
-  public synchronized void pause() {
+  public void pause() {
     paused.set(true);
   }
 
-  public synchronized void unpause() {
+  public void unpause() {
     paused.set(false);
   }
 

transact/src/main/java/dev/dbos/transact/execution/SchedulerService.java

@@ -22,6 +22,7 @@
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 
@@ -67,6 +68,7 @@ public static void validateAnnotatedWorkflowSchedule(RegisteredWorkflow workflow
   private final SystemDatabase systemDatabase;
   private final Duration pollingInterval;
   private final AtomicReference<ScheduledExecutorService> execServiceRef = new AtomicReference<>();
+  private final AtomicBoolean paused = new AtomicBoolean(false);
   private final ConcurrentHashMap<String, ScheduledFuture<?>> workflowScheduleFutures =
       new ConcurrentHashMap<>();
 
@@ -98,6 +100,14 @@ public void close() {
     }
   }
 
+  public void pause() {
+    paused.set(true);
+  }
+
+  public void unpause() {
+    paused.set(false);
+  }
+
   private void pollWorkflowSchedules() {
     try {
       pollWorkflowSchedulesImpl();
@@ -124,20 +134,24 @@ private void pollWorkflowSchedulesImpl() {
     }
 
     // shut down any scheduled future that isn't in the list of current schedules
-    var scheduleIds = schedules.stream().map(s -> s.id()).collect(Collectors.toSet());
+    var currentIds = schedules.stream().map(WorkflowSchedule::id).collect(Collectors.toSet());
     for (var key : workflowScheduleFutures.keySet()) {
-      if (!scheduleIds.contains(key)) {
+      if (!currentIds.contains(key)) {
         cancelWorkflowSchedule(key);
       }
     }
 
     for (var schedule : schedules) {
-
+      var scheduleRunning = workflowScheduleFutures.containsKey(schedule.id());
       if (!schedule.isActive()) {
-        if (workflowScheduleFutures.containsKey(schedule.id())) {
+        // if the schedule is no longer active but we still have a scheduled future for it, cancel
+        // it
+        if (scheduleRunning) {
           cancelWorkflowSchedule(schedule.id());
         }
-      } else {
+      } else if (!scheduleRunning) {
+        // if the schedule is active but we don't yet have a scheduled future for it, schedule it
+        // now
         var optRegWf = dbosExecutor.getWorkflow(schedule.workflowName(), schedule.className());
         if (optRegWf.isEmpty()) {
           logger.error(
@@ -199,12 +213,17 @@ public void schedule() {
                     .ifPresent(
                         cronTime -> {
                           this.nextTime = cronTime.truncatedTo(ChronoUnit.SECONDS);
-                          // prevFuture should be null or a scheduled task that already fired.
-                          // but we still cancel it just to be sure
                           var prevFuture =
                               workflowScheduleFutures.put(
                                   wfSchedule.id(), scheduleTask(this.nextTime, this));
+                          // prevFuture should be null or a scheduled task that already fired.
+                          // cancel it anyway just to be sure
                           if (prevFuture != null) {
+                            if (!prevFuture.isDone()) {
+                              logger.debug(
+                                  "Previous scheduled task for {} has not yet completed",
+                                  wfSchedule.scheduleName());
+                            }
                             prevFuture.cancel(false);
                           }
                         });
@@ -219,6 +238,13 @@ public void run() {
                 }
 
                 try {
+                  if (paused.get()) {
+                    logger.debug(
+                        "Skipping scheduled workflow {} schedule {} because scheduler is paused",
+                        regWorkflow.fullyQualifiedName(),
+                        wfSchedule.scheduleName());
+                    return;
+                  }
                   var args = new Object[] {nextTime.toInstant(), wfSchedule.context()};
                   var workflowId = "sched-%s-%s".formatted(wfSchedule.scheduleName(), nextTime);
                   logger.debug(
@@ -284,6 +310,13 @@ public void run() {
 
               var scheduledTime = nextTime;
               try {
+                if (paused.get()) {
+                  logger.debug(
+                      "Skipping annotated workflow {} schedule {} because scheduler is paused",
+                      workflowName,
+                      swf.cron());
+                  return;
+                }
                 var args = new Object[] {scheduledTime.toInstant(), Instant.now()};
                 var workflowId = "sched-%s-%s".formatted(workflowName, scheduledTime);
                 logger.debug(

transact/src/test/java/dev/dbos/transact/client/ClientScheduleTest.java

@@ -3,6 +3,7 @@
 import static org.junit.jupiter.api.Assertions.*;
 
 import dev.dbos.transact.DBOS;
+import dev.dbos.transact.DBOSTestAccess;
 import dev.dbos.transact.config.DBOSConfig;
 import dev.dbos.transact.utils.PgContainer;
 import dev.dbos.transact.workflow.ScheduleStatus;
@@ -218,6 +219,9 @@ public void clientTriggerScheduleNotFound() {
 
   @Test
   public void clientBackfillSchedule() throws Exception {
+    // Pause scheduler to prevent interference with backfill results
+    DBOSTestAccess.getSchedulerService(dbos).pause();
+
     try (var client = pgContainer.dbosClient()) {
       // Use a cron that won't fire during the test (every minute)
       client.createSchedule(
@@ -241,6 +245,9 @@ public void clientBackfillSchedule() throws Exception {
 
   @Test
   public void clientBackfillScheduleEmptyWindow() {
+    // Pause scheduler to prevent interference with backfill results
+    DBOSTestAccess.getSchedulerService(dbos).pause();
+
     try (var client = pgContainer.dbosClient()) {
       client.createSchedule(
           "backfill-empty", workflowName(), className(), "0/1 * * * * *", null, false, null, null);

transact/src/test/java/dev/dbos/transact/scheduled/ScheduledWorkflowService.java

@@ -5,9 +5,12 @@
 import java.time.Instant;
 import java.util.List;
 import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.CountDownLatch;
 
 interface ScheduledWorkflowService {
   void scheduledRun(Instant scheduled, Object context);
+
+  void latchedRun(Instant scheduled, Object context);
 }
 
 class ScheduledWorkflowImpl implements ScheduledWorkflowService {
@@ -16,6 +19,7 @@ class ScheduledWorkflowImpl implements ScheduledWorkflowService {
   volatile Instant lastScheduled = null;
   volatile Object lastContext = null;
   final List<Instant> allScheduledTimes = new CopyOnWriteArrayList<>();
+  final CountDownLatch latch = new CountDownLatch(3);
 
   @Override
   @Workflow
@@ -26,6 +30,12 @@ public void scheduledRun(Instant scheduled, Object context) {
     allScheduledTimes.add(scheduled);
   }
 
+  @Override
+  @Workflow
+  public void latchedRun(Instant scheduled, Object context) {
+    latch.countDown();
+  }
+
   void reset() {
     counter = 0;
     lastScheduled = null;

transact/src/test/java/dev/dbos/transact/scheduled/WorkflowScheduleTest.java

@@ -3,6 +3,7 @@
 import static org.junit.jupiter.api.Assertions.*;
 
 import dev.dbos.transact.DBOS;
+import dev.dbos.transact.DBOSTestAccess;
 import dev.dbos.transact.config.DBOSConfig;
 import dev.dbos.transact.utils.PgContainer;
 import dev.dbos.transact.workflow.ScheduleStatus;
@@ -13,6 +14,7 @@
 import java.time.ZoneId;
 import java.time.temporal.ChronoUnit;
 import java.util.List;
+import java.util.concurrent.TimeUnit;
 
 import com.zaxxer.hikari.HikariDataSource;
 import org.junit.jupiter.api.AutoClose;
@@ -401,6 +403,9 @@ public void backfillScheduleNotFound() {
   public void backfillScheduleCorrectTimes() throws Exception {
     var impl = registerAndLaunch();
 
+    // Pause scheduler to prevent interference with backfill results
+    DBOSTestAccess.getSchedulerService(dbos).pause();
+
     // Every minute at second 0: "0 * * * * *" (6-field cron)
     dbos.createSchedule(
         "backfill-correct", workflowName(), className(), "0 * * * * *", null, false, null, null);
@@ -430,6 +435,9 @@ public void backfillScheduleCorrectTimes() throws Exception {
   public void backfillScheduleHourly() throws Exception {
     var impl = registerAndLaunch();
 
+    // Pause scheduler to prevent interference with backfill results
+    DBOSTestAccess.getSchedulerService(dbos).pause();
+
     // Every hour at minute 0: "0 0 * * * *" (6-field cron, runs at top of each hour)
     dbos.createSchedule(
         "backfill-hourly", workflowName(), className(), "0 0 * * * *", null, false, null, null);
@@ -461,6 +469,9 @@ public void backfillScheduleHourly() throws Exception {
   public void backfillScheduleDaily() throws Exception {
     var impl = registerAndLaunch();
 
+    // Pause scheduler to prevent interference with backfill results
+    DBOSTestAccess.getSchedulerService(dbos).pause();
+
     // Every day at midnight: "0 0 0 * * *" (6-field cron)
     dbos.createSchedule(
         "backfill-daily", workflowName(), className(), "0 0 0 * * *", null, false, null, null);
@@ -535,18 +546,18 @@ public void scheduleRunsAfterPolling() throws Exception {
     var impl = registerAndLaunch();
 
     dbos.createSchedule(
-        "run-sched", workflowName(), className(), "0/1 * * * * *", null, false, null, null);
+        "run-sched", "latchedRun", className(), "0/1 * * * * *", null, false, null, null);
 
     // Verify schedule was created and is active
     var schedule = dbos.getSchedule("run-sched");
     assertTrue(schedule.isPresent(), "Schedule should be created");
     assertEquals(ScheduleStatus.ACTIVE, schedule.get().status());
 
-    // Allow time for scheduler polls + workflow executions
-    // Schedule triggers at second boundary, so wait long enough for multiple executions
-    Thread.sleep(8000);
-
-    assertTrue(impl.counter >= 2, "Expected at least 2 executions, got " + impl.counter);
-    assertTrue(impl.counter <= 10, "Expected at most 10 executions, got " + impl.counter);
+    // latch initialized with 3 count, with each execution counting down once.
+    // Wait for all 3 counts to be released, which indicates the workflow ran at least 3 times
+    // (scheduler should run it every second).
+    assertTrue(
+        impl.latch.await(5, TimeUnit.SECONDS),
+        "Expected latch to count down to zero within 5 seconds");
   }
 }

transact/src/test/java/dev/dbos/transact/utils/PgContainer.java

@@ -74,7 +74,7 @@ public PgContainer() {
 
   @Override
   public void close() throws Exception {
-    // drop a database we created and return the container too the pool
+    // drop the database we created and return the container too the pool
     var _jdbcUrl = pgContainer.getJdbcUrl();
     try (var conn = DriverManager.getConnection(_jdbcUrl, username(), password());
         var stmt = conn.createStatement()) {