diff --git a/changelog/fixed/patient-porcupine.md b/changelog/fixed/patient-porcupine.md new file mode 100644 index 0000000000..2fa2972c63 --- /dev/null +++ b/changelog/fixed/patient-porcupine.md @@ -0,0 +1 @@ +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. diff --git a/temporalio/worker/_worker.py b/temporalio/worker/_worker.py index 3c6effd80e..3c4c479962 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,12 @@ 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. + WARNING: This setting is experimental. 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 +386,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, @@ -488,6 +496,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 +969,9 @@ 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, 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] self._async_context_run_task = asyncio.create_task(run()) return self @@ -975,8 +986,21 @@ 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() + # 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 + # A plugin's `configure_worker` may drop this key, so default here. + 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], 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 @@ -1016,6 +1040,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 new file mode 100644 index 0000000000..49b89ba292 --- /dev/null +++ b/tests/worker/test_worker_exit.py @@ -0,0 +1,80 @@ +import asyncio +import uuid +from collections.abc import Awaitable, Callable +from datetime import timedelta + +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: + # Plugin cleanup code calls the server after the Worker stops. + 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 can still await after the block. + await client.count_workflows() + + +class StuckCleanupPlugin(SimplePlugin): + def __init__(self) -> None: + super().__init__("test-stuck-cleanup") + self.cancelled = asyncio.Event() + + 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.set() + 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 + # The cancel reaches the cleanup code on a later loop iteration, so wait for it. + await asyncio.wait_for(plugin.cancelled.wait(), timeout=10)