Cancel activities Core stopped tracking at worker shutdown

CHANGELOG.md

@@ -52,6 +52,10 @@ to include examples, links to docs, or any other relevant information.
 
 ### Fixed
 
+- Worker shutdown no longer waits forever for an activity that Core has stopped tracking, such as a
+  local activity whose cancellation was lost when its workflow run was evicted. Once activity polling
+  has shut down, any activity still executing is cancelled, with `worker_shutdown` cancellation
+  details if it has none yet.
 - **Experimental**: External storage metrics now report the wall-clock time storage was in flight.
   Previously each batch's duration was summed, over-reporting the time whenever storage operations
   ran concurrently.

temporalio/worker/_activity.py

@@ -195,6 +195,21 @@ async def drain_poll_queue(self) -> None:
     # Only call this after run()/drain_poll_queue() have returned. This will not
     # raise an exception.
     async def wait_all_completed(self) -> None:
+        # Core tracks no activities once polling has shut down, so cancel stragglers
+        for task_token, activity in self._running_activities.items():
+            if not activity.done:
+                logger.warning(
+                    "Cancelling activity %s still running after worker shutdown",
+                    task_token,
+                )
+                # Cancellation details are set once, so keep any already received
+                if not activity.cancellation_details.details:
+                    activity.cancellation_details.details = (
+                        temporalio.activity.ActivityCancellationDetails(
+                            worker_shutdown=True
+                        )
+                    )
+                activity.cancel(cancelled_by_request=True)
         running_tasks = [v.task for v in self._running_activities.values() if v.task]
         if running_tasks:
             await asyncio.gather(*running_tasks, return_exceptions=False)

tests/worker/test_workflow.py

@@ -27,6 +27,7 @@
     NoReturn,
     cast,
 )
+from unittest.mock import patch
 from urllib.request import urlopen
 
 import pydantic
@@ -36,11 +37,13 @@
 
 import temporalio.activity
 import temporalio.api.sdk.v1
+import temporalio.bridge.worker
 import temporalio.client
 import temporalio.converter
 import temporalio.converter._extstore
 import temporalio.worker
 import temporalio.worker._command_aware_visitor
+import temporalio.worker._workflow
 import temporalio.worker._workflow_instance
 import temporalio.workflow
 from temporalio import activity, workflow
@@ -53,7 +56,10 @@
     PollWorkflowExecutionUpdateResponse,
     ResetStickyTaskQueueRequest,
 )
-from temporalio.bridge.proto.workflow_activation import WorkflowActivation
+from temporalio.bridge.proto.workflow_activation import (
+    RemoveFromCache,
+    WorkflowActivation,
+)
 from temporalio.bridge.proto.workflow_completion import WorkflowActivationCompletion
 from temporalio.client import (
     AsyncActivityCancelledError,
@@ -1135,6 +1141,138 @@ async def activity_started() -> bool:
     )
 
 
+@workflow.defn
+class OrphanedLocalActivityWorkflow:
+    @workflow.run
+    async def run(self) -> None:
+        await workflow.execute_local_activity(
+            "wait_forever_local_activity",
+            start_to_close_timeout=timedelta(minutes=5),
+        )
+
+
+async def test_worker_shutdown_cancels_local_activity_untracked_by_core(
+    client: Client,
+):
+    started = asyncio.Event()
+    workflow_poll_shut_down = asyncio.Event()
+    details: list[temporalio.activity.ActivityCancellationDetails | None] = []
+
+    @activity.defn(name="wait_forever_local_activity")
+    async def wait_forever_local_activity() -> None:
+        started.set()
+        try:
+            await asyncio.sleep(1000)
+        except asyncio.CancelledError:
+            details.append(activity.cancellation_details())
+
+    # Hold the next poll until Core has evicted the run, so its queued cancel is never delivered
+    bridge_worker = temporalio.bridge.worker.Worker
+    orig_poll_activity = bridge_worker.poll_activity_task
+    orig_poll_workflow = bridge_worker.poll_workflow_activation
+    hold_next_poll = False
+
+    async def poll_activity_task(self: temporalio.bridge.worker.Worker):
+        nonlocal hold_next_poll
+        if hold_next_poll:
+            hold_next_poll = False
+            await workflow_poll_shut_down.wait()
+        task = await orig_poll_activity(self)
+        if task.HasField("start") and task.start.is_local:
+            hold_next_poll = True
+        return task
+
+    async def poll_workflow_activation(self: temporalio.bridge.worker.Worker):
+        try:
+            return await orig_poll_workflow(self)
+        except temporalio.bridge.worker.PollShutdownError:  # type: ignore[reportPrivateLocalImportUsage]
+            workflow_poll_shut_down.set()
+            raise
+
+    with (
+        patch.object(bridge_worker, "poll_activity_task", poll_activity_task),
+        patch.object(
+            bridge_worker, "poll_workflow_activation", poll_workflow_activation
+        ),
+    ):
+        worker = new_worker(
+            client,
+            OrphanedLocalActivityWorkflow,
+            activities=[wait_forever_local_activity],
+        )
+        run_task = asyncio.create_task(worker.run())
+        handle = await client.start_workflow(
+            OrphanedLocalActivityWorkflow.run,
+            id=f"workflow-{uuid.uuid4()}",
+            task_queue=worker.task_queue,
+            task_timeout=timedelta(seconds=3),
+        )
+        await started.wait()
+        # Terminating fails the workflow task heartbeat, which evicts the run
+        await handle.terminate()
+        await asyncio.wait_for(worker.shutdown(), 20)
+        await run_task
+    assert details == [
+        temporalio.activity.ActivityCancellationDetails(worker_shutdown=True)
+    ]
+
+
+async def test_worker_shutdown_keeps_details_of_local_activity_ignoring_cancel(
+    client: Client,
+):
+    started = asyncio.Event()
+    cancel_seen = asyncio.Event()
+    details: list[temporalio.activity.ActivityCancellationDetails | None] = []
+
+    @activity.defn(name="wait_forever_local_activity")
+    async def wait_forever_local_activity() -> None:
+        started.set()
+        while True:
+            try:
+                await asyncio.sleep(1000)
+            except asyncio.CancelledError:
+                details.append(activity.cancellation_details())
+                cancel_seen.set()
+                if activity.is_worker_shutdown():
+                    raise
+
+    # Finish evicting only after the cancel reached the activity, so Core drops it afterwards
+    workflow_worker = temporalio.worker._workflow._WorkflowWorker
+    orig_evict = workflow_worker._handle_cache_eviction
+
+    async def handle_cache_eviction(
+        self: temporalio.worker._workflow._WorkflowWorker,
+        act: WorkflowActivation,
+        job: RemoveFromCache,
+    ):
+        await cancel_seen.wait()
+        await orig_evict(self, act, job)
+
+    with patch.object(workflow_worker, "_handle_cache_eviction", handle_cache_eviction):
+        worker = new_worker(
+            client,
+            OrphanedLocalActivityWorkflow,
+            activities=[wait_forever_local_activity],
+        )
+        run_task = asyncio.create_task(worker.run())
+        handle = await client.start_workflow(
+            OrphanedLocalActivityWorkflow.run,
+            id=f"workflow-{uuid.uuid4()}",
+            task_queue=worker.task_queue,
+            task_timeout=timedelta(seconds=3),
+        )
+        await started.wait()
+        # Terminating fails the workflow task heartbeat, which evicts the run
+        await handle.terminate()
+        await asyncio.wait_for(cancel_seen.wait(), 20)
+        await asyncio.wait_for(worker.shutdown(), 20)
+        await run_task
+    assert (
+        details
+        == [temporalio.activity.ActivityCancellationDetails(cancel_requested=True)] * 2
+    )
+
+
 @workflow.defn
 class SimpleChildWorkflow:
     @workflow.run