diff --git a/CHANGELOG.md b/CHANGELOG.md index 7769f776d..38adec686 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/temporalio/worker/_activity.py b/temporalio/worker/_activity.py index ded3047fc..f88dd24f1 100644 --- a/temporalio/worker/_activity.py +++ b/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) diff --git a/tests/worker/test_workflow.py b/tests/worker/test_workflow.py index c4fe53271..497ef1b62 100644 --- a/tests/worker/test_workflow.py +++ b/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