diff --git a/CHANGELOG.md b/CHANGELOG.md index 14a225277..ae7f3a930 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -52,6 +52,8 @@ to include examples, links to docs, or any other relevant information. ### Fixed +- Safely evict cached workflows during worker shutdown so asynchronous cleanup finishes before the workflow executor stops. + - **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/_worker.py b/temporalio/worker/_worker.py index 60f824c4d..bfbf189da 100644 --- a/temporalio/worker/_worker.py +++ b/temporalio/worker/_worker.py @@ -922,6 +922,9 @@ async def shutdown(self) -> None: shut down as it runs. This will not return until the worker has completed shutting down. + Cached workflows are safely evicted before the workflow executor stops, + unless safe workflow eviction was explicitly disabled. This only closes + local workflow tasks; it does not cancel workflows on the server. """ self._shutdown_event.set() await self._shutdown_complete_event.wait() diff --git a/temporalio/worker/_workflow.py b/temporalio/worker/_workflow.py index 1b217b4a5..365e882cc 100644 --- a/temporalio/worker/_workflow.py +++ b/temporalio/worker/_workflow.py @@ -251,6 +251,27 @@ async def run(self) -> None: ] if our_tasks: await asyncio.wait(our_tasks) + # Core does not emit eviction jobs for workflows left in the cache + # when polling stops. Close them on their owning loop before GC can + # resume their cleanup on an unrelated thread or event loop. + evictions = [] + for run_id in list(self._running_workflows): + job = temporalio.bridge.proto.workflow_activation.RemoveFromCache( + message="Worker shutdown", + reason=temporalio.bridge.proto.workflow_activation.RemoveFromCache.LANG_REQUESTED, + ) + activation = temporalio.bridge.proto.workflow_activation.WorkflowActivation( + run_id=run_id, + jobs=[ + temporalio.bridge.proto.workflow_activation.WorkflowActivationJob( + remove_from_cache=job + ) + ], + ) + evictions.append( + self._handle_cache_eviction(activation, job, report_to_core=False) + ) + await asyncio.gather(*evictions) # Shutdown the thread pool executor if we created it if not self._workflow_task_executor_user_provided: self._workflow_task_executor.shutdown() @@ -536,6 +557,8 @@ async def _handle_cache_eviction( self, act: temporalio.bridge.proto.workflow_activation.WorkflowActivation, job: temporalio.bridge.proto.workflow_activation.RemoveFromCache, + *, + report_to_core: bool = True, ) -> None: logger.debug( "Evicting workflow with run ID %s, message: %s", act.run_id, job.message @@ -634,18 +657,19 @@ async def _handle_cache_eviction( # Remove from map and send completion if act.run_id in self._running_workflows: del self._running_workflows[act.run_id] - try: - await self._bridge_worker().complete_workflow_activation( - temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion( - run_id=act.run_id, - successful=temporalio.bridge.proto.workflow_completion.Success(), + if report_to_core: + try: + await self._bridge_worker().complete_workflow_activation( + temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion( + run_id=act.run_id, + successful=temporalio.bridge.proto.workflow_completion.Success(), + ) + ) + except Exception: + logger.exception( + "Failed completing eviction activation on workflow with run ID %s", + act.run_id, ) - ) - except Exception: - logger.exception( - "Failed completing eviction activation on workflow with run ID %s", - act.run_id, - ) # Run eviction hook if present if self._on_eviction_hook is not None: diff --git a/tests/worker/test_workflow.py b/tests/worker/test_workflow.py index c4fe53271..813c123b1 100644 --- a/tests/worker/test_workflow.py +++ b/tests/worker/test_workflow.py @@ -10127,3 +10127,52 @@ async def test_workflow_cancel_no_shielded_future_log( assert not any( "exception in shielded future" in record.message for record in caplog.records ) + + +shutdown_cleanup_results: dict[str, bool] = {} + + +@workflow.defn(sandboxed=False) +class ShutdownAsyncCleanupWorkflow: + @workflow.run + async def run(self) -> None: + try: + await workflow.wait_condition(lambda: False) + finally: + await asyncio.sleep(0) + shutdown_cleanup_results[workflow.info().workflow_id] = ( + workflow.unsafe.is_replaying() + ) + + @workflow.query + def ready(self) -> bool: + return True + + +@pytest.mark.parametrize("user_executor", [False, True]) +async def test_worker_shutdown_cleans_cached_workflows( + client: Client, user_executor: bool +): + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + async with new_worker( + client, + ShutdownAsyncCleanupWorkflow, + workflow_task_executor=executor if user_executor else None, + ) as worker: + handles = [ + await client.start_workflow( + ShutdownAsyncCleanupWorkflow.run, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + for _ in range(3) + ] + for handle in handles: + assert await handle.query(ShutdownAsyncCleanupWorkflow.ready) + assert handle.id not in shutdown_cleanup_results + + for handle in handles: + assert shutdown_cleanup_results.pop(handle.id) + assert (await handle.describe()).status == WorkflowExecutionStatus.RUNNING + if user_executor: + assert executor.submit(lambda: True).result()