diff --git a/src/blueapi/worker/task_worker.py b/src/blueapi/worker/task_worker.py index caa39fe7a..713813aa6 100644 --- a/src/blueapi/worker/task_worker.py +++ b/src/blueapi/worker/task_worker.py @@ -8,6 +8,7 @@ from threading import Event, RLock from typing import Any, TypeVar +from bluesky import RunEngineInterrupted from bluesky.protocols import Status from observability_utils.tracing import ( add_span_attributes, @@ -109,6 +110,7 @@ class TaskWorker: _task_channel: Queue # type: ignore _current: TrackableTask | None + _pending_cancel: "CancelSignal | None" _status_lock: RLock _status_snapshot: dict[str, StatusView] _completed_statuses: set[str] @@ -139,6 +141,7 @@ def __init__( self._warnings = [] self._task_channel = Queue(maxsize=1) self._current = None + self._pending_cancel = None self._worker_events = EventPublisher() self._progress_events = EventPublisher() self._data_events = EventPublisher() @@ -180,19 +183,43 @@ def cancel_active_task( Returns: The task_id of the active task """ - if self._current is None: + current = self._current + if current is None: # Persuades type checker that self._current is not None # We only allow this method to be called if a Plan is active raise TransitionError("Attempted to cancel while no active Task") + + if self._ctx.run_engine.state == "paused": + # abort()/stop() block here until cleanup finishes - defer to the + # worker thread, like ResumeSignal does for resume(). Recorded + # separately so a queued ResumeSignal can't race ahead of this + # and finish the task before the cancel is even looked at. + signal = CancelSignal(failure=failure, reason=reason) + self._pending_cancel = signal + try: + self._task_channel.put_nowait(signal) + except Full: + pass # a signal already queued will check _pending_cancel + return current.task_id + + # RE.abort()/stop() are thread-safe and must be called immediately — + # putting only a CancelSignal would no-op if the worker is blocked in do_task() + default_reason = "Task failed for unknown reason" if failure: - default_reason = "Task failed for unknown reason" self._ctx.run_engine.abort(reason or default_reason) - add_span_attributes({"Task aborted": reason or default_reason}) + with self._status_lock: + if current.outcome is None: + current.set_exception(Exception(reason or default_reason)) else: self._ctx.run_engine.stop() - default_reason = "Cancellation successful: Task stopped without error" - add_span_attributes({"Task stopped": reason or default_reason}) - return self._current.task_id + with self._status_lock: + if current.outcome is None: + current.set_result(None) + self._task_channel.put(CancelSignal(failure=failure, reason=reason)) + add_span_attributes( + {"Task aborted" if failure else "Task stopped": reason or ""} + ) + return current.task_id @start_as_current_span(TRACER) def get_tasks(self) -> list[TrackableTask]: @@ -415,7 +442,7 @@ def resume(self): Command the worker to resume """ LOGGER.info("Requesting to resume the worker") - self._ctx.run_engine.resume() + self._task_channel.put(ResumeSignal()) @start_as_current_span(TRACER) def _cycle_with_error_handling(self) -> None: @@ -441,11 +468,23 @@ def process_task(): LOGGER.info( "Task ran successfully - returned: %s", result, extra=meta ) - self._current.set_result(result) + with self._status_lock: + # A concurrent cancel_active_task() may have + # already set the outcome. + if self._current.outcome is None: + self._current.set_result(result) + except RunEngineInterrupted: + # Raised for both a normal pause (outcome still None) + # and an abort interrupting a running plan (outcome + # already TaskError) - only the latter is a failure. + if isinstance(self._current.outcome, TaskError): + self._report_error(Exception(self._current.outcome.message)) except Exception as e: LOGGER.error("Task failed", extra=meta) - self._current.set_exception(e) - self._report_error(e) + with self._status_lock: + if self._current.outcome is None: + self._current.set_exception(e) + raise with plan_tag_filter_context(next_task.task.name, LOGGER): if self._current_task_otel_context is not None: @@ -463,10 +502,43 @@ def process_task(): else: process_task() + elif isinstance(next_task, ResumeSignal): + pending_cancel = self._pending_cancel + if pending_cancel is not None: + # A cancel was requested after this resume was already + # queued - it takes priority, so this resume never runs. + self._apply_cancel(pending_cancel) + elif self._ctx.run_engine.state == "paused": + if self._current is not None: + try: + result = self._ctx.run_engine.resume() + self._current.set_result(result) + except RunEngineInterrupted: + pass + else: + LOGGER.warning( + "Received resume signal but no active task, ignoring" + ) + else: + LOGGER.warning( + "Received resume signal but RunEngine is not paused, ignoring" + ) + + elif isinstance(next_task, CancelSignal): + self._apply_cancel(self._pending_cancel or next_task) + elif isinstance(next_task, KillSignal): # If we receive a kill signal we begin to shut the worker down. # Note that the kill signal is explicitly not a type of task as we don't # want it to be part of the worker's public API + self._pending_cancel = None + if self._current is not None and self._ctx.run_engine.state == "paused": + self._apply_cancel( + CancelSignal( + failure=True, + reason="Worker is stopping while the task was paused", + ) + ) self._stopping.set() add_span_attributes({"server shutting down": "true"}) else: @@ -474,17 +546,25 @@ def process_task(): except Exception as err: self._report_error(err) finally: - if self._current_task_otel_context is not None: + if ( + self._current_task_otel_context is not None + and self._ctx.run_engine.state not in ["panicked", "paused"] + ): self._current_task_otel_context = None if self._current is not None: - self._current.is_complete = True - self._pending_tasks.pop(self._current.task_id) - self._completed_tasks[self._current.task_id] = self._current - self._report_status() - self._errors.clear() - self._warnings.clear() - self._completed_statuses.clear() + # Don't finalize while paused - the task isn't done yet, it may + # still be resumed or cancelled. + if self._ctx.run_engine.state != "paused": + self._current.is_complete = True + self._pending_tasks.pop(self._current.task_id) + self._completed_tasks[self._current.task_id] = self._current + if self._ctx.run_engine.state != "paused": + self._report_status() + self._current = None + self._errors.clear() + self._warnings.clear() + self._completed_statuses.clear() @property def worker_events(self) -> EventStream[WorkerEvent, int]: @@ -534,6 +614,34 @@ def _report_error(self, err: Exception) -> None: self._current.errors.append(str(err)) self._errors.append(str(err)) + def _apply_cancel(self, signal: "CancelSignal") -> None: + # Only ever runs on the worker thread, so no lock is needed here - + # unlike the caller-thread path in cancel_active_task(). + self._pending_cancel = None + default_reason = "Task failed for unknown reason" + if self._current is not None: + if signal.failure: + reason = signal.reason or default_reason + self._ctx.run_engine.abort(reason) + self._current.set_exception(Exception(reason)) + else: + self._ctx.run_engine.stop() + self._current.set_result(None) + + if signal.failure: + error_message = signal.reason or default_reason + add_span_attributes({"Task aborted": error_message}) + LOGGER.error("Task failed: %s", error_message) + if self._current is not None: + self._report_error(Exception(error_message)) + else: + add_span_attributes( + { + "Task stopped": signal.reason + or "Cancellation successful: Task stopped without error" + } + ) + @start_as_current_span(TRACER) def _report_status( self, @@ -545,7 +653,8 @@ def _report_status( task_status = TaskStatus( task_id=self._current.task_id, task_complete=self._current.is_complete, - task_failed=bool(self._current.errors), + task_failed=bool(self._current.errors) + or isinstance(self._current.outcome, TaskError), result=self._current.outcome, ) correlation_id = self._current.task_id @@ -603,7 +712,7 @@ def _on_document(self, name: str, document: Mapping[str, Any]) -> None: ) else: - raise KeyError( + raise RuntimeError( "Trying to emit a document despite the fact that the RunEngine is idle" ) @@ -688,6 +797,25 @@ class KillSignal: ... +@dataclass +class ResumeSignal: + """ + Object put in the worker's task queue to tell it to resume if paused. + """ + + ... + + +@dataclass +class CancelSignal: + """ + Object put in the worker's task queue to tell it to cancel the current task. + """ + + failure: bool + reason: str | None + + def run_worker_in_own_thread( worker: TaskWorker, executor: ThreadPoolExecutor | None = None ) -> Future: diff --git a/tests/unit_tests/worker/test_task_worker.py b/tests/unit_tests/worker/test_task_worker.py index f03ba7c5f..f5ef9e1fc 100644 --- a/tests/unit_tests/worker/test_task_worker.py +++ b/tests/unit_tests/worker/test_task_worker.py @@ -1,6 +1,7 @@ import dataclasses import itertools import threading +import time from collections.abc import Callable, Iterable from concurrent.futures import Future from pathlib import Path @@ -8,6 +9,7 @@ from typing import Any, TypeVar from unittest.mock import ANY, MagicMock, Mock, patch +import bluesky.plan_stubs as plan_stubs import pydantic import pytest from bluesky.protocols import Movable, Readable, Status @@ -35,7 +37,7 @@ WorkerEvent, WorkerState, ) -from blueapi.worker.event import TaskResult, TaskStatusEnum +from blueapi.worker.event import TaskError, TaskResult, TaskStatusEnum _SIMPLE_TASK = Task(name="sleep", params={"time": 0.0}) _LONG_TASK = Task(name="sleep", params={"time": 1.0}) @@ -44,6 +46,9 @@ params={"movable": "fake_device", "value": 4.0}, ) _FAILING_TASK = Task(name="failing_plan", params={}) +_PAUSING_TASK = Task(name="pausing_plan", params={}) +_PAUSING_TASK_SLOW_CLEANUP = Task(name="pausing_plan_with_slow_cleanup", params={}) +_ABORT_CLEANUP_DELAY = 1.0 _TASK_WITH_METADATA = Task( name="sleep", params={"time": 0.0}, @@ -79,6 +84,17 @@ def failing_plan() -> MsgGenerator: raise KeyError("I failed") +def pausing_plan() -> MsgGenerator: + yield from plan_stubs.pause() + + +def pausing_plan_with_slow_cleanup() -> MsgGenerator: + try: + yield from plan_stubs.pause() + finally: + time.sleep(_ABORT_CLEANUP_DELAY) + + @dataclasses.dataclass class ComplexReturn: foo: int @@ -344,9 +360,9 @@ def test_plan_failure_recorded_in_active_task(worker: TaskWorker) -> None: assert events[-1].task_status.task_failed assert events[-1].errors == ["'I failed'"] - active_task = worker.get_active_task() - assert active_task is not None - assert active_task.errors == ["'I failed'"] + completed_task = worker._completed_tasks.get(task_id) + assert completed_task is not None + assert completed_task.errors == ["'I failed'"] def test_task_not_run_twice(worker: TaskWorker) -> None: @@ -455,6 +471,19 @@ def begin_task_and_wait_until_complete( return events.result(timeout=timeout) +def begin_task_and_wait_until_paused( + worker: TaskWorker, + task_id: str, + timeout: float = 5.0, +) -> None: + paused_future: Future[list[WorkerEvent]] = take_events( + worker.worker_events, + lambda e: e.state == WorkerState.PAUSED, + ) + worker.begin_task(task_id) + paused_future.result(timeout=timeout) + + # # Event stream helpers # @@ -922,3 +951,280 @@ def test_task_result_serialization(plan_result, task_result, type_name): res = TaskResult.from_result(plan_result) assert res.result == task_result assert res.type == type_name + + +def test_pause_does_not_publish_error_event(worker: TaskWorker) -> None: + worker._ctx.register_plan(pausing_plan) + task_id = worker.submit_task(_PAUSING_TASK) + + events_future: Future[list[WorkerEvent]] = take_events( + worker.worker_events, + lambda e: e.state == WorkerState.PAUSED, + ) + worker.begin_task(task_id) + events = events_future.result(timeout=5.0) + + assert all(e.errors == [] for e in events) + + +def test_paused_task_remains_in_pending(worker: TaskWorker) -> None: + worker._ctx.register_plan(pausing_plan) + task_id = worker.submit_task(_PAUSING_TASK) + + begin_task_and_wait_until_paused(worker, task_id) + + assert task_id in worker._pending_tasks + assert task_id not in worker._completed_tasks + + +def test_worker_state_is_paused(worker: TaskWorker) -> None: + worker._ctx.register_plan(pausing_plan) + task_id = worker.submit_task(_PAUSING_TASK) + + begin_task_and_wait_until_paused(worker, task_id) + + assert worker.state == WorkerState.PAUSED + + +def test_resume_after_pause_completes_task(worker: TaskWorker) -> None: + worker._ctx.register_plan(pausing_plan) + task_id = worker.submit_task(_PAUSING_TASK) + + begin_task_and_wait_until_paused(worker, task_id) + + complete_future: Future[list[WorkerEvent]] = take_events( + worker.worker_events, + lambda e: e.is_complete(), + ) + worker.resume() + events = complete_future.result(timeout=5.0) + + assert events[-1].task_status is not None + assert events[-1].task_status.task_complete + assert isinstance(events[-1].task_status.result, TaskResult) + assert task_id in worker._completed_tasks + assert task_id not in worker._pending_tasks + + +def test_cancel_active_task_abort(worker: TaskWorker) -> None: + worker._ctx.register_plan(pausing_plan) + task_id = worker.submit_task(_PAUSING_TASK) + + begin_task_and_wait_until_paused(worker, task_id) + + cancel_future: Future[list[WorkerEvent]] = take_events( + worker.worker_events, + lambda e: ( + e.state == WorkerState.IDLE + and e.task_status is not None + and e.task_status.task_complete + ), + ) + worker.cancel_active_task(failure=True) + events = cancel_future.result(timeout=5.0) + + assert events[-1].errors == ["Task failed for unknown reason"] + assert events[-1].task_status is not None + assert events[-1].task_status.task_failed + assert task_id in worker._completed_tasks + assert task_id not in worker._pending_tasks + + +def test_cancel_active_task_graceful(worker: TaskWorker) -> None: + worker._ctx.register_plan(pausing_plan) + task_id = worker.submit_task(_PAUSING_TASK) + + begin_task_and_wait_until_paused(worker, task_id) + + cancel_future: Future[list[WorkerEvent]] = take_events( + worker.worker_events, + lambda e: ( + e.state == WorkerState.IDLE + and e.task_status is not None + and e.task_status.task_complete + ), + ) + worker.cancel_active_task(failure=False) + events = cancel_future.result(timeout=5.0) + + assert events[-1].errors == [] + assert events[-1].task_status is not None + assert events[-1].task_status.task_complete + assert isinstance(events[-1].task_status.result, TaskResult) + assert task_id in worker._completed_tasks + assert task_id not in worker._pending_tasks + + +def test_cancel_running_task_records_failure(worker: TaskWorker) -> None: + # _LONG_TASK uses asyncio.sleep so the RE event loop is free to process the + # abort signal — unlike FakeDevice which blocks the loop with a sync wait. + task_id = worker.submit_task(_LONG_TASK) + + running_future: Future[list[WorkerEvent]] = take_events( + worker.worker_events, + lambda e: e.state == WorkerState.RUNNING, + ) + worker.begin_task(task_id) + running_future.result(timeout=5.0) + + cancel_future: Future[list[WorkerEvent]] = take_events( + worker.worker_events, + lambda e: e.is_complete(), + ) + worker.cancel_active_task(failure=True, reason="mid-run abort") + events = cancel_future.result(timeout=5.0) + + assert events[-1].task_status is not None + assert events[-1].task_status.task_failed + assert isinstance(events[-1].task_status.result, TaskError) + assert task_id in worker._completed_tasks + assert task_id not in worker._pending_tasks + + +def test_cancel_active_task_does_not_overwrite_existing_outcome( + worker: TaskWorker, +) -> None: + task_id = worker.submit_task(_LONG_TASK) + + running_future: Future[list[WorkerEvent]] = take_events( + worker.worker_events, + lambda e: e.state == WorkerState.RUNNING, + ) + worker.begin_task(task_id) + running_future.result(timeout=5.0) + + current = worker._current + assert current is not None + with worker._status_lock: + current.set_exception(Exception("worker thread got there first")) + + worker.cancel_active_task(failure=True, reason="caller's reason") + + assert isinstance(current.outcome, TaskError) + assert current.outcome.message == "worker thread got there first" + + +def test_cancel_active_task_does_not_block_caller_when_paused( + worker: TaskWorker, +) -> None: + + worker._ctx.register_plan(pausing_plan_with_slow_cleanup) + task_id = worker.submit_task(_PAUSING_TASK_SLOW_CLEANUP) + + begin_task_and_wait_until_paused(worker, task_id) + + start = time.monotonic() + worker.cancel_active_task(failure=True) + elapsed = time.monotonic() - start + + assert elapsed < _ABORT_CLEANUP_DELAY / 2, ( + f"cancel_active_task(failure=True) blocked its caller for " + f"{elapsed:.2f}s while the RunEngine finished aborting a paused " + "task. It must enqueue the cancellation and return immediately, " + "the same way resume() does, instead of calling RE.abort() " + "synchronously on the caller's thread." + ) + + complete_future: Future[list[WorkerEvent]] = take_events( + worker.worker_events, + lambda e: e.is_complete(), + ) + events = complete_future.result(timeout=5.0) + assert events[-1].task_status is not None + assert events[-1].task_status.task_failed + assert task_id in worker._completed_tasks + + +def test_cancel_wins_race_with_concurrent_resume(worker: TaskWorker) -> None: + # Regression test: resume() and cancel_active_task() are both deferred to + # the worker thread while paused. If resume() is queued first and runs to + # completion before the cancel is looked at, the cancel must still win - + # not be silently dropped in favour of a "successful" task. + worker._ctx.register_plan(pausing_plan) + task_id = worker.submit_task(_PAUSING_TASK) + + begin_task_and_wait_until_paused(worker, task_id) + + complete_future: Future[list[WorkerEvent]] = take_events( + worker.worker_events, + lambda e: e.is_complete(), + ) + worker.resume() + worker.cancel_active_task(failure=True, reason="changed my mind") + events = complete_future.result(timeout=5.0) + + assert events[-1].task_status is not None + assert events[-1].task_status.task_failed + assert isinstance(events[-1].task_status.result, TaskError) + assert events[-1].errors == ["changed my mind"] + assert task_id in worker._completed_tasks + assert task_id not in worker._pending_tasks + + +def test_second_cancel_while_paused_supersedes_first(worker: TaskWorker) -> None: + + worker._ctx.register_plan(pausing_plan) + task_id = worker.submit_task(_PAUSING_TASK) + + begin_task_and_wait_until_paused(worker, task_id) + + complete_future: Future[list[WorkerEvent]] = take_events( + worker.worker_events, + lambda e: e.is_complete(), + ) + worker.cancel_active_task(failure=False, reason="first, graceful") + worker.cancel_active_task(failure=True, reason="second, urgent") + events = complete_future.result(timeout=5.0) + + assert events[-1].task_status is not None + assert events[-1].task_status.task_failed + assert isinstance(events[-1].task_status.result, TaskError) + assert events[-1].errors == ["second, urgent"] + assert task_id in worker._completed_tasks + assert task_id not in worker._pending_tasks + + +def test_stop_while_paused_completes_task(worker: TaskWorker) -> None: + worker._ctx.register_plan(pausing_plan) + task_id = worker.submit_task(_PAUSING_TASK) + + begin_task_and_wait_until_paused(worker, task_id) + + worker.stop() + + assert task_id in worker._completed_tasks + assert task_id not in worker._pending_tasks + assert worker._completed_tasks[task_id].is_complete + assert worker._completed_tasks[task_id].outcome is not None + + +def test_resume_when_not_paused_does_nothing(worker: TaskWorker) -> None: + task_id = worker.submit_task(_SIMPLE_TASK) + begin_task_and_wait_until_complete(worker, task_id) + + complete_future: Future[list[WorkerEvent]] = take_events( + worker.worker_events, + lambda e: e.state == WorkerState.IDLE, + ) + worker.resume() + events = complete_future.result(timeout=5.0) + + assert all(e.errors == [] for e in events) + + +def test_can_run_task_after_resume(worker: TaskWorker) -> None: + worker._ctx.register_plan(pausing_plan) + task_id = worker.submit_task(_PAUSING_TASK) + begin_task_and_wait_until_paused(worker, task_id) + + complete_future: Future[list[WorkerEvent]] = take_events( + worker.worker_events, + lambda e: e.is_complete(), + ) + worker.resume() + complete_future.result(timeout=5.0) + + task_id_2 = worker.submit_task(_SIMPLE_TASK) + events = begin_task_and_wait_until_complete(worker, task_id_2) + assert events[-1].task_status is not None + assert events[-1].task_status.task_complete