Test & Improve Queue Semantics (#202)

- Cancelling a workflow removes it from any queue it is in. This way,
cancelled workflows don't block queues.
- Resume sets status to `PENDING` and resets recovery attempts. It also
removes the workflow from any queue it is in. This way, resume can
safely start DLQ'ed, `CANCELLED`, and `ENQUEUED` workflows.
- Lots and lots of tests.

Author: kraftp

dbos/_dbos.py

@@ -800,14 +800,13 @@ def recover_pending_workflows(
     @classmethod
     def cancel_workflow(cls, workflow_id: str) -> None:
         """Cancel a workflow by ID."""
-        _get_dbos_instance()._sys_db.set_workflow_status(
-            workflow_id, WorkflowStatusString.CANCELLED
-        )
+        _get_dbos_instance()._sys_db.cancel_workflow(workflow_id)
 
     @classmethod
-    def resume_workflow(cls, workflow_id: str) -> None:
+    def resume_workflow(cls, workflow_id: str) -> WorkflowHandle[Any]:
         """Resume a workflow by ID."""
-        execute_workflow_by_id(_get_dbos_instance(), workflow_id, False)
+        _get_dbos_instance()._sys_db.resume_workflow(workflow_id)
+        return execute_workflow_by_id(_get_dbos_instance(), workflow_id, False)
 
     @classproperty
     def logger(cls) -> Logger:

dbos/_sys_db.py

@@ -390,20 +390,55 @@ def update_workflow_status(
         if status["workflow_uuid"] in self._temp_txn_wf_ids:
             self._exported_temp_txn_wf_status.add(status["workflow_uuid"])
 
-    def set_workflow_status(
+    def cancel_workflow(
         self,
-        workflow_uuid: str,
-        status: WorkflowStatusString,
+        workflow_id: str,
     ) -> None:
         with self.engine.begin() as c:
-            stmt = (
+            # Remove the workflow from the queues table so it does not block the table
+            c.execute(
+                sa.delete(SystemSchema.workflow_queue).where(
+                    SystemSchema.workflow_queue.c.workflow_uuid == workflow_id
+                )
+            )
+            # Set the workflow's status to CANCELLED
+            c.execute(
                 sa.update(SystemSchema.workflow_status)
-                .where(SystemSchema.workflow_status.c.workflow_uuid == workflow_uuid)
+                .where(SystemSchema.workflow_status.c.workflow_uuid == workflow_id)
                 .values(
-                    status=status,
+                    status=WorkflowStatusString.CANCELLED.value,
+                )
+            )
+
+    def resume_workflow(
+        self,
+        workflow_id: str,
+    ) -> None:
+        with self.engine.begin() as c:
+            # Check the status of the workflow. If it is complete, do nothing.
+            row = c.execute(
+                sa.select(
+                    SystemSchema.workflow_status.c.status,
+                ).where(SystemSchema.workflow_status.c.workflow_uuid == workflow_id)
+            ).fetchone()
+            if (
+                row is None
+                or row[0] == WorkflowStatusString.SUCCESS.value
+                or row[0] == WorkflowStatusString.ERROR.value
+            ):
+                return
+            # Remove the workflow from the queues table so resume can safely be called on an ENQUEUED workflow
+            c.execute(
+                sa.delete(SystemSchema.workflow_queue).where(
+                    SystemSchema.workflow_queue.c.workflow_uuid == workflow_id
                 )
             )
-            c.execute(stmt)
+            # Set the workflow's status to PENDING and clear its recovery attempts.
+            c.execute(
+                sa.update(SystemSchema.workflow_status)
+                .where(SystemSchema.workflow_status.c.workflow_uuid == workflow_id)
+                .values(status=WorkflowStatusString.PENDING.value, recovery_attempts=0)
+            )
 
     def get_workflow_status(
         self, workflow_uuid: str

dbos/_workflow_commands.py

@@ -1,23 +1,14 @@
-import importlib
-import os
-import sys
-from typing import Any, List, Optional, cast
+from typing import List, Optional, cast
 
 import typer
-from rich import print
 
-from dbos import DBOS
-
-from . import _serialization, load_config
-from ._core import execute_workflow_by_id
-from ._dbos_config import ConfigFile, _is_valid_app_name
+from . import _serialization
+from ._dbos_config import ConfigFile
 from ._sys_db import (
     GetWorkflowsInput,
     GetWorkflowsOutput,
     SystemDatabase,
     WorkflowStatuses,
-    WorkflowStatusInternal,
-    WorkflowStatusString,
 )
 
 
@@ -41,7 +32,7 @@ class WorkflowInformation:
     queue_name: Optional[str]
 
 
-def _list_workflows(
+def list_workflows(
     config: ConfigFile,
     li: int,
     user: Optional[str],
@@ -91,17 +82,13 @@ def _list_workflows(
             sys_db.destroy()
 
 
-def _get_workflow(
+def get_workflow(
     config: ConfigFile, uuid: str, request: bool
 ) -> Optional[WorkflowInformation]:
-    sys_db = None
-
     try:
         sys_db = SystemDatabase(config)
-
         info = _get_workflow_info(sys_db, uuid, request)
         return info
-
     except Exception as e:
         typer.echo(f"Error getting workflow: {e}")
         return None
@@ -110,18 +97,13 @@ def _get_workflow(
             sys_db.destroy()
 
 
-def _cancel_workflow(config: ConfigFile, uuid: str) -> None:
-    # config = load_config()
-    sys_db = None
-
+def cancel_workflow(config: ConfigFile, uuid: str) -> None:
     try:
         sys_db = SystemDatabase(config)
-        sys_db.set_workflow_status(uuid, WorkflowStatusString.CANCELLED)
-        return
-
+        sys_db.cancel_workflow(uuid)
     except Exception as e:
         typer.echo(f"Failed to connect to DBOS system database: {e}")
-        return None
+        raise e
     finally:
         if sys_db:
             sys_db.destroy()

dbos/cli/cli.py

@@ -19,7 +19,7 @@
 from .._app_db import ApplicationDatabase
 from .._dbos_config import _is_valid_app_name
 from .._sys_db import SystemDatabase, reset_system_database
-from .._workflow_commands import _cancel_workflow, _get_workflow, _list_workflows
+from .._workflow_commands import cancel_workflow, get_workflow, list_workflows
 from ..cli._github_init import create_template_from_github
 from ._template_init import copy_template, get_project_name, get_templates_directory
 
@@ -282,7 +282,7 @@ def list(
     ] = None,
 ) -> None:
     config = load_config()
-    workflows = _list_workflows(
+    workflows = list_workflows(
         config, limit, user, starttime, endtime, status, request, appversion
     )
     print(jsonpickle.encode(workflows, unpicklable=False))
@@ -301,7 +301,7 @@ def get(
     ] = True,
 ) -> None:
     config = load_config()
-    print(jsonpickle.encode(_get_workflow(config, uuid, request), unpicklable=False))
+    print(jsonpickle.encode(get_workflow(config, uuid, request), unpicklable=False))
 
 
 @workflow.command(
@@ -315,7 +315,7 @@ def cancel(
     ] = None,
 ) -> None:
     config = load_config()
-    _cancel_workflow(config, uuid)
+    cancel_workflow(config, uuid)
     print(f"Workflow {uuid} has been cancelled")
 
 

tests/test_admin_server.py

@@ -151,54 +151,59 @@ def test_admin_diff_port(cleanup_test_databases: None) -> None:
 
 
 def test_admin_workflow_resume(dbos: DBOS, config: ConfigFile) -> None:
+    counter: int = 0
 
     @DBOS.workflow()
     def simple_workflow() -> None:
-        print("Executed Simple workflow")
-        return
+        nonlocal counter
+        counter += 1
 
-    # run the workflow
+    # Run the workflow and flush its results
     simple_workflow()
-    time.sleep(1)
+    assert counter == 1
+    dbos._sys_db.wait_for_buffer_flush()
 
-    # get the workflow list
-    output = _workflow_commands._list_workflows(
+    # Verify the workflow has succeeded
+    output = _workflow_commands.list_workflows(
         config, 10, None, None, None, None, False, None
     )
     assert len(output) == 1, f"Expected list length to be 1, but got {len(output)}"
-
     assert output[0] != None, "Expected output to be not None"
-
     wfUuid = output[0].workflowUUID
-
-    info = _workflow_commands._get_workflow(config, wfUuid, True)
+    info = _workflow_commands.get_workflow(config, wfUuid, True)
     assert info is not None, "Expected output to be not None"
-
     assert info.status == "SUCCESS", f"Expected status to be SUCCESS"
 
+    # Cancel the workflow. Verify it was cancelled
     response = requests.post(
         f"http://localhost:3001/workflows/{wfUuid}/cancel", json=[], timeout=5
     )
     assert response.status_code == 204
+    info = _workflow_commands.get_workflow(config, wfUuid, True)
+    assert info is not None
+    assert info.status == "CANCELLED", f"Expected status to be CANCELLED"
 
-    info = _workflow_commands._get_workflow(config, wfUuid, True)
-    if info is not None:
-        assert info.status == "CANCELLED", f"Expected status to be CANCELLED"
-    else:
-        assert False, "Expected info to be not None"
-
+    # Resume the workflow. Verify that it succeeds again.
     response = requests.post(
         f"http://localhost:3001/workflows/{wfUuid}/resume", json=[], timeout=5
     )
     assert response.status_code == 204
+    dbos._sys_db.wait_for_buffer_flush()
+    assert counter == 2
+    info = _workflow_commands.get_workflow(config, wfUuid, True)
+    assert info is not None
+    assert info.status == "SUCCESS", f"Expected status to be SUCCESS"
 
-    time.sleep(1)
-
-    info = _workflow_commands._get_workflow(config, wfUuid, True)
-    if info is not None:
-        assert info.status == "SUCCESS", f"Expected status to be SUCCESS"
-    else:
-        assert False, "Expected info to be not None"
+    # Resume the workflow. Verify it does not run and status remains SUCCESS
+    response = requests.post(
+        f"http://localhost:3001/workflows/{wfUuid}/resume", json=[], timeout=5
+    )
+    assert response.status_code == 204
+    dbos._sys_db.wait_for_buffer_flush()
+    info = _workflow_commands.get_workflow(config, wfUuid, True)
+    assert info is not None
+    assert info.status == "SUCCESS", f"Expected status to be SUCCESS"
+    assert counter == 2
 
 
 def test_admin_workflow_restart(dbos: DBOS, config: ConfigFile) -> None:
@@ -213,7 +218,7 @@ def simple_workflow() -> None:
     time.sleep(1)
 
     # get the workflow list
-    output = _workflow_commands._list_workflows(
+    output = _workflow_commands.list_workflows(
         config, 10, None, None, None, None, False, None
     )
     assert len(output) == 1, f"Expected list length to be 1, but got {len(output)}"
@@ -222,7 +227,7 @@ def simple_workflow() -> None:
 
     wfUuid = output[0].workflowUUID
 
-    info = _workflow_commands._get_workflow(config, wfUuid, True)
+    info = _workflow_commands.get_workflow(config, wfUuid, True)
     assert info is not None, "Expected output to be not None"
 
     assert info.status == "SUCCESS", f"Expected status to be SUCCESS"
@@ -232,7 +237,7 @@ def simple_workflow() -> None:
     )
     assert response.status_code == 204
 
-    info = _workflow_commands._get_workflow(config, wfUuid, True)
+    info = _workflow_commands.get_workflow(config, wfUuid, True)
     if info is not None:
         assert info.status == "CANCELLED", f"Expected status to be CANCELLED"
     else:
@@ -245,13 +250,13 @@ def simple_workflow() -> None:
 
     time.sleep(1)
 
-    info = _workflow_commands._get_workflow(config, wfUuid, True)
+    info = _workflow_commands.get_workflow(config, wfUuid, True)
     if info is not None:
         assert info.status == "CANCELLED", f"Expected status to be CANCELLED"
     else:
         assert False, "Expected info to be not None"
 
-    output = _workflow_commands._list_workflows(
+    output = _workflow_commands.list_workflows(
         config, 10, None, None, None, None, False, None
     )
     assert len(output) == 2, f"Expected list length to be 2, but got {len(output)}"
@@ -261,7 +266,7 @@ def simple_workflow() -> None:
     else:
         new_wfUuid = output[0].workflowUUID
 
-    info = _workflow_commands._get_workflow(config, new_wfUuid, True)
+    info = _workflow_commands.get_workflow(config, new_wfUuid, True)
     if info is not None:
         assert info.status == "SUCCESS", f"Expected status to be SUCCESS"
     else: