dbos 1.1.0a2__tar.gz → 1.1.0a4__tar.gz

{dbos-1.1.0a2 → dbos-1.1.0a4}/PKG-INFO

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

{dbos-1.1.0a2 → dbos-1.1.0a4}/dbos/_context.py

@@ -392,6 +392,7 @@ class SetWorkflowTimeout:
             else None
         )
         self.saved_workflow_timeout: Optional[int] = None
+        self.saved_workflow_deadline_epoch_ms: Optional[int] = None
 
     def __enter__(self) -> SetWorkflowTimeout:
         # Code to create a basic context
@@ -402,6 +403,8 @@ class SetWorkflowTimeout:
         ctx = assert_current_dbos_context()
         self.saved_workflow_timeout = ctx.workflow_timeout_ms
         ctx.workflow_timeout_ms = self.workflow_timeout_ms
+        self.saved_workflow_deadline_epoch_ms = ctx.workflow_deadline_epoch_ms
+        ctx.workflow_deadline_epoch_ms = None
         return self
 
     def __exit__(
@@ -411,6 +414,9 @@ class SetWorkflowTimeout:
         traceback: Optional[TracebackType],
     ) -> Literal[False]:
         assert_current_dbos_context().workflow_timeout_ms = self.saved_workflow_timeout
+        assert_current_dbos_context().workflow_deadline_epoch_ms = (
+            self.saved_workflow_deadline_epoch_ms
+        )
         # Code to clean up the basic context if we created it
         if self.created_ctx:
             _clear_local_dbos_context()

{dbos-1.1.0a2 → dbos-1.1.0a4}/dbos/_error.py

@@ -62,6 +62,7 @@ class DBOSErrorCode(Enum):
     WorkflowCancelled = 10
     UnexpectedStep = 11
     QueueDeduplicated = 12
+    AwaitedWorkflowCancelled = 13
     ConflictingRegistrationError = 25
 
 
@@ -206,6 +207,19 @@ class DBOSQueueDeduplicatedError(DBOSException):
         )
 
 
+class DBOSAwaitedWorkflowCancelledError(DBOSException):
+    def __init__(self, workflow_id: str):
+        self.workflow_id = workflow_id
+        super().__init__(
+            f"Awaited workflow {workflow_id} was cancelled",
+            dbos_error_code=DBOSErrorCode.AwaitedWorkflowCancelled.value,
+        )
+
+    def __reduce__(self) -> Any:
+        # Tell jsonpickle how to reconstruct this object
+        return (self.__class__, (self.workflow_id,))
+
+
 #######################################
 ## BaseException
 #######################################

{dbos-1.1.0a2 → dbos-1.1.0a4}/dbos/_sys_db.py

@@ -32,6 +32,7 @@ from dbos._utils import INTERNAL_QUEUE_NAME
 from . import _serialization
 from ._context import get_local_dbos_context
 from ._error import (
+    DBOSAwaitedWorkflowCancelledError,
     DBOSConflictingWorkflowError,
     DBOSDeadLetterQueueError,
     DBOSNonExistentWorkflowError,
@@ -96,6 +97,10 @@ class WorkflowStatus:
     executor_id: Optional[str]
     # The application version on which this workflow was started
     app_version: Optional[str]
+    # The start-to-close timeout of the workflow in ms
+    workflow_timeout_ms: Optional[int]
+    # The deadline of a workflow, computed by adding its timeout to its start time.
+    workflow_deadline_epoch_ms: Optional[int]
 
     # INTERNAL FIELDS
 
@@ -222,6 +227,47 @@ class StepInfo(TypedDict):
 _dbos_null_topic = "__null__topic__"
 
 
+class ConditionCount(TypedDict):
+    condition: threading.Condition
+    count: int
+
+
+class ThreadSafeConditionDict:
+    def __init__(self) -> None:
+        self._dict: Dict[str, ConditionCount] = {}
+        self._lock = threading.Lock()
+
+    def get(self, key: str) -> Optional[threading.Condition]:
+        with self._lock:
+            if key not in self._dict:
+                # Key does not exist, return None
+                return None
+            return self._dict[key]["condition"]
+
+    def set(
+        self, key: str, value: threading.Condition
+    ) -> tuple[bool, threading.Condition]:
+        with self._lock:
+            if key in self._dict:
+                # Key already exists, do not overwrite. Increment the wait count.
+                cc = self._dict[key]
+                cc["count"] += 1
+                return False, cc["condition"]
+            self._dict[key] = ConditionCount(condition=value, count=1)
+            return True, value
+
+    def pop(self, key: str) -> None:
+        with self._lock:
+            if key in self._dict:
+                cc = self._dict[key]
+                cc["count"] -= 1
+                if cc["count"] == 0:
+                    # No more threads waiting on this condition, remove it
+                    del self._dict[key]
+            else:
+                dbos_logger.warning(f"Key {key} not found in condition dictionary.")
+
+
 class SystemDatabase:
 
     def __init__(
@@ -248,8 +294,8 @@ class SystemDatabase:
         self._engine_kwargs = engine_kwargs
 
         self.notification_conn: Optional[psycopg.connection.Connection] = None
-        self.notifications_map: Dict[str, threading.Condition] = {}
-        self.workflow_events_map: Dict[str, threading.Condition] = {}
+        self.notifications_map = ThreadSafeConditionDict()
+        self.workflow_events_map = ThreadSafeConditionDict()
 
         # Now we can run background processes
         self._run_background_processes = True
@@ -720,9 +766,9 @@ class SystemDatabase:
                         error = row[2]
                         raise _serialization.deserialize_exception(error)
                     elif status == WorkflowStatusString.CANCELLED.value:
-                        # Raise a normal exception here, not the cancellation exception
+                        # Raise AwaitedWorkflowCancelledError here, not the cancellation exception
                         # because the awaiting workflow is not being cancelled.
-                        raise Exception(f"Awaited workflow {workflow_id} was cancelled")
+                        raise DBOSAwaitedWorkflowCancelledError(workflow_id)
                 else:
                     pass  # CB: I guess we're assuming the WF will show up eventually.
             time.sleep(1)
@@ -796,6 +842,8 @@ class SystemDatabase:
             SystemSchema.workflow_inputs.c.inputs,
             SystemSchema.workflow_status.c.output,
             SystemSchema.workflow_status.c.error,
+            SystemSchema.workflow_status.c.workflow_deadline_epoch_ms,
+            SystemSchema.workflow_status.c.workflow_timeout_ms,
         ).join(
             SystemSchema.workflow_inputs,
             SystemSchema.workflow_status.c.workflow_uuid
@@ -877,6 +925,8 @@ class SystemDatabase:
             info.input = inputs
             info.output = output
             info.error = exception
+            info.workflow_deadline_epoch_ms = row[18]
+            info.workflow_timeout_ms = row[19]
 
             infos.append(info)
         return infos
@@ -906,6 +956,8 @@ class SystemDatabase:
             SystemSchema.workflow_inputs.c.inputs,
             SystemSchema.workflow_status.c.output,
             SystemSchema.workflow_status.c.error,
+            SystemSchema.workflow_status.c.workflow_deadline_epoch_ms,
+            SystemSchema.workflow_status.c.workflow_timeout_ms,
         ).select_from(
             SystemSchema.workflow_queue.join(
                 SystemSchema.workflow_status,
@@ -983,6 +1035,8 @@ class SystemDatabase:
             info.input = inputs
             info.output = output
             info.error = exception
+            info.workflow_deadline_epoch_ms = row[18]
+            info.workflow_timeout_ms = row[19]
 
             infos.append(info)
 
@@ -1288,7 +1342,12 @@ class SystemDatabase:
         condition = threading.Condition()
         # Must acquire first before adding to the map. Otherwise, the notification listener may notify it before the condition is acquired and waited.
         condition.acquire()
-        self.notifications_map[payload] = condition
+        success, _ = self.notifications_map.set(payload, condition)
+        if not success:
+            # This should not happen, but if it does, it means the workflow is executed concurrently.
+            condition.release()
+            self.notifications_map.pop(payload)
+            raise DBOSWorkflowConflictIDError(workflow_uuid)
 
         # Check if the key is already in the database. If not, wait for the notification.
         init_recv: Sequence[Any]
@@ -1381,11 +1440,11 @@ class SystemDatabase:
                             f"Received notification on channel: {channel}, payload: {notify.payload}"
                         )
                         if channel == "dbos_notifications_channel":
-                            if (
-                                notify.payload
-                                and notify.payload in self.notifications_map
-                            ):
-                                condition = self.notifications_map[notify.payload]
+                            if notify.payload:
+                                condition = self.notifications_map.get(notify.payload)
+                                if condition is None:
+                                    # No condition found for this payload
+                                    continue
                                 condition.acquire()
                                 condition.notify_all()
                                 condition.release()
@@ -1393,11 +1452,11 @@ class SystemDatabase:
                                     f"Signaled notifications condition for {notify.payload}"
                                 )
                         elif channel == "dbos_workflow_events_channel":
-                            if (
-                                notify.payload
-                                and notify.payload in self.workflow_events_map
-                            ):
-                                condition = self.workflow_events_map[notify.payload]
+                            if notify.payload:
+                                condition = self.workflow_events_map.get(notify.payload)
+                                if condition is None:
+                                    # No condition found for this payload
+                                    continue
                                 condition.acquire()
                                 condition.notify_all()
                                 condition.release()
@@ -1535,8 +1594,13 @@ class SystemDatabase:
 
         payload = f"{target_uuid}::{key}"
         condition = threading.Condition()
-        self.workflow_events_map[payload] = condition
         condition.acquire()
+        success, existing_condition = self.workflow_events_map.set(payload, condition)
+        if not success:
+            # Wait on the existing condition
+            condition.release()
+            condition = existing_condition
+            condition.acquire()
 
         # Check if the key is already in the database. If not, wait for the notification.
         init_recv: Sequence[Any]
@@ -1776,8 +1840,13 @@ class SystemDatabase:
                         # If a timeout is set, set the deadline on dequeue
                         workflow_deadline_epoch_ms=sa.case(
                             (
-                                SystemSchema.workflow_status.c.workflow_timeout_ms.isnot(
-                                    None
+                                sa.and_(
+                                    SystemSchema.workflow_status.c.workflow_timeout_ms.isnot(
+                                        None
+                                    ),
+                                    SystemSchema.workflow_status.c.workflow_deadline_epoch_ms.is_(
+                                        None
+                                    ),
                                 ),
                                 sa.func.extract("epoch", sa.func.now()) * 1000
                                 + SystemSchema.workflow_status.c.workflow_timeout_ms,

{dbos-1.1.0a2 → dbos-1.1.0a4}/pyproject.toml

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

{dbos-1.1.0a2 → dbos-1.1.0a4}/tests/test_concurrency.py

@@ -2,7 +2,7 @@ import threading
 import time
 import uuid
 from concurrent.futures import Future, ThreadPoolExecutor
-from typing import Tuple
+from typing import Tuple, cast
 
 from sqlalchemy import text
 
@@ -108,3 +108,67 @@ def test_concurrent_conflict_uuid(dbos: DBOS) -> None:
 
     assert future1.result() == wfuuid
     assert future2.result() == wfuuid
+
+
+def test_concurrent_recv(dbos: DBOS) -> None:
+    condition = threading.Condition()
+    counter = 0
+
+    @DBOS.workflow()
+    def test_workflow(topic: str) -> str:
+        nonlocal counter
+        condition.acquire()
+        counter += 1
+        if counter % 2 == 1:
+            # Wait for the other one to notify
+            condition.wait()
+        else:
+            # Notify the other one
+            condition.notify()
+        condition.release()
+        m = cast(str, DBOS.recv(topic, 5))
+        return m
+
+    def test_thread(id: str, topic: str) -> str:
+        with SetWorkflowID(id):
+            return test_workflow(topic)
+
+    wfuuid = str(uuid.uuid4())
+    topic = "test_topic"
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        future1 = executor.submit(test_thread, wfuuid, topic)
+        future2 = executor.submit(test_thread, wfuuid, topic)
+
+        expected_message = "test message"
+        DBOS.send(wfuuid, expected_message, topic)
+        # Both should return the same message
+        assert future1.result() == future2.result()
+        assert future1.result() == expected_message
+        # Make sure the notification map is empty
+        assert not dbos._sys_db.notifications_map._dict
+
+
+def test_concurrent_getevent(dbos: DBOS) -> None:
+    @DBOS.workflow()
+    def test_workflow(event_name: str, value: str) -> str:
+        DBOS.set_event(event_name, value)
+        return value
+
+    def test_thread(id: str, event_name: str) -> str:
+        return cast(str, DBOS.get_event(id, event_name, 5))
+
+    wfuuid = str(uuid.uuid4())
+    event_name = "test_event"
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        future1 = executor.submit(test_thread, wfuuid, event_name)
+        future2 = executor.submit(test_thread, wfuuid, event_name)
+
+        expected_message = "test message"
+        with SetWorkflowID(wfuuid):
+            test_workflow(event_name, expected_message)
+
+        # Both should return the same message
+        assert future1.result() == future2.result()
+        assert future1.result() == expected_message
+        # Make sure the event map is empty
+        assert not dbos._sys_db.workflow_events_map._dict

{dbos-1.1.0a2 → dbos-1.1.0a4}/tests/test_dbos.py

@@ -24,6 +24,7 @@ from dbos import (
 # Private API because this is a test
 from dbos._context import assert_current_dbos_context, get_local_dbos_context
 from dbos._error import (
+    DBOSAwaitedWorkflowCancelledError,
     DBOSConflictingRegistrationError,
     DBOSMaxStepRetriesExceeded,
     DBOSWorkflowCancelledError,
@@ -1507,7 +1508,14 @@ def test_workflow_timeout(dbos: DBOS) -> None:
             with SetWorkflowID(wfid):
                 blocked_workflow()
         assert assert_current_dbos_context().workflow_deadline_epoch_ms is None
+        start_time = time.time() * 1000
         handle = DBOS.start_workflow(blocked_workflow)
+        status = handle.get_status()
+        assert status.workflow_timeout_ms == 100
+        assert (
+            status.workflow_deadline_epoch_ms is not None
+            and status.workflow_deadline_epoch_ms > start_time
+        )
         with pytest.raises(DBOSWorkflowCancelledError):
             handle.get_result()
 
@@ -1555,13 +1563,11 @@ def test_workflow_timeout(dbos: DBOS) -> None:
         with pytest.raises(DBOSWorkflowCancelledError):
             parent_workflow()
 
-    with pytest.raises(Exception) as exc_info:
+    with pytest.raises(DBOSAwaitedWorkflowCancelledError):
         DBOS.retrieve_workflow(start_child).get_result()
-    assert "was cancelled" in str(exc_info.value)
 
-    with pytest.raises(Exception) as exc_info:
+    with pytest.raises(DBOSAwaitedWorkflowCancelledError):
         DBOS.retrieve_workflow(direct_child).get_result()
-    assert "was cancelled" in str(exc_info.value)
 
     # Verify the context variables are set correctly
     with SetWorkflowTimeout(1.0):

{dbos-1.1.0a2 → dbos-1.1.0a4}/tests/test_failures.py

@@ -9,6 +9,7 @@ from sqlalchemy.exc import InvalidRequestError, OperationalError
 
 from dbos import DBOS, Queue, SetWorkflowID
 from dbos._error import (
+    DBOSAwaitedWorkflowCancelledError,
     DBOSDeadLetterQueueError,
     DBOSMaxStepRetriesExceeded,
     DBOSNotAuthorizedError,
@@ -461,6 +462,11 @@ def test_error_serialization() -> None:
     d = deserialize_exception(serialize_exception(e))
     assert isinstance(d, DBOSQueueDeduplicatedError)
     assert str(d) == str(e)
+    # AwaitedWorkflowCancelledError
+    e = DBOSAwaitedWorkflowCancelledError("id")
+    d = deserialize_exception(serialize_exception(e))
+    assert isinstance(d, DBOSAwaitedWorkflowCancelledError)
+    assert str(d) == str(e)
 
     # Test safe_deserialize
     class BadException(Exception):

{dbos-1.1.0a2 → dbos-1.1.0a4}/tests/test_queue.py

@@ -26,6 +26,7 @@ from dbos import (
 )
 from dbos._context import assert_current_dbos_context
 from dbos._dbos import WorkflowHandleAsync
+from dbos._error import DBOSAwaitedWorkflowCancelledError, DBOSWorkflowCancelledError
 from dbos._schemas.system_database import SystemSchema
 from dbos._sys_db import WorkflowStatusString
 from dbos._utils import GlobalParams
@@ -853,9 +854,8 @@ def test_cancelling_queued_workflows(dbos: DBOS) -> None:
 
     # Complete the blocked workflow
     blocking_event.set()
-    with pytest.raises(Exception) as exc_info:
+    with pytest.raises(DBOSAwaitedWorkflowCancelledError):
         blocked_handle.get_result()
-    assert "was cancelled" in str(exc_info.value)
 
     # Verify all queue entries eventually get cleaned up.
     assert queue_entries_are_cleaned_up(dbos)
@@ -891,9 +891,8 @@ def test_timeout_queue(dbos: DBOS) -> None:
 
     # Verify the blocked workflows are cancelled
     for handle in handles:
-        with pytest.raises(Exception) as exc_info:
+        with pytest.raises(DBOSAwaitedWorkflowCancelledError):
             handle.get_result()
-        assert "was cancelled" in str(exc_info.value)
 
     # Verify the normal workflow succeeds
     normal_handle.get_result()
@@ -911,17 +910,14 @@ def test_timeout_queue(dbos: DBOS) -> None:
 
     with SetWorkflowTimeout(1.0):
         handle = queue.enqueue(parent_workflow)
-    with pytest.raises(Exception) as exc_info:
+    with pytest.raises(DBOSAwaitedWorkflowCancelledError):
         handle.get_result()
-    assert "was cancelled" in str(exc_info.value)
 
-    with pytest.raises(Exception) as exc_info:
+    with pytest.raises(DBOSAwaitedWorkflowCancelledError):
         DBOS.retrieve_workflow(child_id).get_result()
-    assert "was cancelled" in str(exc_info.value)
 
     # Verify if a parent called with a timeout enqueues a blocked child
     # then exits the deadline propagates and the child is cancelled.
-    child_id = str(uuid.uuid4())
     queue = Queue("regular_queue")
 
     @DBOS.workflow()
@@ -931,9 +927,41 @@ def test_timeout_queue(dbos: DBOS) -> None:
 
     with SetWorkflowTimeout(1.0):
         child_id = exiting_parent_workflow()
-    with pytest.raises(Exception) as exc_info:
+    with pytest.raises(DBOSAwaitedWorkflowCancelledError):
+        DBOS.retrieve_workflow(child_id).get_result()
+
+    # Verify if a parent called with a timeout enqueues a child that
+    # never starts because the queue is blocked, the deadline propagates
+    # and both parent and child are cancelled.
+    child_id = str(uuid.uuid4())
+    queue = Queue("stuck_queue", concurrency=1)
+
+    start_event = threading.Event()
+    blocking_event = threading.Event()
+
+    @DBOS.workflow()
+    def stuck_workflow() -> None:
+        start_event.set()
+        blocking_event.wait()
+
+    stuck_handle = queue.enqueue(stuck_workflow)
+    start_event.wait()
+
+    @DBOS.workflow()
+    def blocked_parent_workflow() -> None:
+        with SetWorkflowID(child_id):
+            queue.enqueue(blocking_workflow)
+        while True:
+            DBOS.sleep(0.1)
+
+    with SetWorkflowTimeout(1.0):
+        handle = DBOS.start_workflow(blocked_parent_workflow)
+    with pytest.raises(DBOSWorkflowCancelledError):
+        handle.get_result()
+    with pytest.raises(DBOSAwaitedWorkflowCancelledError):
         DBOS.retrieve_workflow(child_id).get_result()
-    assert "was cancelled" in str(exc_info.value)
+    blocking_event.set()
+    stuck_handle.get_result()
 
     # Verify all queue entries eventually get cleaned up.
     assert queue_entries_are_cleaned_up(dbos)
@@ -1341,3 +1369,80 @@ def test_worker_concurrency_across_versions(dbos: DBOS, client: DBOSClient) -> N
     # Change the version, verify the other version complets
     GlobalParams.app_version = other_version
     assert other_version_handle.get_result()
+
+
+def test_timeout_queue_recovery(dbos: DBOS) -> None:
+    queue = Queue("test_queue")
+    evt = threading.Event()
+
+    @DBOS.workflow()
+    def blocking_workflow() -> None:
+        evt.set()
+        while True:
+            DBOS.sleep(0.1)
+
+    timeout = 3.0
+    enqueue_time = time.time()
+    with SetWorkflowTimeout(timeout):
+        original_handle = queue.enqueue(blocking_workflow)
+
+    # Verify the workflow's timeout is properly configured
+    evt.wait()
+    original_status = original_handle.get_status()
+    assert original_status.workflow_timeout_ms == timeout * 1000
+    assert (
+        original_status.workflow_deadline_epoch_ms is not None
+        and original_status.workflow_deadline_epoch_ms > enqueue_time * 1000
+    )
+
+    # Recover the workflow. Verify its deadline remains the same
+    evt.clear()
+    handles = DBOS._recover_pending_workflows()
+    assert len(handles) == 1
+    evt.wait()
+    recovered_handle = handles[0]
+    recovered_status = recovered_handle.get_status()
+    assert recovered_status.workflow_timeout_ms == timeout * 1000
+    assert (
+        recovered_status.workflow_deadline_epoch_ms
+        == original_status.workflow_deadline_epoch_ms
+    )
+
+    with pytest.raises(DBOSAwaitedWorkflowCancelledError):
+        original_handle.get_result()
+
+    with pytest.raises(DBOSAwaitedWorkflowCancelledError):
+        recovered_handle.get_result()
+
+
+def test_unsetting_timeout(dbos: DBOS) -> None:
+
+    queue = Queue("test_queue")
+
+    @DBOS.workflow()
+    def child() -> str:
+        for _ in range(5):
+            DBOS.sleep(1)
+        return DBOS.workflow_id
+
+    @DBOS.workflow()
+    def parent(child_one: str, child_two: str) -> None:
+        with SetWorkflowID(child_two):
+            with SetWorkflowTimeout(None):
+                queue.enqueue(child)
+
+        with SetWorkflowID(child_one):
+            queue.enqueue(child)
+
+    child_one, child_two = str(uuid.uuid4()), str(uuid.uuid4())
+    with SetWorkflowTimeout(1.0):
+        queue.enqueue(parent, child_one, child_two).get_result()
+
+    # Verify child one, which has a propagated timeout, is cancelled
+    handle: WorkflowHandle[str] = DBOS.retrieve_workflow(child_one)
+    with pytest.raises(DBOSAwaitedWorkflowCancelledError):
+        handle.get_result()
+
+    # Verify child two, which doesn't have a timeout, succeeds
+    handle = DBOS.retrieve_workflow(child_two)
+    assert handle.get_result() == child_two

{dbos-1.1.0a2 → dbos-1.1.0a4}/tests/test_workflow_introspection.py

@@ -41,6 +41,8 @@ def test_list_workflow(dbos: DBOS) -> None:
     assert output.app_version == GlobalParams.app_version
     assert output.app_id == ""
     assert output.recovery_attempts == 1
+    assert output.workflow_timeout_ms is None
+    assert output.workflow_deadline_epoch_ms is None
 
     # Test searching by status
     outputs = DBOS.list_workflows(status="PENDING")
@@ -222,6 +224,8 @@ def test_queued_workflows(dbos: DBOS) -> None:
         assert workflow.created_at is not None and workflow.created_at > 0
         assert workflow.updated_at is not None and workflow.updated_at > 0
         assert workflow.recovery_attempts == 1
+        assert workflow.workflow_timeout_ms is None
+        assert workflow.workflow_deadline_epoch_ms is None
 
     # Test sort_desc inverts the order
     workflows = DBOS.list_queued_workflows(sort_desc=True)