diff --git a/CHANGELOG.md b/CHANGELOG.md index 5a8995696..b3aaf0595 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -28,6 +28,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. ### Security ## [1.34.0] - 2026-09-30 diff --git a/temporalio/worker/_activity.py b/temporalio/worker/_activity.py index f72df9cf4..cbb504857 100644 --- a/temporalio/worker/_activity.py +++ b/temporalio/worker/_activity.py @@ -195,9 +195,33 @@ 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) + # A cancel above can land while an activity is still encoding its + # result, ending its task with CancelledError, so never let a task + # exception escape and stall shutdown + for result in await asyncio.gather(*running_tasks, return_exceptions=True): + if isinstance(result, BaseException) and not isinstance( + result, asyncio.CancelledError + ): + logger.warning( + "Activity task raised during worker shutdown", exc_info=result + ) def _handle_cancel_activity_task( self, diff --git a/tests/worker/test_workflow.py b/tests/worker/test_workflow.py index d751fdb8b..299636624 100644 --- a/tests/worker/test_workflow.py +++ b/tests/worker/test_workflow.py @@ -16,7 +16,8 @@ import typing import uuid from abc import ABC, abstractmethod -from collections.abc import Awaitable, Mapping, Sequence +from collections.abc import Awaitable, Iterator, Mapping, Sequence +from contextlib import contextmanager from dataclasses import dataclass from datetime import datetime, timedelta, timezone from enum import IntEnum @@ -27,6 +28,7 @@ NoReturn, cast, ) +from unittest.mock import patch from urllib.request import urlopen import pydantic @@ -36,11 +38,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 +57,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, @@ -1136,6 +1143,200 @@ 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), + ) + + +@contextmanager +def _orphan_local_activity_cancel() -> Iterator[None]: + """Make Core drop a running local activity's cancel on eviction. + + Holds the activity poll that follows a local activity start until workflow + polling has shut down, so the cancel Core queues on eviction is never + delivered and the activity keeps running untracked. + """ + workflow_poll_shut_down = asyncio.Event() + 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 + ), + ): + yield + + +async def test_worker_shutdown_cancels_local_activity_untracked_by_core( + client: Client, +): + started = 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()) + + with _orphan_local_activity_cancel(): + 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 + ) + + +class _SlowPayloadCodec(PayloadCodec): + """Yields while encoding, like a codec that calls a remote KMS.""" + + async def encode(self, payloads: Sequence[Payload]) -> list[Payload]: + if payloads: + await asyncio.sleep(5) + return list(payloads) + + async def decode(self, payloads: Sequence[Payload]) -> list[Payload]: + return list(payloads) + + +async def test_worker_shutdown_completes_while_untracked_activity_encodes_failure( + client: Client, +): + # The shutdown cancel lands while the orphaned activity is still encoding its + # failure through a slow codec; shutdown must still complete + started = asyncio.Event() + + @activity.defn(name="wait_forever_local_activity") + async def raise_on_shutdown_local_activity() -> None: + started.set() + await activity.wait_for_worker_shutdown() + # Details send the failure through the codec + raise ApplicationError("failing at shutdown", {"detail": "x"}) + + config = client.config() + config["data_converter"] = dataclasses.replace( + DataConverter.default, payload_codec=_SlowPayloadCodec() + ) + codec_client = Client(**config) + with _orphan_local_activity_cancel(): + worker = new_worker( + codec_client, + OrphanedLocalActivityWorkflow, + activities=[raise_on_shutdown_local_activity], + ) + run_task = asyncio.create_task(worker.run()) + handle = await codec_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 + + @workflow.defn class SimpleChildWorkflow: @workflow.run