pyworkflow-engine 0.1.20__tar.gz → 0.1.21__tar.gz

{pyworkflow_engine-0.1.20 → pyworkflow_engine-0.1.21}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyworkflow-engine
-Version: 0.1.20
+Version: 0.1.21
 Summary: A Python implementation of durable, event-sourced workflows inspired by Vercel Workflow
 Author: PyWorkflow Contributors
 License: MIT

{pyworkflow_engine-0.1.20 → pyworkflow_engine-0.1.21}/pyproject.toml

@@ -7,7 +7,7 @@ packages = [{include = "pyworkflow"}]
 
 [project]
 name = "pyworkflow-engine"
-version = "0.1.20"
+version = "0.1.21"
 description = "A Python implementation of durable, event-sourced workflows inspired by Vercel Workflow"
 readme = "README.md"
 requires-python = ">=3.11"

{pyworkflow_engine-0.1.20 → pyworkflow_engine-0.1.21}/pyworkflow/__init__.py

@@ -29,7 +29,7 @@ Quick Start:
     >>> run_id = await start(my_workflow, "Alice")
 """
 
-__version__ = "0.1.20"
+__version__ = "0.1.21"
 
 # Configuration
 from pyworkflow.config import (

{pyworkflow_engine-0.1.20 → pyworkflow_engine-0.1.21}/pyworkflow/celery/tasks.py

@@ -1154,8 +1154,14 @@ async def _recover_workflow_on_worker(
         recovery_attempt=run.recovery_attempts,
     )
 
-    # Update status to RUNNING (from RUNNING or INTERRUPTED)
-    await storage.update_run_status(run_id=run_id, status=RunStatus.RUNNING)
+    # Atomically claim the run for recovery.
+    # The run may be in INTERRUPTED or RUNNING state after a worker crash.
+    # Try INTERRUPTED -> RUNNING first (most common recovery path).
+    # If the run is already RUNNING, just set it to RUNNING (idempotent).
+    claimed = await storage.try_claim_run(run_id, RunStatus.INTERRUPTED, RunStatus.RUNNING)
+    if not claimed:
+        # May already be RUNNING from a previous partial recovery - update status directly
+        await storage.update_run_status(run_id=run_id, status=RunStatus.RUNNING)
 
     # Load event log for replay
     events = await storage.get_events(run_id)
@@ -2048,14 +2054,14 @@ async def _resume_workflow_on_worker(
         )
         return None
 
-    # Prevent duplicate resume execution
+    # Atomically claim the run: SUSPENDED -> RUNNING
     # Multiple resume tasks can be scheduled for the same workflow (e.g., race
-    # condition between step completion and suspension handler). Only proceed
-    # if the workflow is actually SUSPENDED. If status is RUNNING, another
-    # resume task got there first.
-    if run.status != RunStatus.SUSPENDED:
+    # condition between step completion and suspension handler). Only one
+    # succeeds; duplicates see the claim fail and return.
+    claimed = await storage.try_claim_run(run_id, RunStatus.SUSPENDED, RunStatus.RUNNING)
+    if not claimed:
         logger.info(
-            f"Workflow status is {run.status.value}, not SUSPENDED - skipping duplicate resume",
+            "Workflow status is not SUSPENDED (already claimed) - skipping duplicate resume",
             run_id=run_id,
             workflow_name=run.workflow_name,
         )
@@ -2075,6 +2081,8 @@ async def _resume_workflow_on_worker(
                 workflow_name=run.workflow_name,
                 triggered_by_hook_id=triggered_by_hook_id,
             )
+            # Revert status back to SUSPENDED since we won't actually resume
+            await storage.update_run_status(run_id=run_id, status=RunStatus.SUSPENDED)
             return None
 
     # Check for cancellation flag
@@ -2084,7 +2092,7 @@ async def _resume_workflow_on_worker(
         f"Resuming workflow execution on worker: {run.workflow_name}",
         run_id=run_id,
         workflow_name=run.workflow_name,
-        current_status=run.status.value,
+        current_status="running",
         cancellation_requested=cancellation_requested,
     )
 
@@ -2103,9 +2111,6 @@ async def _resume_workflow_on_worker(
     args = deserialize_args(run.input_args)
     kwargs = deserialize_kwargs(run.input_kwargs)
 
-    # Update status to running
-    await storage.update_run_status(run_id=run_id, status=RunStatus.RUNNING)
-
     # Execute workflow with event replay
     try:
         result = await execute_workflow_with_context(

{pyworkflow_engine-0.1.20 → pyworkflow_engine-0.1.21}/pyworkflow/core/step.py

@@ -595,6 +595,7 @@ async def _dispatch_step_to_celery(
     """
     from pyworkflow.celery.tasks import execute_step_task
     from pyworkflow.core.exceptions import SuspensionSignal
+    from pyworkflow.engine.events import EventType
 
     logger.info(
         f"Dispatching step to Celery worker: {step_name}",
@@ -602,6 +603,25 @@ async def _dispatch_step_to_celery(
         step_id=step_id,
     )
 
+    # Defense-in-depth: check if STEP_STARTED was already recorded for this step.
+    # This guards against duplicate dispatch when two resume tasks race and both
+    # replay past the same step. If already started, re-suspend to wait.
+    events = await ctx.storage.get_events(ctx.run_id)
+    already_started = any(
+        evt.type == EventType.STEP_STARTED and evt.data.get("step_id") == step_id for evt in events
+    )
+    if already_started:
+        logger.info(
+            f"Step {step_name} already has STEP_STARTED event, re-suspending",
+            run_id=ctx.run_id,
+            step_id=step_id,
+        )
+        raise SuspensionSignal(
+            reason=f"step_dispatch:{step_id}",
+            step_id=step_id,
+            step_name=step_name,
+        )
+
     # Validate event limits before recording step event
     await ctx.validate_event_limits()
 

{pyworkflow_engine-0.1.20 → pyworkflow_engine-0.1.21}/pyworkflow/storage/base.py

@@ -358,6 +358,35 @@ class StorageBackend(ABC):
         """
         pass
 
+    # Atomic Status Transition
+
+    async def try_claim_run(
+        self, run_id: str, from_status: RunStatus, to_status: RunStatus
+    ) -> bool:
+        """
+        Atomically transition run status if the current status matches.
+
+        This is a compare-and-swap operation: the status is only updated
+        if the current status equals `from_status`. Returns True if the
+        transition was applied, False if the current status did not match
+        (meaning another task already claimed this run).
+
+        Args:
+            run_id: Workflow run identifier
+            from_status: Expected current status
+            to_status: New status to set
+
+        Returns:
+            True if the transition succeeded, False otherwise
+        """
+        # Default implementation using get_run + update_run_status.
+        # Backends should override with truly atomic implementations.
+        run = await self.get_run(run_id)
+        if not run or run.status != from_status:
+            return False
+        await self.update_run_status(run_id, to_status)
+        return True
+
     # Cancellation Flag Operations
 
     @abstractmethod

{pyworkflow_engine-0.1.20 → pyworkflow_engine-0.1.21}/pyworkflow/storage/cassandra.py

@@ -1209,6 +1209,31 @@ class CassandraStorageBackend(StorageBackend):
         # Apply offset and limit
         return hooks[offset : offset + limit]
 
+    # Atomic Status Transition
+
+    async def try_claim_run(
+        self, run_id: str, from_status: RunStatus, to_status: RunStatus
+    ) -> bool:
+        """Atomically transition run status using lightweight transaction (IF)."""
+        session = self._ensure_connected()
+
+        result = session.execute(
+            SimpleStatement(
+                """
+                UPDATE workflow_runs
+                SET status = %s, updated_at = %s
+                WHERE run_id = %s
+                IF status = %s
+                """,
+                consistency_level=ConsistencyLevel.SERIAL,
+            ),
+            (to_status.value, datetime.now(UTC), run_id, from_status.value),
+        )
+
+        # Cassandra LWT returns [applied] column
+        row = result.one()
+        return row is not None and row[0] is True
+
     # Cancellation Flag Operations
 
     async def set_cancellation_flag(self, run_id: str) -> None:

{pyworkflow_engine-0.1.20 → pyworkflow_engine-0.1.21}/pyworkflow/storage/dynamodb.py

@@ -889,6 +889,37 @@ class DynamoDBStorageBackend(StorageBackend):
 
             return [self._item_to_hook(self._item_to_dict(item)) for item in items]
 
+    # Atomic Status Transition
+
+    async def try_claim_run(
+        self, run_id: str, from_status: RunStatus, to_status: RunStatus
+    ) -> bool:
+        """Atomically transition run status using conditional update."""
+        async with self._get_client() as client:
+            try:
+                now = datetime.now(UTC).isoformat()
+                await client.update_item(
+                    TableName=self.table_name,
+                    Key={
+                        "PK": {"S": f"RUN#{run_id}"},
+                        "SK": {"S": "#METADATA"},
+                    },
+                    UpdateExpression="SET #status = :new_status, updated_at = :now, GSI1SK = :gsi1sk",
+                    ConditionExpression="#status = :expected_status",
+                    ExpressionAttributeNames={"#status": "status"},
+                    ExpressionAttributeValues={
+                        ":new_status": {"S": to_status.value},
+                        ":expected_status": {"S": from_status.value},
+                        ":now": {"S": now},
+                        ":gsi1sk": {"S": f"{to_status.value}#{now}"},
+                    },
+                )
+                return True
+            except ClientError as e:
+                if e.response["Error"]["Code"] == "ConditionalCheckFailedException":
+                    return False
+                raise
+
     # Cancellation Flag Operations
 
     async def set_cancellation_flag(self, run_id: str) -> None:

{pyworkflow_engine-0.1.20 → pyworkflow_engine-0.1.21}/pyworkflow/storage/file.py

@@ -592,6 +592,34 @@ class FileStorageBackend(StorageBackend):
         hook_data_list = await asyncio.to_thread(_list)
         return [Hook.from_dict(data) for data in hook_data_list]
 
+    # Atomic Status Transition
+
+    async def try_claim_run(
+        self, run_id: str, from_status: RunStatus, to_status: RunStatus
+    ) -> bool:
+        """Atomically transition run status using file lock."""
+        run_file = self.runs_dir / f"{run_id}.json"
+
+        if not run_file.exists():
+            return False
+
+        lock_file = self.locks_dir / f"{run_id}.lock"
+        lock = FileLock(str(lock_file))
+
+        def _try_claim() -> bool:
+            with lock:
+                if not run_file.exists():
+                    return False
+                data = json.loads(run_file.read_text())
+                if data.get("status") != from_status.value:
+                    return False
+                data["status"] = to_status.value
+                data["updated_at"] = datetime.now(UTC).isoformat()
+                run_file.write_text(json.dumps(data, indent=2))
+                return True
+
+        return await asyncio.to_thread(_try_claim)
+
     # Cancellation Flag Operations
 
     async def set_cancellation_flag(self, run_id: str) -> None:

{pyworkflow_engine-0.1.20 → pyworkflow_engine-0.1.21}/pyworkflow/storage/memory.py

@@ -366,6 +366,20 @@ class InMemoryStorageBackend(StorageBackend):
             # Apply pagination
             return hooks[offset : offset + limit]
 
+    # Atomic Status Transition
+
+    async def try_claim_run(
+        self, run_id: str, from_status: RunStatus, to_status: RunStatus
+    ) -> bool:
+        """Atomically transition run status using lock-protected check-and-set."""
+        with self._lock:
+            run = self._runs.get(run_id)
+            if not run or run.status != from_status:
+                return False
+            run.status = to_status
+            run.updated_at = datetime.now(UTC)
+            return True
+
     # Cancellation Flag Operations
 
     async def set_cancellation_flag(self, run_id: str) -> None:

{pyworkflow_engine-0.1.20 → pyworkflow_engine-0.1.21}/pyworkflow/storage/mysql.py

@@ -767,6 +767,26 @@ class MySQLStorageBackend(StorageBackend):
 
         return [self._row_to_hook(row) for row in rows]
 
+    # Atomic Status Transition
+
+    async def try_claim_run(
+        self, run_id: str, from_status: RunStatus, to_status: RunStatus
+    ) -> bool:
+        """Atomically transition run status using conditional UPDATE."""
+        pool = self._ensure_connected()
+
+        async with pool.acquire() as conn, conn.cursor() as cur:
+            await cur.execute(
+                """
+                UPDATE workflow_runs
+                SET status = %s, updated_at = %s
+                WHERE run_id = %s AND status = %s
+                """,
+                (to_status.value, datetime.now(UTC), run_id, from_status.value),
+            )
+
+            return cur.rowcount > 0
+
     # Cancellation Flag Operations
 
     async def set_cancellation_flag(self, run_id: str) -> None:

{pyworkflow_engine-0.1.20 → pyworkflow_engine-0.1.21}/pyworkflow/storage/postgres.py

@@ -862,6 +862,30 @@ class PostgresStorageBackend(StorageBackend):
 
         return [self._row_to_hook(row) for row in rows]
 
+    # Atomic Status Transition
+
+    async def try_claim_run(
+        self, run_id: str, from_status: RunStatus, to_status: RunStatus
+    ) -> bool:
+        """Atomically transition run status using conditional UPDATE."""
+        pool = await self._get_pool()
+
+        async with pool.acquire() as conn:
+            result = await conn.execute(
+                """
+                UPDATE workflow_runs
+                SET status = $1, updated_at = $2
+                WHERE run_id = $3 AND status = $4
+                """,
+                to_status.value,
+                datetime.now(UTC),
+                run_id,
+                from_status.value,
+            )
+
+        # asyncpg returns 'UPDATE N' where N is rows affected
+        return result == "UPDATE 1"
+
     # Cancellation Flag Operations
 
     async def set_cancellation_flag(self, run_id: str) -> None:

{pyworkflow_engine-0.1.20 → pyworkflow_engine-0.1.21}/pyworkflow/storage/sqlite.py

@@ -750,6 +750,26 @@ class SQLiteStorageBackend(StorageBackend):
 
         return [self._row_to_hook(row) for row in rows]
 
+    # Atomic Status Transition
+
+    async def try_claim_run(
+        self, run_id: str, from_status: RunStatus, to_status: RunStatus
+    ) -> bool:
+        """Atomically transition run status using conditional UPDATE."""
+        db = self._ensure_connected()
+
+        cursor = await db.execute(
+            """
+            UPDATE workflow_runs
+            SET status = ?, updated_at = ?
+            WHERE run_id = ? AND status = ?
+            """,
+            (to_status.value, datetime.now(UTC).isoformat(), run_id, from_status.value),
+        )
+        await db.commit()
+
+        return cursor.rowcount > 0
+
     # Cancellation Flag Operations
 
     async def set_cancellation_flag(self, run_id: str) -> None: