Drain cooperative local cleanup and fence it at shutdown deadline

Author: durable-workflow-ops

src/durable_workflow/worker.py

@@ -17,6 +17,7 @@
 
 import asyncio
 import contextlib
+import contextvars
 import hashlib
 import inspect
 import json
@@ -28,6 +29,7 @@
 import types
 import uuid
 from collections.abc import Awaitable, Callable, Coroutine, Iterable, Mapping
+from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime, timezone
 from functools import wraps
 from types import FunctionType
@@ -1056,6 +1058,8 @@ def __init__(
         self.max_concurrent_worker_sessions = max_concurrent_worker_sessions
         self._worker_sessions: dict[str, WorkerSession] = {}
         self._stop = asyncio.Event()
+        self._local_activity_shutdown = asyncio.Event()
+        self._local_activity_executor: ThreadPoolExecutor | None = None
         self._wf_semaphore = asyncio.Semaphore(max_concurrent_workflow_tasks)
         self._act_semaphore = asyncio.Semaphore(max_concurrent_activity_tasks)
         self._shutdown_timeout = shutdown_timeout
@@ -1395,6 +1399,8 @@ async def _resolve_workflow_nexus_commands(
         raise RuntimeError("workflow yielded too many consecutive Nexus service calls")
 
     async def _renew_local_workflow_lease(self, task: dict[str, Any]) -> None:
+        if self._cooperative_cancellation_supported and self._local_activity_shutdown.is_set():
+            raise LocalActivityExecutionAborted("worker shutdown abandoned its local workflow claim")
         task_id = str(task["task_id"])
         attempt = int(task.get("workflow_task_attempt", 1))
         try:
@@ -1405,6 +1411,8 @@ async def _renew_local_workflow_lease(self, task: dict[str, Any]) -> None:
             )
         except Exception as exc:
             raise LocalActivityExecutionAborted("workflow task lease renewal failed") from exc
+        if self._cooperative_cancellation_supported and self._local_activity_shutdown.is_set():
+            raise LocalActivityExecutionAborted("worker shutdown abandoned its local workflow claim")
         if not isinstance(response, Mapping) or any((
             response.get("task_id") != task_id,
             response.get("lease_owner") != self.worker_id,
@@ -1548,16 +1556,22 @@ async def observe_lease() -> None:
                 await self._renew_local_workflow_lease(task)
 
         invocation = asyncio.create_task(self._execute_activity_callable(
-            task, command.activity_type, tuple(command.arguments), handler,
+            task, command.activity_type, tuple(command.arguments), handler, run_sync_in_thread=True,
         ))
         observation = asyncio.create_task(observe_lease())
+        shutdown = asyncio.create_task(self._local_activity_shutdown.wait())
         try:
-            done, _ = await asyncio.wait([invocation, observation], return_when=asyncio.FIRST_COMPLETED)
+            done, _ = await asyncio.wait([invocation, observation, shutdown], return_when=asyncio.FIRST_COMPLETED)
+            if shutdown in done:
+                raise LocalActivityExecutionAborted("worker shutdown abandoned its local workflow claim")
             if observation in done:
                 await observation  # Propagate transport observation or lost lease, never a workflow cancellation.
                 raise LocalActivityExecutionAborted("local lease observer stopped without an acknowledgment")
             return await invocation
         finally:
+            shutdown.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await shutdown
             observation.cancel()
             with contextlib.suppress(asyncio.CancelledError, Exception):
                 await observation
@@ -1651,7 +1665,10 @@ def check_boundary(
                 now = time.monotonic()
                 if state["lease_aborted"]:
                     raise LocalActivityExecutionAborted("local activity lost its workflow task lease")
-                if self._stop.is_set() or (
+                if self._cooperative_cancellation_supported and self._local_activity_shutdown.is_set():
+                    state["lease_aborted"] = True
+                    raise LocalActivityExecutionAborted("worker shutdown abandoned its local workflow claim")
+                if (self._stop.is_set() and not self._cooperative_cancellation_supported) or (
                     task.get("cancel_requested") is True and task.get("cancellation_request") is None
                     and task.get("_delivered_cancellation_request_id") is None
                 ):
@@ -2320,6 +2337,7 @@ async def _execute_activity_callable(
         activity_type: str,
         args: tuple[Any, ...],
         fn: Callable[..., Any],
+        *, run_sync_in_thread: bool = False,
     ) -> Any:
         context = ActivityInterceptorContext(
             worker_id=self.worker_id,
@@ -2330,7 +2348,17 @@ async def _execute_activity_callable(
         )
 
         async def call_activity(ctx: ActivityInterceptorContext) -> Any:
-            result = fn(*ctx.args)
+            if run_sync_in_thread and not inspect.iscoroutinefunction(fn):
+                if self._local_activity_executor is None:
+                    # Replay threads wait for local results, so sharing their pool can deadlock.
+                    self._local_activity_executor = ThreadPoolExecutor(
+                        max_workers=self.max_concurrent_workflow_tasks, thread_name_prefix="dw-local-activity",
+                    )
+                result = await asyncio.get_running_loop().run_in_executor(
+                    self._local_activity_executor, contextvars.copy_context().run, fn, *ctx.args,
+                )
+            else:
+                result = fn(*ctx.args)
             if asyncio.iscoroutine(result):
                 return await result
             return result
@@ -3527,6 +3555,8 @@ async def _shutdown(self) -> None:
                 in_flight,
                 timeout=self._remaining_shutdown_time(deadline),
             )
+            if pending:
+                self._local_activity_shutdown.set()
             for t in pending:
                 t.cancel()
             if pending:
@@ -3541,6 +3571,9 @@ async def _shutdown(self) -> None:
                 )
             await asyncio.gather(*in_flight, return_exceptions=True)
 
+        if self._local_activity_executor is not None:
+            self._local_activity_executor.shutdown(wait=False, cancel_futures=True)
+
         for session in self._worker_sessions.values():
             if not session.active:
                 continue

tests/test_cooperative_cancellation_worker.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import asyncio
+import threading
+from concurrent.futures import ThreadPoolExecutor
 from copy import deepcopy
 from typing import Any
 from unittest.mock import AsyncMock
@@ -11,6 +13,7 @@
 from durable_workflow.client import Client
 from durable_workflow.errors import ServerError, WorkflowCancelled
 from durable_workflow.worker import Worker
+from durable_workflow.workflow import LocalActivityExecutionAborted
 from tests.test_cooperative_cancellation import marker, observation, request
 from tests.test_worker import compatible_cluster_info
 
@@ -70,8 +73,13 @@ def lease_ack(*, observed: bool = False) -> dict[str, Any]:
 
 
 class ClaimServer:
-    def __init__(self, *, history: list[dict[str, Any]] | None = None) -> None:
+    def __init__(
+        self, *, history: list[dict[str, Any]] | None = None,
+        lease_owner: str = "cooperative-worker", workflow_task_attempt: int = 4,
+    ) -> None:
         self.client = AsyncMock(spec=Client)
+        self.lease_owner = lease_owner
+        self.workflow_task_attempt = workflow_task_attempt
         self.history = list(history if history is not None else [request()])
         self.trace: list[str] = []
         self.delivery_error: Exception | None = None
@@ -83,15 +91,17 @@ def __init__(self, *, history: list[dict[str, Any]] | None = None) -> None:
             },
         })
         self.client.register_worker.return_value = {"registered": True}
-        self.client.heartbeat_workflow_task.return_value = lease_ack()
+        self.client.heartbeat_workflow_task.return_value = {
+            **lease_ack(), "lease_owner": lease_owner, "workflow_task_attempt": workflow_task_attempt,
+        }
         self.client.workflow_task_history.side_effect = self.page
         self.client.deliver_workflow_cancellation.side_effect = self.deliver
         self.client.complete_workflow_task.side_effect = self.complete
 
     async def page(self, **kwargs: Any) -> dict[str, Any]:
         assert kwargs == {
             "task_id": "task-1", "next_history_page_token": "opaque-first-page",
-            "lease_owner": "cooperative-worker", "workflow_task_attempt": 4,
+            "lease_owner": self.lease_owner, "workflow_task_attempt": self.workflow_task_attempt,
         }
         self.trace.append("history")
         return {"history_events": deepcopy(self.history), "next_history_page_token": None}
@@ -115,7 +125,7 @@ async def complete(self, **kwargs: Any) -> dict[str, Any]:
     async def worker(self, monkeypatch: pytest.MonkeyPatch, **kwargs: Any) -> Worker:
         monkeypatch.setenv("DURABLE_WORKFLOW_WORKER_PROTOCOL_VERSION", "1.20")
         worker = Worker(
-            self.client, task_queue="queue", worker_id="cooperative-worker",
+            self.client, task_queue="queue", worker_id=self.lease_owner,
             workflows=kwargs.pop("workflows", [CancellationWorkflow]),
             capabilities=["cooperative_cancellation"], **kwargs,
         )
@@ -324,6 +334,175 @@ async def compensate(request_id: str) -> str:
     assert server.client.deliver_workflow_cancellation.await_count == 1
 
 
+async def test_worker_stop_drains_shielded_local_cleanup_within_its_grace_period(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    server = ClaimServer()
+    entered = asyncio.Event()
+    release = asyncio.Event()
+
+    @activity.defn(name="cleanup")
+    async def compensate(request_id: str) -> str:
+        entered.set()
+        await release.wait()
+        await activity.context().heartbeat()
+        return request_id
+
+    worker = await server.worker(
+        monkeypatch, workflows=[LocalCleanupWorkflow], activities=[compensate], shutdown_timeout=1,
+    )
+    task = claimed_task()
+    task.update(workflow_type="worker-cancellation-local-cleanup", arguments=serializer.envelope([], codec="avro"))
+    execution = worker._track(worker._run_workflow_task(task))
+    await asyncio.wait_for(entered.wait(), timeout=1)
+    stopping = asyncio.create_task(worker.stop())
+    await asyncio.sleep(0)
+    assert worker._stop.is_set()
+    release.set()
+    commands = await asyncio.wait_for(execution, timeout=1)
+    await asyncio.wait_for(stopping, timeout=1)
+    assert commands is not None
+    assert [command["type"] for command in commands] == ["record_local_activity", "complete_workflow"]
+    server.client.deliver_workflow_cancellation.assert_awaited_once()
+    server.client.fail_workflow_task.assert_not_awaited()
+    server.client.deregister_worker_registration.assert_awaited_once_with("cooperative-worker")
+
+
+@pytest.mark.parametrize("handler_kind", ["async", "sync"])
+async def test_shutdown_timeout_abandons_cleanup_and_replacement_replays_original_delivery(
+    monkeypatch: pytest.MonkeyPatch, handler_kind: str,
+) -> None:
+    server = ClaimServer()
+    entered = asyncio.Event()
+    discarded = asyncio.Event()
+    loop = asyncio.get_running_loop()
+    release_thread = threading.Event()
+    fenced_heartbeats: list[bool] = []
+
+    @activity.defn(name="cleanup")
+    async def compensate(request_id: str) -> object:
+        assert request_id == "request-1"
+        entered.set()
+        try:
+            await asyncio.Event().wait()
+        except asyncio.CancelledError:
+            with pytest.raises(LocalActivityExecutionAborted):
+                await activity.context().heartbeat()
+            fenced_heartbeats.append(True)
+            discarded.set()
+            return object()  # A late, unencodable result must never become an activity failure.
+
+    @activity.defn(name="cleanup")
+    def synchronous_compensation(request_id: str) -> object:
+        assert request_id == "request-1"
+        context = activity.context()
+        assert context.info.worker_id == "cooperative-worker"
+        loop.call_soon_threadsafe(entered.set)
+        try:
+            assert release_thread.wait(timeout=2)
+            heartbeat = asyncio.run_coroutine_threadsafe(context.heartbeat(), loop)
+            try:
+                heartbeat.result(timeout=1)
+            except LocalActivityExecutionAborted:
+                fenced_heartbeats.append(True)
+            else:
+                fenced_heartbeats.append(False)
+            return object()
+        finally:
+            loop.call_soon_threadsafe(discarded.set)
+
+    worker = await server.worker(
+        monkeypatch, workflows=[LocalCleanupWorkflow],
+        activities=[compensate if handler_kind == "async" else synchronous_compensation], shutdown_timeout=0.01,
+    )
+    task = claimed_task()
+    task.update(workflow_type="worker-cancellation-local-cleanup", arguments=serializer.envelope([], codec="avro"))
+    execution = worker._track(worker._run_workflow_task(task))
+    await asyncio.wait_for(entered.wait(), timeout=1)
+    original_history = deepcopy(server.history)
+    heartbeats_before_stop = server.client.heartbeat_workflow_task.await_count
+    await asyncio.wait_for(worker.stop(), timeout=1)
+    release_thread.set()
+    await asyncio.wait_for(discarded.wait(), timeout=1)
+    assert fenced_heartbeats == [True]
+    assert server.client.heartbeat_workflow_task.await_count == heartbeats_before_stop
+    assert execution.cancelled() or execution.result() is None
+    assert server.history == original_history
+    server.client.complete_workflow_task.assert_not_awaited()
+    server.client.fail_workflow_task.assert_not_awaited()
+    server.client.deregister_worker_registration.assert_awaited_once_with("cooperative-worker")
+
+    replacement = ClaimServer(history=original_history, lease_owner="replacement-worker", workflow_task_attempt=5)
+    cleanup_ids: list[str] = []
+
+    @activity.defn(name="cleanup")
+    async def resumed_cleanup(request_id: str) -> str:
+        cleanup_ids.append(request_id)
+        await activity.context().heartbeat()
+        return request_id
+
+    successor = await replacement.worker(monkeypatch, workflows=[LocalCleanupWorkflow], activities=[resumed_cleanup])
+    reclaimed = deepcopy(task)
+    reclaimed["workflow_task_attempt"] = 5
+    reclaimed["history_events"] = deepcopy(original_history)
+    commands = await asyncio.wait_for(successor._run_workflow_task(reclaimed), timeout=1)
+    assert cleanup_ids == ["request-1"]
+    assert commands is not None
+    assert [command["type"] for command in commands] == ["record_local_activity", "complete_workflow"]
+    assert replacement.history == original_history
+    replacement.client.deliver_workflow_cancellation.assert_not_awaited()
+    replacement.client.fail_workflow_task.assert_not_awaited()
+    assert replacement.client.complete_workflow_task.await_args.kwargs["lease_owner"] == "replacement-worker"
+    assert replacement.client.complete_workflow_task.await_args.kwargs["workflow_task_attempt"] == 5
+
+
+async def test_active_synchronous_local_call_observes_request_without_blocking_the_worker(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    server = ClaimServer(history=[])
+    entered = asyncio.Event()
+    discarded = asyncio.Event()
+    release = threading.Event()
+    loop = asyncio.get_running_loop()
+    fenced: list[bool] = []
+    loop.set_default_executor(ThreadPoolExecutor(max_workers=1))
+
+    @activity.defn(name="work")
+    def work() -> object:
+        context = activity.context()
+        loop.call_soon_threadsafe(entered.set)
+        try:
+            assert release.wait(timeout=2)
+            heartbeat = asyncio.run_coroutine_threadsafe(context.heartbeat(), loop)
+            try:
+                heartbeat.result(timeout=1)
+            except LocalActivityExecutionAborted:
+                fenced.append(True)
+            else:
+                fenced.append(False)
+            return object()
+        finally:
+            loop.call_soon_threadsafe(discarded.set)
+
+    worker = await server.worker(monkeypatch, activities=[work], heartbeat_interval=0.01)
+    execution = worker._track(worker._run_workflow_task(claimed_task("local_activity", observed=False)))
+    await asyncio.wait_for(entered.wait(), timeout=1)
+    server.history.append(request())
+    server.client.heartbeat_workflow_task.return_value = lease_ack(observed=True)
+    try:
+        commands = await asyncio.wait_for(execution, timeout=1)
+        assert commands is not None and [command["type"] for command in commands] == ["start_timer"]
+        heartbeats_after_delivery = server.client.heartbeat_workflow_task.await_count
+    finally:
+        release.set()
+    await asyncio.wait_for(discarded.wait(), timeout=1)
+    assert fenced == [True]
+    assert server.client.heartbeat_workflow_task.await_count == heartbeats_after_delivery
+    assert server.client.deliver_workflow_cancellation.await_args.kwargs["call_kind"] == "local_activity"
+    server.client.fail_workflow_task.assert_not_awaited()
+    await worker.stop()
+
+
 @pytest.mark.parametrize("value", ["", None, {}, "not-a-page"])
 async def test_invalid_or_missing_refresh_token_never_delivers(monkeypatch: pytest.MonkeyPatch, value: Any) -> None:
     server = ClaimServer()