planar 0.5.0__py3-none-any.whl

planar/workflows/test_concurrency_detection.py

@@ -0,0 +1,120 @@
+import asyncio
+import multiprocessing
+import multiprocessing.connection
+from multiprocessing.connection import Connection
+from uuid import UUID
+
+from planar.db import DatabaseManager, new_session
+from planar.session import engine_var, get_engine, session_var
+from planar.workflows.decorators import step, workflow
+from planar.workflows.exceptions import LockResourceFailed
+from planar.workflows.execution import (
+    _DEFAULT_LOCK_DURATION,
+    execute,
+)
+from planar.workflows.lock import lock_workflow
+from planar.workflows.models import Workflow, WorkflowStatus
+
+# bidirectional communication between the test process and the worker processes.
+conn: Connection
+
+
+@step(max_retries=0)
+async def dummy_step():
+    conn.send("waiting")
+    # Wait until "proceed" is received from the queue.
+    if conn.recv() != "proceed":
+        raise Exception('Expected "proceed"')
+    return "success"
+
+
+@workflow()
+async def dummy_workflow():
+    # Run the dummy step and return its result.
+    result = await dummy_step()
+    return result
+
+
+# copy of the resume_workflow function which allows more fine grained control from
+# the test process. This is fine because our goal is to test concurrency detection
+# implemented by the execute function.
+async def resume_with_semaphores(workflow_id: UUID):
+    engine = get_engine()
+    async with new_session(engine) as session:
+        tok = session_var.set(session)
+        try:
+            async with session.begin():
+                workflow = await session.get(Workflow, workflow_id)
+            if not workflow:
+                raise ValueError(f"Workflow {workflow_id} not found")
+            conn.send("ready")
+            # Wait until "start" is received on stdin.
+            if conn.recv() != "start":
+                raise Exception('Expected "start"')
+            async with lock_workflow(
+                workflow,
+                _DEFAULT_LOCK_DURATION,
+                retry_count=0,
+            ):
+                await execute(workflow)
+            conn.send("completed")
+        except LockResourceFailed:
+            conn.send("conflict")
+        finally:
+            session_var.reset(tok)
+
+
+# This worker function will be launched as a separate process.
+# It takes the workflow id, db_url and a multiprocess Pipe.
+def worker(wf_id: UUID, db_url: str, connection: Connection):
+    global conn
+    conn = connection
+    # Create a new engine for this process.
+    db_manager = DatabaseManager(db_url)
+    db_manager.connect()
+    engine = db_manager.get_engine()
+    engine_var.set(engine)
+    # Run the resume_with_semaphores coroutine.
+    # We use asyncio.run so that the worker’s event loop is independent.
+    asyncio.run(resume_with_semaphores(wf_id))
+
+
+async def test_concurrent_workflow_execution(tmp_db_url, tmp_db_engine):
+    async with new_session(tmp_db_engine) as session:
+        session_var.set(session)
+        wf: Workflow = await dummy_workflow.start()
+        wf_id = wf.id
+
+        # Launch two separate processes that attempt to resume the workflow concurrently.
+        p1_parent, p1_worker = multiprocessing.Pipe(duplex=True)
+        p2_parent, p2_worker = multiprocessing.Pipe(duplex=True)
+        p1 = multiprocessing.Process(target=worker, args=(wf_id, tmp_db_url, p1_worker))
+        p2 = multiprocessing.Process(target=worker, args=(wf_id, tmp_db_url, p2_worker))
+        p1.start()
+        p2.start()
+        # wait for both workers to fetch the workflow from the database.
+        assert p1_parent.recv() == "ready"
+        assert p2_parent.recv() == "ready"
+        # allow worker 1 to proceed.
+        p1_parent.send("start")
+        # wait for worker 1 to start the workflow and pause in the dummy step.
+        assert p1_parent.recv() == "waiting"
+        # allow worker 2 to proceed.
+        p2_parent.send("start")
+        # worker 2 should fail and will send a "conflict" message.
+        assert p2_parent.recv() == "conflict"
+        # allow worker 1 to proceed
+        p1_parent.send("proceed")
+        # worker 1 should complete the workflow and send a "completed" message.
+        assert p1_parent.recv() == "completed"
+        # cleanup workers
+        p1.join()
+        p2.join()
+
+        await session.refresh(wf)
+        assert wf, f"Workflow {wf_id} not found"
+        # Assert that the workflow completed successfully.
+        assert wf.status == WorkflowStatus.SUCCEEDED, (
+            f"Unexpected workflow status: {wf.status}"
+        )
+        assert wf.result == "success", f"Unexpected workflow result: {wf.result}"

planar/workflows/test_lock_timeout.py

@@ -0,0 +1,140 @@
+import asyncio
+from datetime import timedelta
+
+from planar.db import new_session
+from planar.session import session_var
+from planar.testing.synchronizable_tracer import SynchronizableTracer, TraceSpec
+from planar.utils import utc_now
+from planar.workflows.decorators import workflow
+from planar.workflows.execution import execute
+from planar.workflows.models import (
+    LockedResource,
+    Workflow,
+    WorkflowStatus,
+    workflow_exec_lock_key,
+)
+from planar.workflows.orchestrator import WorkflowOrchestrator
+from planar.workflows.step_core import Suspend, suspend
+from planar.workflows.tracing import tracer_var
+
+
+# Define a long-running workflow.
+@workflow()
+async def long_running_workflow():
+    # Simulate a long-running operation by sleeping 1 second.
+    await asyncio.sleep(1)
+    return "finished"
+
+
+async def test_lock_timer_extension(mem_db_engine):
+    tracer = SynchronizableTracer()
+    tracer_var.set(tracer)
+    lock_acquired = tracer.instrument(
+        TraceSpec(function_name="lock_resource", message="commit")
+    )
+    lock_heartbeat = tracer.instrument(
+        TraceSpec(function_name="lock_heartbeat", message="commit")
+    )
+
+    async with new_session(mem_db_engine) as session:
+        # This test verifies that when a workflow is executing, the heartbeat task
+        # (lock_heartbeat) extends the workflow's lock_until field. We run a
+        # long-running workflow (which sleeps for 1 second) with a short lock
+        # duration and heartbeat interval. While the workflow is running we query
+        # the stored workflow record and ensure that lock_until is updated
+        # (extended) by the heartbeat.
+
+        session_var.set(session)
+        # Start the workflow.
+        # Run workflow execution in the background with short durations so
+        # heartbeat kicks in quickly.
+        async with WorkflowOrchestrator.ensure_started(
+            lock_duration=timedelta(seconds=1)
+        ) as orchestrator:
+            wf: Workflow = await long_running_workflow.start()
+            wf_id = wf.id
+            lock_key = workflow_exec_lock_key(wf_id)
+
+            await lock_acquired.wait()
+
+            async with session.begin():
+                locked_resource = await session.get(LockedResource, lock_key)
+            assert locked_resource, "Expected a locked resource record"
+            lock_time_1 = locked_resource.lock_until
+            assert lock_time_1, "Lock time should be set"
+
+            # Wait a bit longer to allow another heartbeat cycle.
+            await lock_heartbeat.wait()
+            async with session.begin():
+                await session.refresh(locked_resource)
+            lock_time_2 = locked_resource.lock_until
+            assert lock_time_2, "Lock time should be set"
+
+            # The lock_time_2 should be later than lock_time_1 if the heartbeat is working.
+            assert lock_time_2 > lock_time_1, (
+                f"Expected lock_until to be extended by heartbeat: {lock_time_1} vs {lock_time_2}"
+            )
+
+            # Let the workflow finish.
+            await orchestrator.wait_for_completion(wf_id)
+
+            # Verify the workflow completed successfully.
+            await session.refresh(wf)
+            assert wf.status == WorkflowStatus.SUCCEEDED
+            assert wf.result == "finished"
+
+
+@workflow()
+async def crashed_worker_workflow():
+    # This workflow uses suspend() to simulate work that is paused.
+    # The first execution returns a Suspend object.
+    # When resumed it completes and returns "completed".
+    # First step: suspend (simulate waiting, e.g. because a worker had locked it).
+    await suspend(interval=timedelta(seconds=5))
+    # After the suspension it resumes here.
+    return "completed"
+
+
+async def test_orchestrator_resumes_crashed_worker(mem_db_engine):
+    # This test simulates the scenario where a worker has “crashed” after
+    # locking a workflow. We start a workflow that suspends. Then we add a LockedResource
+    # record with an expired lock_until time to simulate a crashed
+
+    # Invoking the workflow_orchestrator (which polls for suspended workflows
+    # whose wakeup time is reached or that have expired locks) should cause the
+    # the workflow to be resumed. Finally, we verify that the workflow
+    # completes successfully. Start the workflow – its first execution will
+    # suspend.
+    async with new_session(mem_db_engine) as session:
+        session_var.set(session)
+        wf = await crashed_worker_workflow.start()
+
+        result = await execute(wf)
+        assert isinstance(result, Suspend), (
+            "Expected the workflow to suspend on first execution."
+        )
+        # Simulate a crashed worker by directly changing the workflow record.
+        await session.refresh(wf)
+        # Force wakeup_at and lock_until to be in the past.
+        past_time = utc_now() - timedelta(seconds=1)
+        wf.wakeup_at = past_time
+        session.add(LockedResource(lock_key=f"workflow:{wf.id}", lock_until=past_time))
+        # Ensure it is marked as running, which would not normally be picked by
+        # the orchestrator
+        await session.commit()
+
+        # Now run the orchestrator, which polls for suspended workflows with
+        # wakeup_at <= now.
+        # We use a short poll interval.
+        async with WorkflowOrchestrator.ensure_started(
+            poll_interval=0.2
+        ) as orchestrator:
+            await orchestrator.wait_for_completion(wf.id)
+
+        await session.refresh(wf)
+        assert wf.status == WorkflowStatus.SUCCEEDED, (
+            f"Expected workflow status 'success' but got {wf.status}"
+        )
+        assert wf.result == "completed", (
+            f"Expected workflow result 'completed' but got {wf.result}"
+        )