dbos 1.1.0a2__py3-none-any.whl → 1.1.0a3__py3-none-any.whl

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 @@ 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
@@ -1288,7 +1329,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 +1427,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 +1439,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 +1581,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]

{dbos-1.1.0a2.dist-info → dbos-1.1.0a3.dist-info}/METADATA

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

{dbos-1.1.0a2.dist-info → dbos-1.1.0a3.dist-info}/RECORD

@@ -1,7 +1,7 @@
-dbos-1.1.0a2.dist-info/METADATA,sha256=X9RcRNberl8VXnm_V76fVezhh20QcbKpoFto07FXjhs,13267
-dbos-1.1.0a2.dist-info/WHEEL,sha256=tSfRZzRHthuv7vxpI4aehrdN9scLjk-dCJkPLzkHxGg,90
-dbos-1.1.0a2.dist-info/entry_points.txt,sha256=_QOQ3tVfEjtjBlr1jS4sHqHya9lI2aIEIWkz8dqYp14,58
-dbos-1.1.0a2.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
+dbos-1.1.0a3.dist-info/METADATA,sha256=wwmIbv3PGphti-I6d0es4O3hZMHu7m7hXZbMfrko2WM,13267
+dbos-1.1.0a3.dist-info/WHEEL,sha256=tSfRZzRHthuv7vxpI4aehrdN9scLjk-dCJkPLzkHxGg,90
+dbos-1.1.0a3.dist-info/entry_points.txt,sha256=_QOQ3tVfEjtjBlr1jS4sHqHya9lI2aIEIWkz8dqYp14,58
+dbos-1.1.0a3.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
 dbos/__init__.py,sha256=NssPCubaBxdiKarOWa-wViz1hdJSkmBGcpLX_gQ4NeA,891
 dbos/__main__.py,sha256=G7Exn-MhGrVJVDbgNlpzhfh8WMX_72t3_oJaFT9Lmt8,653
 dbos/_admin_server.py,sha256=A_28_nJ1nBBYDmCxtklJR9O2v14JRMtD1rAo_D4y8Kc,9764
@@ -47,7 +47,7 @@ dbos/_schemas/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dbos/_schemas/application_database.py,sha256=SypAS9l9EsaBHFn9FR8jmnqt01M74d9AF1AMa4m2hhI,1040
 dbos/_schemas/system_database.py,sha256=3Z0L72bOgHnusK1hBaETWU9RfiLBP0QnS-fdu41i0yY,5835
 dbos/_serialization.py,sha256=bWuwhXSQcGmiazvhJHA5gwhrRWxtmFmcCFQSDJnqqkU,3666
-dbos/_sys_db.py,sha256=P4cxqDIF8EXfYuWA36gt-JHneDVO_lMwtnNrT0WVWB4,81311
+dbos/_sys_db.py,sha256=cgUpkgVF8jz_dcseShiJl4jFqoPlF24UadjVZ7LQ2qc,83235
 dbos/_templates/dbos-db-starter/README.md,sha256=GhxhBj42wjTt1fWEtwNriHbJuKb66Vzu89G4pxNHw2g,930
 dbos/_templates/dbos-db-starter/__package/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dbos/_templates/dbos-db-starter/__package/main.py.dbos,sha256=aQnBPSSQpkB8ERfhf7gB7P9tsU6OPKhZscfeh0yiaD8,2702
@@ -67,4 +67,4 @@ dbos/cli/cli.py,sha256=HinoCGrAUTiSeq7AAoCFfhdiE0uDw7vLMuDMN1_YTLI,20705
 dbos/dbos-config.schema.json,sha256=CjaspeYmOkx6Ip_pcxtmfXJTn_YGdSx_0pcPBF7KZmo,6060
 dbos/py.typed,sha256=QfzXT1Ktfk3Rj84akygc7_42z0lRpCq0Ilh8OXI6Zas,44
 version/__init__.py,sha256=L4sNxecRuqdtSFdpUGX3TtBi9KL3k7YsZVIvv-fv9-A,1678
-dbos-1.1.0a2.dist-info/RECORD,,
+dbos-1.1.0a3.dist-info/RECORD,,