Support concurrent recv and get_event (#345)

This PR fixes a race condition when multiple workflow instances are
waiting on `recv` or multiple callers are waiting on `get_event`.

The solution is to maintain a thread safe map.
- For `recv`, only one workflow instance should be waiting, because
`recv` consumes messages. If a workflow is already waiting for `recv`,
directly raise a `DBOSWorkflowConflictIDError` and wait for the existing
workflow to finish.
- For `get_event`, multiple callers can wait on the same event. In this
case, we maintain a reference counter and only delete the condition
variable when nobody is waiting.
- Add unit tests for concurrent recv and get_event.

Author: qianl15

dbos/_sys_db.py

@@ -222,6 +222,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 +289,8 @@ def __init__(
         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
@@ -1288,7 +1329,12 @@ def recv(
         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,23 +1427,23 @@ def _notification_listener(self) -> None:
                             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()
                                 dbos_logger.debug(
                                     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 +1581,13 @@ def get_event(
 
         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]

tests/test_concurrency.py

@@ -2,7 +2,7 @@
 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_txn_thread(id: str) -> str:
 
     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