From aa3982c77d8216e7ceab321ba74c604863d828e9 Mon Sep 17 00:00:00 2001 From: Nijanthan Hariharan Date: Thu, 22 Oct 2026 12:38:32 -0700 Subject: [PATCH] Fix workflow eviction after activation timeouts --- changelog/fixed/patient-puffin.md | 1 + temporalio/worker/_workflow.py | 63 +++--- tests/worker/test_workflow.py | 37 ++-- tests/worker/test_workflow_eviction.py | 176 +++++++++++++++++ tests/worker/test_workflow_eviction_unit.py | 209 ++++++++++++++++++++ 5 files changed, 449 insertions(+), 37 deletions(-) create mode 100644 changelog/fixed/patient-puffin.md create mode 100644 tests/worker/test_workflow_eviction.py create mode 100644 tests/worker/test_workflow_eviction_unit.py diff --git a/changelog/fixed/patient-puffin.md b/changelog/fixed/patient-puffin.md new file mode 100644 index 0000000000..baa52ebe70 --- /dev/null +++ b/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. diff --git a/temporalio/worker/_workflow.py b/temporalio/worker/_workflow.py index f50955c4cf..fa1263fd92 100644 --- a/temporalio/worker/_workflow.py +++ b/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 diff --git a/tests/worker/test_workflow.py b/tests/worker/test_workflow.py index 4186bb6676..ea07e1d5a8 100644 --- a/tests/worker/test_workflow.py +++ b/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 diff --git a/tests/worker/test_workflow_eviction.py b/tests/worker/test_workflow_eviction.py new file mode 100644 index 0000000000..e08168816d --- /dev/null +++ b/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() diff --git a/tests/worker/test_workflow_eviction_unit.py b/tests/worker/test_workflow_eviction_unit.py new file mode 100644 index 0000000000..e734812c9e --- /dev/null +++ b/tests/worker/test_workflow_eviction_unit.py @@ -0,0 +1,209 @@ +from __future__ import annotations + +import asyncio +import concurrent.futures +import threading +from collections.abc import AsyncIterator +from unittest.mock import AsyncMock, Mock + +import pytest + +import temporalio.bridge.worker +from temporalio.bridge.proto.workflow_activation import ( + RemoveFromCache, + WorkflowActivation, +) +from temporalio.bridge.proto.workflow_completion import WorkflowActivationCompletion +from temporalio.common import MetricMeter +from temporalio.converter import DataConverter +from temporalio.worker._workflow import _RunningWorkflow, _WorkflowWorker +from temporalio.worker._workflow_instance import ( + UnsandboxedWorkflowRunner, + WorkflowInstance, +) +from tests.helpers import assert_eventually + + +class EvictionHarness: + def __init__(self) -> None: + self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=2) + self.release = threading.Event() + self.started = asyncio.Event() + self.finished = asyncio.Event() + self.loop = asyncio.get_running_loop() + self.calls = 0 + self.instance = Mock(spec=WorkflowInstance) + self.instance.activate.side_effect = self.activate + self.workflow = _RunningWorkflow(self.instance, "workflow-id") + self.bridge = Mock(spec=temporalio.bridge.worker.Worker) + self.bridge.complete_workflow_activation = AsyncMock() + self.worker = _WorkflowWorker( + bridge_worker=lambda: self.bridge, + namespace="default", + task_queue="test-eviction", + workflows=[], + workflow_task_executor=self.executor, + max_concurrent_workflow_tasks=2, + workflow_runner=UnsandboxedWorkflowRunner(), + unsandboxed_workflow_runner=UnsandboxedWorkflowRunner(), + data_converter=DataConverter.default, + interceptors=[], + workflow_failure_exception_types=[], + patch_activation_callback=None, + debug_mode=False, + disable_eager_activity_execution=False, + metric_meter=MetricMeter.noop, + on_eviction_hook=None, + disable_safe_eviction=False, + should_enforce_versioning_behavior=False, + assert_local_activity_valid=lambda _: None, + encode_headers=False, + max_workflow_task_external_storage_concurrency=1, + ) + self.worker._deadlock_timeout_seconds = 1 + self.worker._running_workflows["run-id"] = self.workflow + self.job = RemoveFromCache(message="test") + self.activation = WorkflowActivation(run_id="run-id") + self.activation.jobs.add().remove_from_cache.CopyFrom(self.job) + self.tasks: list[asyncio.Task[None]] = [] + + def activate(self, act: WorkflowActivation) -> WorkflowActivationCompletion: + self.calls += 1 + self.loop.call_soon_threadsafe(self.started.set) + self.release.wait() + self.loop.call_soon_threadsafe(self.finished.set) + return WorkflowActivationCompletion(run_id=act.run_id) + + def evict(self) -> asyncio.Task[None]: + task = asyncio.create_task( + self.worker._handle_cache_eviction(self.activation, self.job) + ) + self.tasks.append(task) + return task + + async def wait_for_timeout(self) -> None: + async def check() -> None: + assert self.worker._could_not_evict_count == 1 + + await assert_eventually(check) + + def assert_not_evicted(self) -> None: + assert self.worker._running_workflows["run-id"] is self.workflow + self.bridge.complete_workflow_activation.assert_not_awaited() + + +@pytest.fixture +async def eviction() -> AsyncIterator[EvictionHarness]: + harness = EvictionHarness() + try: + yield harness + finally: + harness.release.set() + await asyncio.gather(*harness.tasks, return_exceptions=True) + harness.executor.shutdown() + + +async def test_eviction_consumes_success_during_retry_delay( + eviction: EvictionHarness, +) -> None: + task = eviction.evict() + await eviction.wait_for_timeout() + eviction.assert_not_evicted() + eviction.release.set() + await eviction.finished.wait() + await task + assert eviction.calls == 1 + assert "run-id" not in eviction.worker._running_workflows + eviction.bridge.complete_workflow_activation.assert_awaited_once() + assert eviction.worker._could_not_evict_count == 0 + + +@pytest.mark.parametrize( + "error", + [ + RuntimeError("activation failed"), + asyncio.TimeoutError(), + asyncio.CancelledError(), + ], +) +async def test_eviction_retries_completed_failure( + eviction: EvictionHarness, error: BaseException +) -> None: + eviction.instance.activate.side_effect = [ + error, + WorkflowActivationCompletion(run_id="run-id"), + ] + await eviction.evict() + assert eviction.instance.activate.call_count == 2 + eviction.bridge.complete_workflow_activation.assert_awaited_once() + assert eviction.worker._could_not_evict_count == 0 + + +async def test_eviction_cancellation_waiting_for_deadlocked_activation( + eviction: EvictionHarness, +) -> None: + activation = eviction.loop.run_in_executor( + eviction.executor, eviction.workflow.activate, eviction.activation + ) + eviction.workflow.deadlocked_activation_task = activation + await eviction.started.wait() + task = eviction.evict() + await asyncio.sleep(0) + task.cancel() + # A cancelled handler must not turn the still-running thread into a completed + # activation and start cleanup on the same workflow instance. + eviction.release.set() + with pytest.raises(asyncio.CancelledError): + await task + assert not activation.cancelled() + await activation + assert eviction.calls == 1 + eviction.assert_not_evicted() + + +async def test_eviction_cancellation_waiting_for_cleanup( + eviction: EvictionHarness, +) -> None: + task = eviction.evict() + await eviction.started.wait() + task.cancel() + eviction.release.set() + with pytest.raises(asyncio.CancelledError): + await task + assert eviction.calls == 1 + eviction.assert_not_evicted() + + +async def test_activation_cancellation_preserves_completion_for_eviction( + eviction: EvictionHarness, +) -> None: + eviction.instance.get_info.return_value.workflow_type = "TestWorkflow" + task = asyncio.create_task( + eviction.worker._handle_activation(WorkflowActivation(run_id="run-id")) + ) + eviction.tasks.append(task) + await eviction.started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + activation = eviction.workflow.deadlocked_activation_task + assert activation is not None + assert not activation.done() + eviction.assert_not_evicted() + eviction.release.set() + await activation + await eviction.evict() + assert eviction.calls == 2 + eviction.bridge.complete_workflow_activation.assert_awaited_once() + + +async def test_activation_execution_timeout_is_not_a_deadlock( + eviction: EvictionHarness, +) -> None: + eviction.instance.get_info.return_value.workflow_type = "TestWorkflow" + eviction.instance.activate.side_effect = asyncio.TimeoutError("execution timeout") + await eviction.worker._handle_activation(WorkflowActivation(run_id="run-id")) + assert eviction.workflow.deadlocked_activation_task is None + eviction.bridge.complete_workflow_activation.assert_awaited_once() + completion = eviction.bridge.complete_workflow_activation.call_args.args[0] + assert completion.failed.failure.message == "execution timeout"