fix(worker): evict cached workflows during shutdown

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.

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()

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:

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()