Better Scheduler (#484)

Update the scheduler to enqueue a task only after a short random
interval has passed, and only if the task is not already enqueued. This
prevents the problem in a distributed setting of hundreds of workers
attempting to enqueue a scheduled task simultaneously (which is safe but
not performant).

Author: kraftp

dbos/_scheduler.py

@@ -1,3 +1,4 @@
+import random
 import threading
 import traceback
 from datetime import datetime, timezone
@@ -15,28 +16,40 @@
 
 ScheduledWorkflow = Callable[[datetime, datetime], None]
 
-scheduler_queue: Queue
-
 
 def scheduler_loop(
     func: ScheduledWorkflow, cron: str, stop_event: threading.Event
 ) -> None:
+    from dbos._dbos import _get_dbos_instance
+
+    dbos = _get_dbos_instance()
+    scheduler_queue = dbos._registry.get_internal_queue()
     try:
         iter = croniter(cron, datetime.now(timezone.utc), second_at_beginning=True)
-    except Exception as e:
+    except Exception:
         dbos_logger.error(
             f'Cannot run scheduled function {get_dbos_func_name(func)}. Invalid crontab "{cron}"'
         )
+        raise
     while not stop_event.is_set():
-        nextExecTime = iter.get_next(datetime)
-        sleepTime = nextExecTime - datetime.now(timezone.utc)
-        if stop_event.wait(timeout=sleepTime.total_seconds()):
+        next_exec_time = iter.get_next(datetime)
+        sleep_time = (next_exec_time - datetime.now(timezone.utc)).total_seconds()
+        sleep_time = max(0, sleep_time)
+        # To prevent a "thundering herd" problem in a distributed setting,
+        # apply jitter of up to 10% the sleep time, capped at 10 seconds
+        max_jitter = min(sleep_time / 10, 10)
+        jitter = random.uniform(0, max_jitter)
+        if stop_event.wait(timeout=sleep_time + jitter):
             return
         try:
-            with SetWorkflowID(
-                f"sched-{get_dbos_func_name(func)}-{nextExecTime.isoformat()}"
-            ):
-                scheduler_queue.enqueue(func, nextExecTime, datetime.now(timezone.utc))
+            workflowID = (
+                f"sched-{get_dbos_func_name(func)}-{next_exec_time.isoformat()}"
+            )
+            if not dbos._sys_db.get_workflow_status(workflowID):
+                with SetWorkflowID(workflowID):
+                    scheduler_queue.enqueue(
+                        func, next_exec_time, datetime.now(timezone.utc)
+                    )
         except Exception:
             dbos_logger.warning(
                 f"Exception encountered in scheduler thread: {traceback.format_exc()})"
@@ -49,13 +62,10 @@ def scheduled(
     def decorator(func: ScheduledWorkflow) -> ScheduledWorkflow:
         try:
             croniter(cron, datetime.now(timezone.utc), second_at_beginning=True)
-        except Exception as e:
+        except Exception:
             raise ValueError(
                 f'Invalid crontab "{cron}" for scheduled function function {get_dbos_func_name(func)}.'
             )
-
-        global scheduler_queue
-        scheduler_queue = dbosreg.get_internal_queue()
         stop_event = threading.Event()
         dbosreg.register_poller(stop_event, scheduler_loop, func, cron, stop_event)
         return func

tests/test_scheduler.py

@@ -102,7 +102,7 @@ def test_workflow(scheduled: datetime, actual: datetime) -> None:
         wf_counter += 1
 
     time.sleep(5)
-    assert wf_counter > 2 and wf_counter <= 5
+    assert wf_counter > 1 and wf_counter <= 5
 
 
 def test_appdb_downtime(dbos: DBOS, skip_with_sqlite: None) -> None:
@@ -123,7 +123,7 @@ def test_workflow(scheduled: datetime, actual: datetime) -> None:
     time.sleep(2)
     assert dbos._app_db
     simulate_db_restart(dbos._app_db.engine, 2)
-    time.sleep(2)
+    time.sleep(3)
     assert wf_counter > 2
 
 
@@ -138,7 +138,7 @@ def test_workflow(scheduled: datetime, actual: datetime) -> None:
 
     time.sleep(2)
     simulate_db_restart(dbos._sys_db.engine, 2)
-    time.sleep(2)
+    time.sleep(3)
     # We know there should be at least 2 occurrences from the 4 seconds when the DB was up.
     #  There could be more than 4, depending on the pace the machine...
     assert wf_counter >= 2
@@ -154,7 +154,7 @@ def test_transaction(scheduled: datetime, actual: datetime) -> None:
         txn_counter += 1
 
     time.sleep(5)
-    assert txn_counter > 2 and txn_counter <= 5
+    assert txn_counter > 1 and txn_counter <= 5
 
 
 def test_scheduled_step(dbos: DBOS) -> None:
@@ -205,8 +205,8 @@ def test_transaction() -> None:
         nonlocal txn_counter
         txn_counter += 1
 
-    time.sleep(3)
-    assert wf_counter >= 1 and wf_counter <= 3
+    time.sleep(4)
+    assert wf_counter >= 1 and wf_counter <= 4
     max_tries = 10
     for i in range(max_tries):
         try:
@@ -223,7 +223,7 @@ def test_transaction() -> None:
         evt.set()
 
     # Wait for workflows to finish
-    time.sleep(2)
+    time.sleep(3)
 
     dbos._sys_db.update_workflow_outcome(workflow_id, "PENDING")