dbos 2.1.0a1__tar.gz → 2.1.0a3__tar.gz

{dbos-2.1.0a1 → dbos-2.1.0a3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dbos
-Version: 2.1.0a1
+Version: 2.1.0a3
 Summary: Ultra-lightweight durable execution in Python
 Author-Email: "DBOS, Inc." <contact@dbos.dev>
 License: MIT

{dbos-2.1.0a1 → dbos-2.1.0a3}/dbos/_core.py

@@ -1278,21 +1278,24 @@ def recv(dbos: "DBOS", topic: Optional[str] = None, timeout_seconds: float = 60)
 def set_event(dbos: "DBOS", key: str, value: Any) -> None:
     cur_ctx = get_local_dbos_context()
     if cur_ctx is not None:
-        # Must call it within a workflow
-        assert (
-            cur_ctx.is_workflow()
-        ), "set_event() must be called from within a workflow"
-        attributes: TracedAttributes = {
-            "name": "set_event",
-        }
-        with EnterDBOSStep(attributes):
-            ctx = assert_current_dbos_context()
-            dbos._sys_db.set_event(
-                ctx.workflow_id, ctx.curr_step_function_id, key, value
+        if cur_ctx.is_workflow():
+            # If called from a workflow function, run as a step
+            attributes: TracedAttributes = {
+                "name": "set_event",
+            }
+            with EnterDBOSStep(attributes):
+                ctx = assert_current_dbos_context()
+                dbos._sys_db.set_event_from_workflow(
+                    ctx.workflow_id, ctx.curr_step_function_id, key, value
+                )
+        elif cur_ctx.is_step():
+            dbos._sys_db.set_event_from_step(cur_ctx.workflow_id, key, value)
+        else:
+            raise DBOSException(
+                "set_event() must be called from within a workflow or step"
             )
     else:
-        # Cannot call it from outside of a workflow
-        raise DBOSException("set_event() must be called from within a workflow")
+        raise DBOSException("set_event() must be called from within a workflow or step")
 
 
 def get_event(

{dbos-2.1.0a1 → dbos-2.1.0a3}/dbos/_scheduler.py

@@ -1,3 +1,4 @@
+import random
 import threading
 import traceback
 from datetime import datetime, timezone
@@ -15,28 +16,40 @@ from ._registrations import get_dbos_func_name
 
 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

{dbos-2.1.0a1 → dbos-2.1.0a3}/dbos/_sys_db.py

@@ -1524,7 +1524,7 @@ class SystemDatabase(ABC):
         return duration
 
     @db_retry()
-    def set_event(
+    def set_event_from_workflow(
         self,
         workflow_uuid: str,
         function_id: int,
@@ -1566,6 +1566,26 @@ class SystemDatabase(ABC):
             }
             self._record_operation_result_txn(output, conn=c)
 
+    def set_event_from_step(
+        self,
+        workflow_uuid: str,
+        key: str,
+        message: Any,
+    ) -> None:
+        with self.engine.begin() as c:
+            c.execute(
+                self.dialect.insert(SystemSchema.workflow_events)
+                .values(
+                    workflow_uuid=workflow_uuid,
+                    key=key,
+                    value=_serialization.serialize(message),
+                )
+                .on_conflict_do_update(
+                    index_elements=["workflow_uuid", "key"],
+                    set_={"value": _serialization.serialize(message)},
+                )
+            )
+
     def get_all_events(self, workflow_id: str) -> Dict[str, Any]:
         """
         Get all events currently present for a workflow ID.

{dbos-2.1.0a1 → dbos-2.1.0a3}/pyproject.toml

@@ -34,7 +34,7 @@ classifiers = [
     "Topic :: Software Development :: Libraries :: Python Modules",
     "Framework :: AsyncIO",
 ]
-version = "2.1.0a1"
+version = "2.1.0a3"
 
 [project.license]
 text = "MIT"

{dbos-2.1.0a1 → dbos-2.1.0a3}/tests/test_dbos.py

@@ -965,41 +965,57 @@ def test_send_recv_temp_wf(dbos: DBOS) -> None:
 def test_set_get_events(dbos: DBOS) -> None:
     @DBOS.workflow()
     def test_setevent_workflow() -> None:
-        dbos.set_event("key1", "value1")
-        dbos.set_event("key2", "value2")
-        dbos.set_event("key3", None)
+        DBOS.set_event("key1", "value1")
+        DBOS.set_event("key2", "value2")
+        DBOS.set_event("key3", None)
+        set_event_step()
+
+    @DBOS.step()
+    def set_event_step() -> None:
+        DBOS.set_event("key4", "value4")
 
     @DBOS.workflow()
     def test_getevent_workflow(
-        target_uuid: str, key: str, timeout_seconds: float = 10
+        target_uuid: str, key: str, timeout: float = 0.0
     ) -> Optional[str]:
-        msg = dbos.get_event(target_uuid, key, timeout_seconds)
+        msg = dbos.get_event(target_uuid, key, timeout)
         return str(msg) if msg is not None else None
 
-    wfuuid = str(uuid.uuid4())
-    with SetWorkflowID(wfuuid):
+    wfid = str(uuid.uuid4())
+    with SetWorkflowID(wfid):
         test_setevent_workflow()
-    with SetWorkflowID(wfuuid):
+    with SetWorkflowID(wfid):
         test_setevent_workflow()
 
-    value1 = test_getevent_workflow(wfuuid, "key1")
+    value1 = test_getevent_workflow(wfid, "key1")
     assert value1 == "value1"
 
-    value2 = test_getevent_workflow(wfuuid, "key2")
+    value2 = test_getevent_workflow(wfid, "key2")
     assert value2 == "value2"
 
     # Run getEvent outside of a workflow
-    value1 = dbos.get_event(wfuuid, "key1")
+    value1 = DBOS.get_event(wfid, "key1", 0)
     assert value1 == "value1"
 
-    value2 = dbos.get_event(wfuuid, "key2")
+    value2 = DBOS.get_event(wfid, "key2", 0)
     assert value2 == "value2"
 
     begin_time = time.time()
-    value3 = test_getevent_workflow(wfuuid, "key3")
+    value3 = test_getevent_workflow(wfid, "key3")
     assert value3 is None
-    duration = time.time() - begin_time
-    assert duration < 1  # None is from the event not from the timeout
+
+    value4 = DBOS.get_event(wfid, "key4", 0)
+    assert value4 == "value4"
+
+    steps = DBOS.list_workflow_steps(wfid)
+    assert len(steps) == 4
+    assert (
+        steps[0]["function_name"]
+        == steps[1]["function_name"]
+        == steps[2]["function_name"]
+        == "DBOS.setEvent"
+    )
+    assert steps[3]["function_name"] == set_event_step.__qualname__
 
     # Test OAOO
     timeout_uuid = str(uuid.uuid4())

{dbos-2.1.0a1 → dbos-2.1.0a3}/tests/test_scheduler.py

@@ -102,7 +102,7 @@ def test_scheduled_workflow(dbos: DBOS) -> 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_appdb_downtime(dbos: DBOS, skip_with_sqlite: None) -> 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_sysdb_downtime(dbos: DBOS, skip_with_sqlite: None) -> 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_scheduled_transaction(dbos: DBOS) -> 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_scheduler_oaoo(dbos: DBOS) -> 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_scheduler_oaoo(dbos: DBOS) -> None:
         evt.set()
 
     # Wait for workflows to finish
-    time.sleep(2)
+    time.sleep(3)
 
     dbos._sys_db.update_workflow_outcome(workflow_id, "PENDING")