dbos 0.23.0a9__tar.gz → 0.23.0a10__tar.gz

{dbos-0.23.0a9 → dbos-0.23.0a10}/PKG-INFO

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

{dbos-0.23.0a9 → dbos-0.23.0a10}/dbos/_recovery.py

@@ -27,8 +27,9 @@ def startup_recovery_thread(
                     pending_workflow.queue_name
                     and pending_workflow.queue_name != "_dbos_internal_queue"
                 ):
-                    dbos._sys_db.clear_queue_assignment(pending_workflow.workflow_uuid)
-                    continue
+                    cleared = dbos._sys_db.clear_queue_assignment(pending_workflow.workflow_uuid)
+                    if cleared:
+                        continue
                 execute_workflow_by_id(dbos, pending_workflow.workflow_uuid)
                 pending_workflows.remove(pending_workflow)
         except DBOSWorkflowFunctionNotFoundError:
@@ -56,10 +57,15 @@ def recover_pending_workflows(
                 and pending_workflow.queue_name != "_dbos_internal_queue"
             ):
                 try:
-                    dbos._sys_db.clear_queue_assignment(pending_workflow.workflow_uuid)
-                    workflow_handles.append(
-                        dbos.retrieve_workflow(pending_workflow.workflow_uuid)
-                    )
+                    cleared = dbos._sys_db.clear_queue_assignment(pending_workflow.workflow_uuid)
+                    if cleared:
+                        workflow_handles.append(
+                            dbos.retrieve_workflow(pending_workflow.workflow_uuid)
+                        )
+                    else:
+                        workflow_handles.append(
+                            execute_workflow_by_id(dbos, pending_workflow.workflow_uuid)
+                        )
                 except Exception as e:
                     dbos.logger.error(e)
             else:

{dbos-0.23.0a9 → dbos-0.23.0a10}/dbos/_sys_db.py

@@ -1460,21 +1460,35 @@ class SystemDatabase:
                     .values(completed_at_epoch_ms=int(time.time() * 1000))
                 )
 
-    def clear_queue_assignment(self, workflow_id: str) -> None:
+
+    def clear_queue_assignment(self, workflow_id: str) -> bool:
         if self._debug_mode:
             raise Exception("called clear_queue_assignment in debug mode")
-        with self.engine.begin() as c:
-            c.execute(
-                sa.update(SystemSchema.workflow_queue)
-                .where(SystemSchema.workflow_queue.c.workflow_uuid == workflow_id)
-                .values(executor_id=None, started_at_epoch_ms=None)
-            )
-            c.execute(
-                sa.update(SystemSchema.workflow_status)
-                .where(SystemSchema.workflow_status.c.workflow_uuid == workflow_id)
-                .values(executor_id=None, status=WorkflowStatusString.ENQUEUED.value)
-            )
 
+        with self.engine.connect() as conn:
+            with conn.begin() as transaction:
+                res = conn.execute(
+                    sa.update(SystemSchema.workflow_queue)
+                    .where(SystemSchema.workflow_queue.c.workflow_uuid == workflow_id)
+                    .values(executor_id=None, started_at_epoch_ms=None)
+                )
+
+                # If no rows were affected, the workflow is not anymore in the queue
+                if res.rowcount == 0:
+                    transaction.rollback()
+                    return False
+
+                res = conn.execute(
+                    sa.update(SystemSchema.workflow_status)
+                    .where(SystemSchema.workflow_status.c.workflow_uuid == workflow_id)
+                    .values(executor_id=None, status=WorkflowStatusString.ENQUEUED.value)
+                )
+                if res.rowcount == 0:
+                    # This should never happen
+                    raise Exception(
+                        f"UNREACHABLE: Workflow {workflow_id} is found in the workflow_queue table but not found in the workflow_status table"
+                    )
+                return True
 
 def reset_system_database(config: ConfigFile) -> None:
     sysdb_name = (

{dbos-0.23.0a9 → dbos-0.23.0a10}/pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
 ]
 requires-python = ">=3.9"
 readme = "README.md"
-version = "0.23.0a9"
+version = "0.23.0a10"
 
 [project.license]
 text = "MIT"

{dbos-0.23.0a9 → dbos-0.23.0a10}/tests/test_queue.py

@@ -19,7 +19,7 @@ from dbos import (
     WorkflowHandle,
 )
 from dbos._schemas.system_database import SystemSchema
-from dbos._sys_db import WorkflowStatusString
+from dbos._sys_db import WorkflowStatusString, _buffer_flush_interval_secs
 from tests.conftest import default_config, queue_entries_are_cleaned_up
 
 
@@ -900,6 +900,40 @@ def test_resuming_queued_workflows(dbos: DBOS) -> None:
     assert queue_entries_are_cleaned_up(dbos)
 
 
+# Test a race condition between removing a task from the queue and flushing the status buffer
+def test_resuming_already_completed_queue_workflow(dbos: DBOS) -> None:
+    dbos._sys_db._run_background_processes = False # Disable buffer flush
+
+    start_event = threading.Event()
+    counter = 0
+    @DBOS.workflow()
+    def test_step() -> None:
+        start_event.set()
+        nonlocal counter
+        counter += 1
+
+    queue = Queue("test_queue")
+    handle = queue.enqueue(test_step)
+    start_event.wait()
+    start_event.clear()
+    time.sleep(_buffer_flush_interval_secs)
+    assert handle.get_status().status == WorkflowStatusString.PENDING.value # Not flushed
+    assert counter == 1 # But, really, it's completed
+    dbos._sys_db._workflow_status_buffer = {} # Clear buffer (simulates a process restart)
+
+    # Recovery picks up on the workflow and recovers it
+    recovered_ids = DBOS.recover_pending_workflows()
+    assert len(recovered_ids) == 1
+    assert recovered_ids[0].get_workflow_id() == handle.get_workflow_id()
+    start_event.wait()
+    assert counter == 2 # The workflow ran again
+    time.sleep(_buffer_flush_interval_secs) # This is actually to wait that _get_wf_invoke_func buffers the status
+    dbos._sys_db._flush_workflow_status_buffer() # Manually flush
+    assert handle.get_status().status == WorkflowStatusString.SUCCESS.value # Is recovered
+    assert handle.get_status().executor_id == "local"
+    assert handle.get_status().recovery_attempts == 2
+
+
 def test_dlq_enqueued_workflows(dbos: DBOS) -> None:
     start_event = threading.Event()
     blocking_event = threading.Event()