From c7ad8023ccb37637c3d74e4e210e215d2960d3dc Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 20 Oct 2026 21:37:20 -0700 Subject: [PATCH 1/7] Awaited the Worker run task on exit instead of cancelling it. __aexit__ cancelled the run task while a plugin's run_worker cleanup still awaited, and the run wrapper then cancelled the caller's task, so every later await raised CancelledError. Now __aexit__ waits, with a bound, for the plugins' code after shutdown. --- temporalio/worker/_worker.py | 19 +++++++++++--- tests/worker/test_worker_exit.py | 44 ++++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+), 3 deletions(-) create mode 100644 tests/worker/test_worker_exit.py diff --git a/temporalio/worker/_worker.py b/temporalio/worker/_worker.py index 3c6effd80e..1092d4add9 100644 --- a/temporalio/worker/_worker.py +++ b/temporalio/worker/_worker.py @@ -488,6 +488,7 @@ def _init_from_config(self, client: temporalio.client.Client, config: WorkerConf self._async_context_inner_task: asyncio.Task | None = None self._async_context_run_task: asyncio.Task | None = None self._async_context_run_exception: BaseException | None = None + self._async_context_exiting = False self._activity_worker: _ActivityWorker | None = None self._runtime = ( @@ -960,7 +961,10 @@ async def run(): await self.run() except BaseException as err: self._async_context_run_exception = err - self._async_context_inner_task.cancel() # type: ignore[union-attr] + # After __aexit__ starts, the caller is past the block, so a cancel would hit + # its later awaits. + if not self._async_context_exiting: + self._async_context_inner_task.cancel() # type: ignore[union-attr] self._async_context_run_task = asyncio.create_task(run()) return self @@ -975,13 +979,22 @@ async def __aexit__(self, exc_type: type[BaseException] | None, *args: Any) -> N if not self._async_context_run_task: raise RuntimeError("Never started") await self.shutdown() - # Cancel our run task - self._async_context_run_task.cancel() + # The run is done after shutdown, except for the code that plugins run after it. Let + # that code finish, so a plugin can still await calls in its cleanup. + self._async_context_exiting = True + run_task = self._async_context_run_task + await asyncio.wait([run_task], timeout=_RUN_TASK_EXIT_TIMEOUT.total_seconds()) + if not run_task.done(): + run_task.cancel() # Only re-raise our exception if present and exc_type is cancel if exc_type is asyncio.CancelledError and self._async_context_run_exception: raise self._async_context_run_exception +# How long __aexit__ waits for the plugins' code after the Worker shut down. +_RUN_TASK_EXIT_TIMEOUT = timedelta(seconds=60) + + class WorkerConfig(TypedDict, total=False): """TypedDict of config originally passed to :py:class:`Worker`.""" diff --git a/tests/worker/test_worker_exit.py b/tests/worker/test_worker_exit.py new file mode 100644 index 0000000000..8950b2e8e1 --- /dev/null +++ b/tests/worker/test_worker_exit.py @@ -0,0 +1,44 @@ +import uuid +from collections.abc import Awaitable, Callable + +from temporalio import workflow +from temporalio.client import Client +from temporalio.plugin import SimplePlugin +from temporalio.worker import Worker + + +@workflow.defn +class NothingWorkflow: + @workflow.run + async def run(self) -> None: + pass + + +class CleanupPlugin(SimplePlugin): + def __init__(self, client: Client) -> None: + super().__init__("test-cleanup") + self.client = client + self.cleanup_result: int | None = None + + async def run_worker( + self, worker: Worker, next: Callable[[Worker], Awaitable[None]] + ) -> None: + try: + await super().run_worker(worker, next) + finally: + # A call after the Worker stops, as a plugin's cleanup makes. + self.cleanup_result = (await self.client.count_workflows()).count + + +async def test_plugin_cleanup_awaits_after_worker_exit(client: Client): + plugin = CleanupPlugin(client) + async with Worker( + client, + task_queue=f"tq-{uuid.uuid4()}", + workflows=[NothingWorkflow], + plugins=[plugin], + ): + pass + assert plugin.cleanup_result is not None + # The caller's own awaits after the block must not be cancelled. + await client.count_workflows() From 1b741445507663234ca59a1911602c6654cb18f6 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Wed, 21 Oct 2026 10:32:49 -0700 Subject: [PATCH 2/7] Made the Worker exit comments literal. --- temporalio/worker/_worker.py | 8 +++----- tests/worker/test_worker_exit.py | 4 ++-- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/temporalio/worker/_worker.py b/temporalio/worker/_worker.py index 1092d4add9..54b2790bc6 100644 --- a/temporalio/worker/_worker.py +++ b/temporalio/worker/_worker.py @@ -961,8 +961,7 @@ async def run(): await self.run() except BaseException as err: self._async_context_run_exception = err - # After __aexit__ starts, the caller is past the block, so a cancel would hit - # its later awaits. + # After __aexit__ starts, a cancel would stop the caller's code after the block. if not self._async_context_exiting: self._async_context_inner_task.cancel() # type: ignore[union-attr] @@ -979,8 +978,7 @@ async def __aexit__(self, exc_type: type[BaseException] | None, *args: Any) -> N if not self._async_context_run_task: raise RuntimeError("Never started") await self.shutdown() - # The run is done after shutdown, except for the code that plugins run after it. Let - # that code finish, so a plugin can still await calls in its cleanup. + # Plugins can run cleanup code after shutdown. Wait for that code instead of cancelling it. self._async_context_exiting = True run_task = self._async_context_run_task await asyncio.wait([run_task], timeout=_RUN_TASK_EXIT_TIMEOUT.total_seconds()) @@ -991,7 +989,7 @@ async def __aexit__(self, exc_type: type[BaseException] | None, *args: Any) -> N raise self._async_context_run_exception -# How long __aexit__ waits for the plugins' code after the Worker shut down. +# How long __aexit__ waits for plugin cleanup code after shutdown. _RUN_TASK_EXIT_TIMEOUT = timedelta(seconds=60) diff --git a/tests/worker/test_worker_exit.py b/tests/worker/test_worker_exit.py index 8950b2e8e1..7f2942ae55 100644 --- a/tests/worker/test_worker_exit.py +++ b/tests/worker/test_worker_exit.py @@ -26,7 +26,7 @@ async def run_worker( try: await super().run_worker(worker, next) finally: - # A call after the Worker stops, as a plugin's cleanup makes. + # Plugin cleanup code calls the server after the Worker stops. self.cleanup_result = (await self.client.count_workflows()).count @@ -40,5 +40,5 @@ async def test_plugin_cleanup_awaits_after_worker_exit(client: Client): ): pass assert plugin.cleanup_result is not None - # The caller's own awaits after the block must not be cancelled. + # The caller can still await after the block. await client.count_workflows() From d194df78de3ea026988d4d018ee9d45f9d0831ca Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Wed, 21 Oct 2026 22:24:12 -0700 Subject: [PATCH 3/7] Added a changelog entry for the Worker exit fix. --- changelog/fixed/patient-porcupine.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changelog/fixed/patient-porcupine.md diff --git a/changelog/fixed/patient-porcupine.md b/changelog/fixed/patient-porcupine.md new file mode 100644 index 0000000000..04a1786884 --- /dev/null +++ b/changelog/fixed/patient-porcupine.md @@ -0,0 +1 @@ +Exiting `async with Worker(...)` now waits for the Worker to finish its shutdown, including the `run_worker` cleanup of plugins, instead of cancelling the task that entered the block. From 7c17d7a4d998b4617c67525495c64cd44331b071 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Wed, 21 Oct 2026 23:22:42 -0700 Subject: [PATCH 4/7] Waited for the Worker run task with no timeout on exit. --- temporalio/worker/_worker.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/temporalio/worker/_worker.py b/temporalio/worker/_worker.py index 54b2790bc6..4e8485daec 100644 --- a/temporalio/worker/_worker.py +++ b/temporalio/worker/_worker.py @@ -981,18 +981,17 @@ async def __aexit__(self, exc_type: type[BaseException] | None, *args: Any) -> N # Plugins can run cleanup code after shutdown. Wait for that code instead of cancelling it. self._async_context_exiting = True run_task = self._async_context_run_task - await asyncio.wait([run_task], timeout=_RUN_TASK_EXIT_TIMEOUT.total_seconds()) - if not run_task.done(): + try: + # wait() doesn't raise for the task's own outcome, only when this call is cancelled. + await asyncio.wait([run_task]) + except asyncio.CancelledError: run_task.cancel() + raise # Only re-raise our exception if present and exc_type is cancel if exc_type is asyncio.CancelledError and self._async_context_run_exception: raise self._async_context_run_exception -# How long __aexit__ waits for plugin cleanup code after shutdown. -_RUN_TASK_EXIT_TIMEOUT = timedelta(seconds=60) - - class WorkerConfig(TypedDict, total=False): """TypedDict of config originally passed to :py:class:`Worker`.""" From 1688d3cb5f4e81b2e93b0952d065b98ef6f7125d Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Wed, 21 Oct 2026 23:27:20 -0700 Subject: [PATCH 5/7] Made the wait for plugin cleanup a Worker option. --- temporalio/worker/_worker.py | 15 ++++++++++++- tests/worker/test_worker_exit.py | 37 ++++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/temporalio/worker/_worker.py b/temporalio/worker/_worker.py index 4e8485daec..0df790ab82 100644 --- a/temporalio/worker/_worker.py +++ b/temporalio/worker/_worker.py @@ -132,6 +132,7 @@ def __init__( max_task_queue_activities_per_second: float | None = None, max_eager_activity_reservations_per_workflow_task: int = 3, graceful_shutdown_timeout: timedelta = timedelta(), + plugin_shutdown_timeout: timedelta | None = timedelta(seconds=60), workflow_failure_exception_types: Sequence[type[BaseException]] = [], shared_state_manager: SharedStateManager | None = None, debug_mode: bool = False, @@ -278,6 +279,11 @@ def __init__( graceful_shutdown_timeout: Amount of time after shutdown is called that activities are given to complete before their tasks are cancelled. + plugin_shutdown_timeout: Amount of time that exiting + ``async with Worker(...)`` waits, after shutdown, for the code + that plugins run after the Worker stops in ``run_worker``. That + code is cancelled after this time. ``None`` waits with no limit. + The default is 60 seconds. workflow_failure_exception_types: The types of exceptions that, if a workflow-thrown exception extends, will cause the workflow/update to fail instead of suspending the workflow via @@ -379,6 +385,7 @@ def __init__( max_task_queue_activities_per_second=max_task_queue_activities_per_second, max_eager_activity_reservations_per_workflow_task=max_eager_activity_reservations_per_workflow_task, graceful_shutdown_timeout=graceful_shutdown_timeout, + plugin_shutdown_timeout=plugin_shutdown_timeout, workflow_failure_exception_types=workflow_failure_exception_types, shared_state_manager=shared_state_manager, debug_mode=debug_mode, @@ -981,12 +988,17 @@ async def __aexit__(self, exc_type: type[BaseException] | None, *args: Any) -> N # Plugins can run cleanup code after shutdown. Wait for that code instead of cancelling it. self._async_context_exiting = True run_task = self._async_context_run_task + timeout = self._config.get("plugin_shutdown_timeout", timedelta(seconds=60)) try: # wait() doesn't raise for the task's own outcome, only when this call is cancelled. - await asyncio.wait([run_task]) + await asyncio.wait( + [run_task], timeout=timeout.total_seconds() if timeout else None + ) except asyncio.CancelledError: run_task.cancel() raise + if not run_task.done(): + run_task.cancel() # Only re-raise our exception if present and exc_type is cancel if exc_type is asyncio.CancelledError and self._async_context_run_exception: raise self._async_context_run_exception @@ -1026,6 +1038,7 @@ class WorkerConfig(TypedDict, total=False): max_task_queue_activities_per_second: float | None max_eager_activity_reservations_per_workflow_task: int graceful_shutdown_timeout: timedelta + plugin_shutdown_timeout: timedelta | None workflow_failure_exception_types: Sequence[type[BaseException]] shared_state_manager: SharedStateManager | None debug_mode: bool diff --git a/tests/worker/test_worker_exit.py b/tests/worker/test_worker_exit.py index 7f2942ae55..05339548f5 100644 --- a/tests/worker/test_worker_exit.py +++ b/tests/worker/test_worker_exit.py @@ -1,5 +1,7 @@ +import asyncio import uuid from collections.abc import Awaitable, Callable +from datetime import timedelta from temporalio import workflow from temporalio.client import Client @@ -42,3 +44,38 @@ async def test_plugin_cleanup_awaits_after_worker_exit(client: Client): assert plugin.cleanup_result is not None # The caller can still await after the block. await client.count_workflows() + + +class StuckCleanupPlugin(SimplePlugin): + def __init__(self) -> None: + super().__init__("test-stuck-cleanup") + self.cancelled = False + + async def run_worker( + self, worker: Worker, next: Callable[[Worker], Awaitable[None]] + ) -> None: + try: + await super().run_worker(worker, next) + finally: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + self.cancelled = True + raise + + +async def test_plugin_cleanup_is_cancelled_after_plugin_shutdown_timeout( + client: Client, +): + plugin = StuckCleanupPlugin() + async with Worker( + client, + task_queue=f"tq-{uuid.uuid4()}", + workflows=[NothingWorkflow], + plugins=[plugin], + plugin_shutdown_timeout=timedelta(milliseconds=100), + ): + pass + # Let the cancelled run task finish its cleanup. + await asyncio.sleep(0) + assert plugin.cancelled From 666ec4405ba00efb36a5561c5a40309322a0345f Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Wed, 21 Oct 2026 23:27:29 -0700 Subject: [PATCH 6/7] Named plugin_shutdown_timeout in the changelog entry. --- changelog/fixed/patient-porcupine.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/changelog/fixed/patient-porcupine.md b/changelog/fixed/patient-porcupine.md index 04a1786884..2fa2972c63 100644 --- a/changelog/fixed/patient-porcupine.md +++ b/changelog/fixed/patient-porcupine.md @@ -1 +1 @@ -Exiting `async with Worker(...)` now waits for the Worker to finish its shutdown, including the `run_worker` cleanup of plugins, instead of cancelling the task that entered the block. +Exiting `async with Worker(...)` now waits for the code that plugins run after the Worker stops in `run_worker`, up to the new `plugin_shutdown_timeout` Worker option, instead of cancelling the task that entered the block. From 70d8721864e0f0a5c958ef150cfcae780880842e Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Wed, 21 Oct 2026 23:30:11 -0700 Subject: [PATCH 7/7] Waited for the cleanup cancel in the exit test instead of one yield. --- tests/worker/test_worker_exit.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/tests/worker/test_worker_exit.py b/tests/worker/test_worker_exit.py index 05339548f5..49b89ba292 100644 --- a/tests/worker/test_worker_exit.py +++ b/tests/worker/test_worker_exit.py @@ -49,7 +49,7 @@ async def test_plugin_cleanup_awaits_after_worker_exit(client: Client): class StuckCleanupPlugin(SimplePlugin): def __init__(self) -> None: super().__init__("test-stuck-cleanup") - self.cancelled = False + self.cancelled = asyncio.Event() async def run_worker( self, worker: Worker, next: Callable[[Worker], Awaitable[None]] @@ -60,7 +60,7 @@ async def run_worker( try: await asyncio.Event().wait() except asyncio.CancelledError: - self.cancelled = True + self.cancelled.set() raise @@ -76,6 +76,5 @@ async def test_plugin_cleanup_is_cancelled_after_plugin_shutdown_timeout( plugin_shutdown_timeout=timedelta(milliseconds=100), ): pass - # Let the cancelled run task finish its cleanup. - await asyncio.sleep(0) - assert plugin.cancelled + # The cancel reaches the cleanup code on a later loop iteration, so wait for it. + await asyncio.wait_for(plugin.cancelled.wait(), timeout=10)