Fix workflow eviction after activation timeouts

changelog/fixed/patient-puffin.md

@@ -0,0 +1 @@
+Preserve workflow activation and eviction executor futures across timeouts so safe eviction waits for running activations and does not repeat successful cleanup.

temporalio/worker/_workflow.py

@@ -9,7 +9,7 @@
 import os
 import sys
 import threading
-from collections.abc import Awaitable, Callable, MutableMapping, Sequence
+from collections.abc import Callable, MutableMapping, Sequence
 from dataclasses import dataclass
 from datetime import timezone
 from types import TracebackType
@@ -411,12 +411,16 @@ async def _handle_activation(
                     act,
                 )
 
-                # Run activation task with deadlock timeout
+                # Cancelling the executor future cannot stop its thread. Keep
+                # it intact so eviction can wait for the actual completion.
                 try:
-                    completion = await asyncio.wait_for(
-                        activate_task, self._deadlock_timeout_seconds
+                    await asyncio.wait(
+                        [activate_task], timeout=self._deadlock_timeout_seconds
                     )
-                except asyncio.TimeoutError:
+                except asyncio.CancelledError:
+                    workflow.deadlocked_activation_task = activate_task
+                    raise
+                if not activate_task.done():
                     # Need to create the deadlock exception up here so it
                     # captures the trace now instead of later after we may have
                     # interrupted it
@@ -432,6 +436,7 @@ async def _handle_activation(
                     # Set the task and raise
                     workflow.deadlocked_activation_task = activate_task
                     raise deadlock_exc from None
+                completion = activate_task.result()
 
         except Exception as err:
             if isinstance(err, _DeadlockError):
@@ -559,9 +564,12 @@ async def _handle_cache_eviction(
                 logger.debug(
                     "Waiting for deadlocked task to complete on run %s", act.run_id
                 )
+                # Cancellation of this handler must not cancel the executor
+                # future or let eviction overlap the still-running activation.
+                await asyncio.wait([workflow.deadlocked_activation_task])
                 try:
-                    await workflow.deadlocked_activation_task
-                except:
+                    workflow.deadlocked_activation_task.result()
+                except BaseException:
                     pass
 
             # Process the activation to evict. It is very important that
@@ -575,28 +583,35 @@ async def _handle_cache_eviction(
             seen_fail = False
             handle_eviction_task: asyncio.Future | None = None
             while True:
+                if not self._debug_mode:
+                    if handle_eviction_task is None:
+                        handle_eviction_task = (
+                            asyncio.get_running_loop().run_in_executor(
+                                self._workflow_task_executor,
+                                workflow.activate,
+                                act,
+                            )
+                        )
+                    # Keep the same execution across timeouts and retry delays.
+                    # Wait outside the retry handler so caller cancellation
+                    # propagates without cancelling the executor future.
+                    await asyncio.wait(
+                        [handle_eviction_task], timeout=self._deadlock_timeout_seconds
+                    )
                 try:
                     if self._debug_mode:
                         await self._activate_inline_for_debug(
                             asyncio.get_running_loop(), workflow, act
                         )
                     else:
-                        # We only create the eviction task if we haven't already or
-                        # it is done. This is because if it already is running and
-                        # timed out, it's still running (and holding on to a
-                        # thread). But if did complete running but failed with
-                        # another error, we want to re-create the task.
-                        if not handle_eviction_task or handle_eviction_task.done():
-                            handle_eviction_task = (
-                                asyncio.get_running_loop().run_in_executor(
-                                    self._workflow_task_executor,
-                                    workflow.activate,
-                                    act,
-                                )
-                            )
-                        await asyncio.wait_for(
-                            handle_eviction_task, self._deadlock_timeout_seconds
-                        )
+                        assert handle_eviction_task is not None
+                        if not handle_eviction_task.done():
+                            raise asyncio.TimeoutError()
+                        # Consume the result before deciding whether to retry:
+                        # completion during backoff can also be a success.
+                        completed_eviction_task = handle_eviction_task
+                        handle_eviction_task = None
+                        completed_eviction_task.result()
                     # Break if it succeeds
                     break
                 except BaseException as err:
@@ -852,7 +867,7 @@ def __init__(
     ):
         self.instance = instance
         self.workflow_id = workflow_id
-        self.deadlocked_activation_task: Awaitable | None = None
+        self.deadlocked_activation_task: asyncio.Future | None = None
         self._deadlock_can_be_interrupted_lock = threading.Lock()
         self._deadlock_can_be_interrupted = False
 

tests/worker/test_workflow.py

@@ -8938,19 +8938,30 @@ async def test_workflow_deadlock_uninterruptible(client: Client):
             id=f"workflow-{uuid.uuid4()}",
             task_queue=worker.task_queue,
         )
-        # Wait for task fail
-        await assert_task_fail_eventually(handle, message_contains="deadlock")
-        # Confirm could not be interrupted
-        assert deadlock_uninterruptible_completed == 0
-
-        # Now complete the event and confirm the workflow does complete
-        deadlock_uninterruptible_event.set()
-
-        async def check_completed():
-            assert deadlock_uninterruptible_completed >= 1
-
-        await assert_eventually(check_completed)
-        completed_sec = time.monotonic()
+        try:
+            # Wait for task fail
+            await assert_task_fail_eventually(handle, message_contains="deadlock")
+            # The thread is still blocked, so eviction must retain the instance
+            # and a completion future that cannot be cancelled by the timeout.
+            assert deadlock_uninterruptible_completed == 0
+            workflow_worker = worker._workflow_worker
+            assert workflow_worker is not None
+            assert handle.result_run_id is not None
+            running = workflow_worker._running_workflows[handle.result_run_id]
+            future = running.deadlocked_activation_task
+            assert isinstance(future, asyncio.Future)
+            assert not future.done()
+
+            # Now complete the event and confirm the workflow does complete
+            deadlock_uninterruptible_event.set()
+
+            async def check_completed():
+                assert deadlock_uninterruptible_completed >= 1
+
+            await assert_eventually(check_completed)
+            completed_sec = time.monotonic()
+        finally:
+            deadlock_uninterruptible_event.set()
     # Confirm worker shutdown didn't hang
     assert time.monotonic() - completed_sec < 20
 

tests/worker/test_workflow_eviction.py

@@ -0,0 +1,176 @@
+import asyncio
+import threading
+import uuid
+from collections.abc import Sequence
+
+import pytest
+
+import temporalio.converter
+import temporalio.converter._extstore
+import temporalio.workflow
+from temporalio.bridge.proto.workflow_activation import WorkflowActivation
+from temporalio.bridge.proto.workflow_completion import WorkflowActivationCompletion
+from temporalio.client import Client
+from temporalio.worker import WorkflowInstance, WorkflowInstanceDetails, WorkflowRunner
+from temporalio.worker._command_aware_visitor import CommandInfo
+from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner
+from tests.helpers import assert_eventually, assert_task_fail_eventually, new_worker
+
+
+@temporalio.workflow.defn
+class EvictionCompletionWorkflow:
+    @temporalio.workflow.run
+    async def run(self) -> str:
+        return "done"
+
+
+class _BlockingRunner(WorkflowRunner):
+    def __init__(self, *, block_eviction: bool) -> None:
+        self._runner = SandboxedWorkflowRunner()
+        self.block_eviction = block_eviction
+        self.started = threading.Event()
+        self.release = threading.Event()
+        self.instances: list[_BlockingInstance] = []
+
+    def prepare_workflow(self, defn: temporalio.workflow._Definition) -> None:
+        self._runner.prepare_workflow(defn)
+
+    def create_instance(self, det: WorkflowInstanceDetails) -> WorkflowInstance:
+        instance = _BlockingInstance(self, self._runner.create_instance(det))
+        self.instances.append(instance)
+        return instance
+
+    def set_worker_level_failure_exception_types(
+        self, types: Sequence[type[BaseException]]
+    ) -> None:
+        self._runner.set_worker_level_failure_exception_types(types)
+
+
+class _BlockingInstance(WorkflowInstance):
+    def __init__(self, runner: _BlockingRunner, instance: WorkflowInstance) -> None:
+        self._runner = runner
+        self._instance = instance
+        self._lock = threading.Lock()
+        self._active = 0
+        self.max_active = 0
+        self.eviction_calls = 0
+
+    def activate(self, act: WorkflowActivation) -> WorkflowActivationCompletion:
+        evicting = any(job.HasField("remove_from_cache") for job in act.jobs)
+        with self._lock:
+            self._active += 1
+            self.max_active = max(self.max_active, self._active)
+            if evicting:
+                self.eviction_calls += 1
+            should_block = (
+                evicting == self._runner.block_eviction
+                and not self._runner.started.is_set()
+            )
+            if should_block:
+                self._runner.started.set()
+        try:
+            # Hold the executor thread outside the sandbox so the deadlock
+            # interruption cannot complete it before the test releases it.
+            if should_block:
+                assert self._runner.release.wait(timeout=15)
+            return self._instance.activate(act)
+        finally:
+            with self._lock:
+                self._active -= 1
+
+    def get_serialization_context(
+        self, command_info: CommandInfo | None
+    ) -> temporalio.converter.SerializationContext | None:
+        return self._instance.get_serialization_context(command_info)
+
+    def get_external_store_context(
+        self, command_info: CommandInfo | None
+    ) -> temporalio.converter._extstore.StorageDriverStoreContext:
+        return self._instance.get_external_store_context(command_info)
+
+    def get_info(self) -> temporalio.workflow.Info:
+        return self._instance.get_info()
+
+    def get_thread_id(self) -> int | None:
+        return self._instance.get_thread_id()
+
+
+async def test_workflow_eviction_waits_for_timed_out_activation(client: Client) -> None:
+    runner = _BlockingRunner(block_eviction=False)
+    async with new_worker(
+        client,
+        EvictionCompletionWorkflow,
+        workflow_runner=runner,
+        max_concurrent_workflow_tasks=2,
+    ) as worker:
+        workflow_worker = worker._workflow_worker
+        assert workflow_worker is not None
+        workflow_worker._deadlock_timeout_seconds = 1
+        try:
+            handle = await client.start_workflow(
+                EvictionCompletionWorkflow.run,
+                id=f"workflow-{uuid.uuid4()}",
+                task_queue=worker.task_queue,
+            )
+            await assert_task_fail_eventually(handle, message_contains="deadlock")
+            assert runner.started.is_set()
+            instance = runner.instances[0]
+            assert instance.eviction_calls == 0
+            assert handle.result_run_id is not None
+            running = workflow_worker._running_workflows[handle.result_run_id]
+            assert isinstance(running.deadlocked_activation_task, asyncio.Future)
+            assert not running.deadlocked_activation_task.done()
+            runner.release.set()
+            assert await asyncio.wait_for(handle.result(), timeout=10) == "done"
+
+            async def check_evicted() -> None:
+                assert instance.eviction_calls == 1
+
+            await assert_eventually(check_evicted)
+            assert instance.max_active == 1
+        finally:
+            runner.release.set()
+
+
+@pytest.mark.parametrize("hold_across_retry", [False, True])
+async def test_workflow_eviction_waits_for_cleanup_completion(
+    client: Client, hold_across_retry: bool
+) -> None:
+    runner = _BlockingRunner(block_eviction=True)
+    async with new_worker(
+        client,
+        EvictionCompletionWorkflow,
+        workflow_runner=runner,
+        max_concurrent_workflow_tasks=2,
+    ) as worker:
+        workflow_worker = worker._workflow_worker
+        assert workflow_worker is not None
+        workflow_worker._deadlock_timeout_seconds = 1
+        evicted = asyncio.Event()
+        workflow_worker._on_eviction_hook = lambda _run_id, _job: evicted.set()
+        try:
+            assert (
+                await client.execute_workflow(
+                    EvictionCompletionWorkflow.run,
+                    id=f"workflow-{uuid.uuid4()}",
+                    task_queue=worker.task_queue,
+                )
+                == "done"
+            )
+
+            async def check_timed_out() -> None:
+                assert workflow_worker._could_not_evict_count == 1
+
+            await assert_eventually(check_timed_out)
+            if hold_across_retry:
+                # The retry delay is two seconds. Keeping the same execution
+                # blocked past it detects an overlapping eviction submission.
+                await asyncio.sleep(2.4)
+            assert not evicted.is_set()
+            runner.release.set()
+            await asyncio.wait_for(evicted.wait(), timeout=10)
+            assert runner.instances[0].eviction_calls == 1
+            assert runner.instances[0].max_active == 1
+            assert workflow_worker._could_not_evict_count == 0
+        finally:
+            runner.release.set()