From 96a72e3df3cad66cf762c423ac4f54922c64a7c5 Mon Sep 17 00:00:00 2001 From: alxj9837 <52531949+Alexj9837@users.noreply.github.com> Date: Tue, 30 Jun 2026 13:36:01 +0000 Subject: [PATCH 01/10] fix(worker): correctly handle pause, resume and cancel via worker thread signals Route resume() and cancel_active_task() through the worker thread's signal queue instead of mutating state directly from the caller, so state stays consistent across pause/resume/cancel transitions. Includes tests and a fix for cancel_active_task() when the RunEngine is already in a paused state. --- src/blueapi/worker/task_worker.py | 146 +++++++++-- tests/unit_tests/worker/test_task_worker.py | 254 +++++++++++++++++++- 2 files changed, 374 insertions(+), 26 deletions(-) diff --git a/src/blueapi/worker/task_worker.py b/src/blueapi/worker/task_worker.py index 1e09f3bab9..c804bb9d6c 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,39 @@ 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}) + 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 + 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, "task_id") def get_task_by_id(self, task_id: str) -> TrackableTask | None: @@ -410,7 +433,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: @@ -436,11 +459,15 @@ def process_task(): LOGGER.info( "Task ran successfully - returned: %s", result, extra=meta ) - self._current.set_result(result) + if self._current.outcome is None: + self._current.set_result(result) + except RunEngineInterrupted: + 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) + raise with plan_tag_filter_context(next_task.task.name, LOGGER): if self._current_task_otel_context is not None: @@ -458,10 +485,33 @@ 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(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 self._stopping.set() add_span_attributes({"server shutting down": "true"}) else: @@ -469,17 +519,23 @@ 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() + 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]: @@ -529,6 +585,32 @@ 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: + 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, @@ -540,7 +622,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 @@ -598,7 +681,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" ) @@ -683,6 +766,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 588c37d8e4..f621946476 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,220 @@ 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_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_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 From c21efe0b96d12be77626f87903ac3b58d0edf674 Mon Sep 17 00:00:00 2001 From: alxj9837 <52531949+Alexj9837@users.noreply.github.com> Date: Wed, 8 Jul 2026 12:03:47 +0000 Subject: [PATCH 02/10] fix(worker): fix pause/resume/cancel bugs found in review - Resolve a paused task before stopping the worker, instead of leaving the RunEngine paused and the task incomplete forever with no thread left to resolve it. - Apply the latest cancel_active_task() request while paused instead of a stale queued one. - Prevent a race between the caller thread and worker thread when recording a cancelled task's outcome. --- src/blueapi/worker/task_worker.py | 38 ++++++++++--- tests/unit_tests/worker/test_task_worker.py | 60 +++++++++++++++++++++ 2 files changed, 92 insertions(+), 6 deletions(-) diff --git a/src/blueapi/worker/task_worker.py b/src/blueapi/worker/task_worker.py index c804bb9d6c..49f1f22b37 100644 --- a/src/blueapi/worker/task_worker.py +++ b/src/blueapi/worker/task_worker.py @@ -207,10 +207,14 @@ def cancel_active_task( default_reason = "Task failed for unknown reason" if failure: self._ctx.run_engine.abort(reason or default_reason) - current.set_exception(Exception(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() - current.set_result(None) + 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 ""} @@ -459,14 +463,22 @@ def process_task(): LOGGER.info( "Task ran successfully - returned: %s", result, extra=meta ) - if self._current.outcome is None: - 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) + 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): @@ -508,10 +520,20 @@ def process_task(): ) elif isinstance(next_task, CancelSignal): - self._apply_cancel(next_task) + 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: @@ -526,6 +548,8 @@ def process_task(): self._current_task_otel_context = None if self._current is not None: + # 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) @@ -586,6 +610,8 @@ def _report_error(self, err: Exception) -> None: 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: diff --git a/tests/unit_tests/worker/test_task_worker.py b/tests/unit_tests/worker/test_task_worker.py index f621946476..385a74bc69 100644 --- a/tests/unit_tests/worker/test_task_worker.py +++ b/tests/unit_tests/worker/test_task_worker.py @@ -1081,6 +1081,29 @@ def test_cancel_running_task_records_failure(worker: TaskWorker) -> None: 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: @@ -1138,6 +1161,43 @@ def test_cancel_wins_race_with_concurrent_resume(worker: TaskWorker) -> None: 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) From 7adfe361991249d44b9ab49e6838b0d9129e7558 Mon Sep 17 00:00:00 2001 From: alxj9837 <52531949+Alexj9837@users.noreply.github.com> Date: Mon, 13 Jul 2026 09:07:36 +0000 Subject: [PATCH 03/10] adding more tests to cover when a resume plan is paused immediaetly. --- tests/unit_tests/worker/test_task_worker.py | 58 +++++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/tests/unit_tests/worker/test_task_worker.py b/tests/unit_tests/worker/test_task_worker.py index 385a74bc69..75b15a6f60 100644 --- a/tests/unit_tests/worker/test_task_worker.py +++ b/tests/unit_tests/worker/test_task_worker.py @@ -48,6 +48,7 @@ _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={}) +_TWICE_PAUSING_TASK = Task(name="twice_pausing_plan", params={}) _ABORT_CLEANUP_DELAY = 1.0 _TASK_WITH_METADATA = Task( name="sleep", @@ -95,6 +96,11 @@ def pausing_plan_with_slow_cleanup() -> MsgGenerator: time.sleep(_ABORT_CLEANUP_DELAY) +def twice_pausing_plan() -> MsgGenerator: + yield from plan_stubs.pause() + yield from plan_stubs.pause() + + @dataclasses.dataclass class ComplexReturn: foo: int @@ -1081,6 +1087,34 @@ def test_cancel_running_task_records_failure(worker: TaskWorker) -> None: assert task_id not in worker._pending_tasks +def test_cancel_running_task_gracefully(worker: TaskWorker) -> None: + # Covers the failure=False branch of cancel_active_task() while the + # RunEngine is actively running (not paused) - RE.stop() is called + # directly on the caller's thread and must not be reported as a failure. + 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=False) + events = cancel_future.result(timeout=5.0) + + assert events[-1].errors == [] + assert events[-1].task_status is not None + assert not events[-1].task_status.task_failed + 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_does_not_overwrite_existing_outcome( worker: TaskWorker, ) -> None: @@ -1212,6 +1246,30 @@ def test_resume_when_not_paused_does_nothing(worker: TaskWorker) -> None: assert all(e.errors == [] for e in events) +def test_resume_re_pauses_when_plan_pauses_again(worker: TaskWorker) -> None: + # Covers the RunEngineInterrupted branch of the ResumeSignal handler: + # RE.resume() raises RunEngineInterrupted rather than completing when the + # plan immediately pauses again, and this must not be treated as a + # failure or leave the task in a half-finished state. + worker._ctx.register_plan(twice_pausing_plan) + task_id = worker.submit_task(_TWICE_PAUSING_TASK) + begin_task_and_wait_until_paused(worker, task_id) + + paused_again_future: Future[list[WorkerEvent]] = take_events( + worker.worker_events, + lambda e: e.state == WorkerState.PAUSED, + ) + worker.resume() + paused_again_future.result(timeout=5.0) + + assert worker.state == WorkerState.PAUSED + current = worker._current + assert current is not None + assert current.outcome is None + assert task_id in worker._pending_tasks + assert task_id not in worker._completed_tasks + + def test_can_run_task_after_resume(worker: TaskWorker) -> None: worker._ctx.register_plan(pausing_plan) task_id = worker.submit_task(_PAUSING_TASK) From 9988e2d8de1c883894143f1fc8dd8f05ba073a36 Mon Sep 17 00:00:00 2001 From: alxj9837 <52531949+Alexj9837@users.noreply.github.com> Date: Tue, 14 Jul 2026 08:05:29 +0000 Subject: [PATCH 04/10] code review suggestions --- src/blueapi/worker/task_worker.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/blueapi/worker/task_worker.py b/src/blueapi/worker/task_worker.py index 49f1f22b37..63cb297201 100644 --- a/src/blueapi/worker/task_worker.py +++ b/src/blueapi/worker/task_worker.py @@ -509,7 +509,10 @@ def process_task(): result = self._ctx.run_engine.resume() self._current.set_result(result) except RunEngineInterrupted: - pass + # Expected when the plan pauses again immediately + # after resuming - not a failure, just leave the + # task's outcome unset so it stays resumable. + LOGGER.debug("RunEngine resume interrupted; ignoring") else: LOGGER.warning( "Received resume signal but no active task, ignoring" @@ -798,7 +801,7 @@ class ResumeSignal: Object put in the worker's task queue to tell it to resume if paused. """ - ... + pass @dataclass From 8549ca6e68180ee236b119842c39fd99ae985e35 Mon Sep 17 00:00:00 2001 From: alxj9837 <52531949+Alexj9837@users.noreply.github.com> Date: Tue, 14 Jul 2026 09:21:13 +0000 Subject: [PATCH 05/10] cancel_active_task() set the outcome after RE.abort()/stop() returned, racing the worker thread's own finalization, which could report a completed task with no result. --- src/blueapi/worker/task_worker.py | 11 +++++++++-- tests/system_tests/test_blueapi_system.py | 6 +++++- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/src/blueapi/worker/task_worker.py b/src/blueapi/worker/task_worker.py index 63cb297201..4b7da41bc3 100644 --- a/src/blueapi/worker/task_worker.py +++ b/src/blueapi/worker/task_worker.py @@ -204,17 +204,24 @@ def cancel_active_task( # 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() + # + # The outcome must be set *before* calling abort()/stop(), not after: + # those calls block until the worker thread's do_task() unblocks (either + # returning normally or raising RunEngineInterrupted), and that thread's + # own handling doesn't set a result for a plain stop. Setting it first + # ensures the outcome is already in place by the time the worker thread + # finalizes and reports the task's final status. default_reason = "Task failed for unknown reason" if failure: - self._ctx.run_engine.abort(reason or default_reason) with self._status_lock: if current.outcome is None: current.set_exception(Exception(reason or default_reason)) + self._ctx.run_engine.abort(reason or default_reason) else: - self._ctx.run_engine.stop() with self._status_lock: if current.outcome is None: current.set_result(None) + self._ctx.run_engine.stop() self._task_channel.put(CancelSignal(failure=failure, reason=reason)) add_span_attributes( {"Task aborted" if failure else "Task stopped": reason or ""} diff --git a/tests/system_tests/test_blueapi_system.py b/tests/system_tests/test_blueapi_system.py index fa44e5ff9e..d00a5cfc8a 100644 --- a/tests/system_tests/test_blueapi_system.py +++ b/tests/system_tests/test_blueapi_system.py @@ -349,10 +349,14 @@ def test_delete_non_existent_task(rest_client: BlueapiRestClient): def test_put_worker_task(rest_client: BlueapiRestClient): - created_task = rest_client.create_task(_SIMPLE_TASK) + # Uses _LONG_TASK, not _SIMPLE_TASK: the active task is cleared as soon as + # it completes, so a near-instant task can finish before get_active_task() + # is called, racily returning None instead of the task just submitted. + created_task = rest_client.create_task(_LONG_TASK) rest_client.update_worker_task(WorkerTask(task_id=created_task.task_id)) active_task = rest_client.get_active_task() assert active_task.task_id == created_task.task_id + rest_client.cancel_current_task(WorkerState.ABORTING) rest_client.clear_task(created_task.task_id) From c4f167198a509c2786e0ee0822dd76e65c94d015 Mon Sep 17 00:00:00 2001 From: alxj9837 <52531949+Alexj9837@users.noreply.github.com> Date: Tue, 14 Jul 2026 09:48:20 +0000 Subject: [PATCH 06/10] reducing comments --- src/blueapi/worker/task_worker.py | 43 ++++++++------------- tests/system_tests/test_blueapi_system.py | 5 +-- tests/unit_tests/worker/test_task_worker.py | 17 +++----- 3 files changed, 24 insertions(+), 41 deletions(-) diff --git a/src/blueapi/worker/task_worker.py b/src/blueapi/worker/task_worker.py index 4b7da41bc3..f6c902587b 100644 --- a/src/blueapi/worker/task_worker.py +++ b/src/blueapi/worker/task_worker.py @@ -190,10 +190,9 @@ def cancel_active_task( 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. + # abort()/stop() block until cleanup finishes, so defer to the worker + # thread. Also recorded in _pending_cancel so a queued ResumeSignal + # can't race ahead and finish the task before this is looked at. signal = CancelSignal(failure=failure, reason=reason) self._pending_cancel = signal try: @@ -202,15 +201,10 @@ def cancel_active_task( 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() - # - # The outcome must be set *before* calling abort()/stop(), not after: - # those calls block until the worker thread's do_task() unblocks (either - # returning normally or raising RunEngineInterrupted), and that thread's - # own handling doesn't set a result for a plain stop. Setting it first - # ensures the outcome is already in place by the time the worker thread - # finalizes and reports the task's final status. + # 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(). The outcome is set beforehand so it's already in place + # once the worker thread's do_task() unblocks and finalizes. default_reason = "Task failed for unknown reason" if failure: with self._status_lock: @@ -471,14 +465,13 @@ def process_task(): "Task ran successfully - returned: %s", result, extra=meta ) with self._status_lock: - # A concurrent cancel_active_task() may have - # already set the outcome. + # cancel_active_task() may have set this concurrently. 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. + # Raised by both a pause (outcome still None) and an + # abort (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: @@ -507,8 +500,8 @@ def 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. + # A cancel queued after this resume 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: @@ -516,9 +509,8 @@ def process_task(): result = self._ctx.run_engine.resume() self._current.set_result(result) except RunEngineInterrupted: - # Expected when the plan pauses again immediately - # after resuming - not a failure, just leave the - # task's outcome unset so it stays resumable. + # Plan paused again immediately - not a failure, + # leave the outcome unset so it stays resumable. LOGGER.debug("RunEngine resume interrupted; ignoring") else: LOGGER.warning( @@ -558,8 +550,7 @@ def process_task(): self._current_task_otel_context = None if self._current is not None: - # Don't finalize while paused - the task isn't done yet, it may - # still be resumed or cancelled. + # Not done yet while paused - 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) @@ -620,8 +611,6 @@ def _report_error(self, err: Exception) -> None: 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: diff --git a/tests/system_tests/test_blueapi_system.py b/tests/system_tests/test_blueapi_system.py index d00a5cfc8a..6ec451944e 100644 --- a/tests/system_tests/test_blueapi_system.py +++ b/tests/system_tests/test_blueapi_system.py @@ -349,9 +349,8 @@ def test_delete_non_existent_task(rest_client: BlueapiRestClient): def test_put_worker_task(rest_client: BlueapiRestClient): - # Uses _LONG_TASK, not _SIMPLE_TASK: the active task is cleared as soon as - # it completes, so a near-instant task can finish before get_active_task() - # is called, racily returning None instead of the task just submitted. + # _LONG_TASK, since a near-instant task could complete (clearing the + # active task) before get_active_task() below is called. created_task = rest_client.create_task(_LONG_TASK) rest_client.update_worker_task(WorkerTask(task_id=created_task.task_id)) active_task = rest_client.get_active_task() diff --git a/tests/unit_tests/worker/test_task_worker.py b/tests/unit_tests/worker/test_task_worker.py index 75b15a6f60..7f9e8809de 100644 --- a/tests/unit_tests/worker/test_task_worker.py +++ b/tests/unit_tests/worker/test_task_worker.py @@ -1088,9 +1088,8 @@ def test_cancel_running_task_records_failure(worker: TaskWorker) -> None: def test_cancel_running_task_gracefully(worker: TaskWorker) -> None: - # Covers the failure=False branch of cancel_active_task() while the - # RunEngine is actively running (not paused) - RE.stop() is called - # directly on the caller's thread and must not be reported as a failure. + # cancel_active_task(failure=False) on a running (not paused) task must + # not be reported as a failure. task_id = worker.submit_task(_LONG_TASK) running_future: Future[list[WorkerEvent]] = take_events( @@ -1170,10 +1169,8 @@ def test_cancel_active_task_does_not_block_caller_when_paused( 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. + # A cancel queued right after a resume must still win, not be dropped in + # favour of the resume completing the task first. worker._ctx.register_plan(pausing_plan) task_id = worker.submit_task(_PAUSING_TASK) @@ -1247,10 +1244,8 @@ def test_resume_when_not_paused_does_nothing(worker: TaskWorker) -> None: def test_resume_re_pauses_when_plan_pauses_again(worker: TaskWorker) -> None: - # Covers the RunEngineInterrupted branch of the ResumeSignal handler: - # RE.resume() raises RunEngineInterrupted rather than completing when the - # plan immediately pauses again, and this must not be treated as a - # failure or leave the task in a half-finished state. + # If the plan pauses again immediately on resume, that's not a failure + # and must not leave the task half-finished. worker._ctx.register_plan(twice_pausing_plan) task_id = worker.submit_task(_TWICE_PAUSING_TASK) begin_task_and_wait_until_paused(worker, task_id) From 5e78865d9b482fdf4a59cf5f20f501ef0a9a2536 Mon Sep 17 00:00:00 2001 From: Alex J <52531949+Alexj9837@users.noreply.github.com> Date: Wed, 12 Aug 2026 12:34:25 +0100 Subject: [PATCH 07/10] Refactor test_put_worker_task to accept long_task fixed a merge conflict error --- src/blueapi/worker/task_worker.py | 29 +++++++++++------------ tests/system_tests/test_blueapi_system.py | 12 ++++++---- 2 files changed, 22 insertions(+), 19 deletions(-) diff --git a/src/blueapi/worker/task_worker.py b/src/blueapi/worker/task_worker.py index c7cc03e672..c55c5f190c 100644 --- a/src/blueapi/worker/task_worker.py +++ b/src/blueapi/worker/task_worker.py @@ -556,8 +556,9 @@ def process_task(): 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() + finished_task = self._current self._current = None + self._report_status(finished_task) self._errors.clear() self._warnings.clear() self._completed_statuses.clear() @@ -602,7 +603,7 @@ def _on_state_change( old_state = WorkerState.UNKNOWN LOGGER.debug(f"Notifying state change {old_state} -> {new_state}") self._state = new_state - self._report_status() + self._report_status(self._current) def _report_error(self, err: Exception) -> None: LOGGER.error(err, exc_info=True) @@ -637,26 +638,24 @@ def _apply_cancel(self, signal: "CancelSignal") -> None: ) @start_as_current_span(TRACER) - def _report_status( - self, - ) -> None: + def _report_status(self, current: TrackableTask | None) -> None: task_status: TaskStatus | None errors = self._errors warnings = self._warnings - if self._current is not None: + if current is not None: task_status = TaskStatus( - task_id=self._current.task_id, - task_complete=self._current.is_complete, - task_failed=bool(self._current.errors) - or isinstance(self._current.outcome, TaskError), - result=self._current.outcome, + task_id=current.task_id, + task_complete=current.is_complete, + task_failed=bool(current.errors) + or isinstance(current.outcome, TaskError), + result=current.outcome, ) - correlation_id = self._current.task_id + correlation_id = current.task_id add_span_attributes( { - "task_id": self._current.task_id, - "task_complete": self._current.is_complete, - "task_failed": self._current.errors, + "task_id": current.task_id, + "task_complete": current.is_complete, + "task_failed": current.errors, } ) else: diff --git a/tests/system_tests/test_blueapi_system.py b/tests/system_tests/test_blueapi_system.py index cef97971a7..970bd909e2 100644 --- a/tests/system_tests/test_blueapi_system.py +++ b/tests/system_tests/test_blueapi_system.py @@ -409,10 +409,10 @@ def test_delete_non_existent_task(rest_client: BlueapiRestClient): rest_client.clear_task("Not-exists") -def test_put_worker_task(rest_client: BlueapiRestClient): - # _LONG_TASK, since a near-instant task could complete (clearing the +def test_put_worker_task(rest_client: BlueapiRestClient, long_task: TaskRequest): + # long_task, since a near-instant task could complete (clearing the # active task) before get_active_task() below is called. - created_task = rest_client.create_task(_LONG_TASK) + created_task = rest_client.create_task(long_task) rest_client.update_worker_task(WorkerTask(task_id=created_task.task_id)) active_task = rest_client.get_active_task() assert active_task.task_id == created_task.task_id @@ -809,7 +809,9 @@ def test_any_user_can_retrieve_active_task( task_id = ( client_factory[AdminUser.admin] .create_and_start_task( - task_factory(AdminUser.admin, VALID_INSTRUMENT_SESSION[AdminUser.admin]) + task_factory( + AdminUser.admin, VALID_INSTRUMENT_SESSION[AdminUser.admin], time=1 + ) ) .task_id ) @@ -817,6 +819,8 @@ def test_any_user_can_retrieve_active_task( for user in User: assert client_factory[user].get_active_task().task_id == task_id + client_factory[AdminUser.admin].abort() + def test_non_admin_can_only_start_own_tasks( client_factory: dict[ValidUser, BlueapiClient], From 3d2eb4f717e5044a845df9a8e167c3a099776ac0 Mon Sep 17 00:00:00 2001 From: alexj9837 <52531949+Alexj9837@users.noreply.github.com> Date: Wed, 19 Aug 2026 14:30:40 +0000 Subject: [PATCH 08/10] extracting suggestion from #1622 , going for a blocking run egine design change due to race conditions. --- docs/reference/openapi.yaml | 5 +- src/blueapi/config.py | 2 +- src/blueapi/worker/task.py | 12 +- src/blueapi/worker/task_worker.py | 197 +++++++++++++------- tests/unit_tests/service/test_interface.py | 2 - tests/unit_tests/service/test_rest_api.py | 12 +- tests/unit_tests/worker/test_task_worker.py | 115 ++++++------ 7 files changed, 203 insertions(+), 142 deletions(-) diff --git a/docs/reference/openapi.yaml b/docs/reference/openapi.yaml index e0cd8282df..9182125167 100644 --- a/docs/reference/openapi.yaml +++ b/docs/reference/openapi.yaml @@ -365,7 +365,7 @@ components: title: Errors type: array is_complete: - default: false + readOnly: true title: Is Complete type: boolean is_pending: @@ -389,6 +389,7 @@ components: required: - task_id - task + - is_complete title: TrackableTask type: object ValidationError: @@ -449,7 +450,7 @@ info: name: Apache 2.0 url: https://www.apache.org/licenses/LICENSE-2.0.html title: BlueAPI Control - version: 1.5.0 + version: 1.5.1 openapi: 3.1.0 paths: /api/v1/devices: diff --git a/src/blueapi/config.py b/src/blueapi/config.py index a181f4c344..b93704c57b 100644 --- a/src/blueapi/config.py +++ b/src/blueapi/config.py @@ -324,7 +324,7 @@ class ApplicationConfig(BlueapiBaseModel): """ #: API version to publish in OpenAPI schema - REST_API_VERSION: ClassVar[str] = "1.5.0" + REST_API_VERSION: ClassVar[str] = "1.5.1" LICENSE_INFO: ClassVar[dict[str, str]] = { "name": "Apache 2.0", diff --git a/src/blueapi/worker/task.py b/src/blueapi/worker/task.py index 9ce373c769..632b719658 100644 --- a/src/blueapi/worker/task.py +++ b/src/blueapi/worker/task.py @@ -2,6 +2,7 @@ from collections.abc import Mapping from typing import Any +from bluesky.run_engine import RunEngineResult from pydantic import BaseModel, Field, TypeAdapter from blueapi.core import BlueskyContext @@ -29,7 +30,7 @@ def prepare_params(self, ctx: BlueskyContext) -> Mapping[str, Any]: # Re-create dict manually to avoid nesting in model_dump output return {field: getattr(model, field) for field in model.__pydantic_fields__} - def do_task(self, ctx: BlueskyContext) -> None: + def do_task(self, ctx: BlueskyContext) -> Any: LOGGER.info( f"Asked to run plan {self.name} with {self.params} and " f"metadata {self.metadata} for all runs" @@ -39,9 +40,12 @@ def do_task(self, ctx: BlueskyContext) -> None: prepared_params = self.prepare_params(ctx) ctx.run_engine.md.update(self.metadata) result = ctx.run_engine(func(**prepared_params)) - if isinstance(result, tuple): # pragma: no cover - # this is never true if the run_engine is configured correctly - return None + if not isinstance(result, RunEngineResult): + # this is unreachable unless something has misconfigured it. + raise RuntimeError( + "RunEngine did not return a RunEngineResult - is " + "call_returns_result set on this RunEngine instance?" + ) return result.plan_result diff --git a/src/blueapi/worker/task_worker.py b/src/blueapi/worker/task_worker.py index c55c5f190c..cea92ca325 100644 --- a/src/blueapi/worker/task_worker.py +++ b/src/blueapi/worker/task_worker.py @@ -5,7 +5,7 @@ from dataclasses import dataclass from functools import partial from queue import Full, Queue -from threading import Event, RLock +from threading import Condition, Event, RLock from typing import Any, TypeVar from bluesky import RunEngineInterrupted @@ -20,7 +20,7 @@ from opentelemetry.baggage import get_baggage from opentelemetry.context import Context, get_current from opentelemetry.trace import SpanKind -from pydantic import Field +from pydantic import Field, computed_field, model_validator from pydantic.json_schema import SkipJsonSchema from blueapi.core import ( @@ -69,11 +69,23 @@ class TrackableTask(BlueapiBaseModel): task_id: str task: Task request_id: str | SkipJsonSchema[None] = None - is_complete: bool = False is_pending: bool = True errors: list[str] = Field(default_factory=list) outcome: TaskResult | TaskError | None = None + @computed_field # type: ignore[prop-decorator] + @property + def is_complete(self) -> bool: + return self.outcome is not None + + @model_validator(mode="before") + @classmethod + def remove_complete(cls, values: Any) -> Any: + # fix pydantic falling over itself with computed fields + if isinstance(values, dict): + values.pop("is_complete", None) + return values + def set_result(self, result: Any): self.outcome = TaskResult.from_result(result) @@ -112,6 +124,7 @@ class TaskWorker: _current: TrackableTask | None _pending_cancel: "CancelSignal | None" _status_lock: RLock + _state_change: Condition _status_snapshot: dict[str, StatusView] _completed_statuses: set[str] _worker_events: EventPublisher[WorkerEvent] @@ -146,6 +159,7 @@ def __init__( self._progress_events = EventPublisher() self._data_events = EventPublisher() self._status_lock = RLock() + self._state_change = Condition() self._status_snapshot = {} self._completed_statuses = set() self._started = Event() @@ -183,43 +197,44 @@ def cancel_active_task( Returns: The task_id of the active task """ + 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 until cleanup finishes, so defer to the worker - # thread. Also recorded in _pending_cancel so a queued ResumeSignal - # can't race ahead and finish the task before this is 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(). The outcome is set beforehand so it's already in place - # once the worker thread's do_task() unblocks and finalizes. - default_reason = "Task failed for unknown reason" - if failure: - with self._status_lock: - if current.outcome is None: - current.set_exception(Exception(reason or default_reason)) - self._ctx.run_engine.abort(reason or default_reason) + if self._current_task_otel_context is None: + self._current_task_otel_context = get_current() + + reason = reason or ("Task aborted" if failure else "Task stopped") + if self._state != WorkerState.RUNNING: + with self._state_change: + self._task_channel.put(AbortSignal(failure, reason)) + if not self._state_change.wait_for( + lambda: self._state == WorkerState.IDLE, + timeout=self._start_stop_timeout, + ): + raise TimeoutError( + "Worker did not finish cancelling within " + f"{self._start_stop_timeout} seconds" + ) else: - with self._status_lock: - if current.outcome is None: - current.set_result(None) - self._ctx.run_engine.stop() - self._task_channel.put(CancelSignal(failure=failure, reason=reason)) - add_span_attributes( - {"Task aborted" if failure else "Task stopped": reason or ""} - ) + # The worker thread's own do_task() may be finishing concurrently + # on another thread - don't clobber whatever it already recorded. + try: + if failure: + with self._status_lock: + if current.outcome is None: + current.outcome = TaskError(type="Abort", message=reason) + self._ctx.run_engine.abort(reason=reason) + else: + with self._status_lock: + if current.outcome is None: + current.outcome = TaskResult.from_result(None) + self._ctx.run_engine.stop() + except TransitionError: + return current.task_id return current.task_id @start_as_current_span(TRACER, "task_id") @@ -233,6 +248,9 @@ def get_task_by_id(self, task_id: str) -> TrackableTask | None: Optional[TrackableTask[T]]: The task matching the ID, None if the task ID is unknown to the worker. """ + current = self._current + if current is not None and current.task_id == task_id: + return current return self._pending_tasks.get(task_id, None) or self._completed_tasks[task_id] @start_as_current_span(TRACER) @@ -246,11 +264,8 @@ def get_tasks(self, status: TaskStatusEnum | None = None) -> list[TrackableTask] list[TrackableTask]: A list of tasks that match the given status. """ if status == TaskStatusEnum.RUNNING: - return [ - task - for task in self._pending_tasks.values() - if not task.is_pending and not task.is_complete - ] + current = self._current + return [current] if current is not None else [] elif status == TaskStatusEnum.PENDING: return [task for task in self._pending_tasks.values() if task.is_pending] elif status == TaskStatusEnum.COMPLETE: @@ -343,8 +358,9 @@ def mark_task_as_started(event: WorkerEvent, _: str | None) -> None: self._current_task_otel_context = get_current() """ Cache the current trace context as the one for this task id """ self._task_channel.put_nowait(trackable_task) - task_started.wait(timeout=5.0) - if not task_started.is_set(): + if task_started.wait(timeout=5.0): + self._pending_tasks.pop(trackable_task.task_id) + else: raise TimeoutError("Failed to start plan within timeout") except Full as f: LOGGER.error("Cannot submit task while another is running") @@ -438,7 +454,25 @@ def resume(self): Command the worker to resume """ LOGGER.info("Requesting to resume the worker") - self._task_channel.put(ResumeSignal()) + self._current_task_otel_context = get_current() + current = self._current + + def resumed() -> bool: + if self._state == WorkerState.PAUSED: + return False + if self._state == WorkerState.IDLE: + return current is None or current.task_id in self._completed_tasks + return True + + with self._state_change: + self._task_channel.put(ResumeSignal()) + if not self._state_change.wait_for( + resumed, + timeout=self._start_stop_timeout, + ): + raise TimeoutError( + f"Worker did not resume within {self._start_stop_timeout} seconds" + ) @start_as_current_span(TRACER) def _cycle_with_error_handling(self) -> None: @@ -503,7 +537,7 @@ def process_task(): # A cancel queued after this resume takes priority, so # this resume never runs. self._apply_cancel(pending_cancel) - elif self._ctx.run_engine.state == "paused": + elif self._state == WorkerState.PAUSED: if self._current is not None: try: result = self._ctx.run_engine.resume() @@ -521,7 +555,7 @@ def process_task(): "Received resume signal but RunEngine is not paused, ignoring" ) - elif isinstance(next_task, CancelSignal): + elif isinstance(next_task, (CancelSignal, AbortSignal)): self._apply_cancel(self._pending_cancel or next_task) elif isinstance(next_task, KillSignal): @@ -529,7 +563,7 @@ def process_task(): # 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": + if self._current is not None and self._state == WorkerState.PAUSED: self._apply_cancel( CancelSignal( failure=True, @@ -543,21 +577,21 @@ def process_task(): except Exception as err: self._report_error(err) finally: - if ( - self._current_task_otel_context is not None - and self._ctx.run_engine.state not in ["panicked", "paused"] - ): + if self._current_task_otel_context is not None and self._state not in [ + WorkerState.PANICKED, + WorkerState.PAUSED, + ]: self._current_task_otel_context = None if self._current is not None: # Not done yet while paused - 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) + if self._state != WorkerState.PAUSED: self._completed_tasks[self._current.task_id] = self._current - if self._ctx.run_engine.state != "paused": + if self._state != WorkerState.PAUSED: finished_task = self._current self._current = None + with self._state_change: + self._state_change.notify_all() self._report_status(finished_task) self._errors.clear() self._warnings.clear() @@ -602,7 +636,9 @@ def _on_state_change( else: old_state = WorkerState.UNKNOWN LOGGER.debug(f"Notifying state change {old_state} -> {new_state}") - self._state = new_state + with self._state_change: + self._state = new_state + self._state_change.notify_all() self._report_status(self._current) def _report_error(self, err: Exception) -> None: @@ -611,25 +647,39 @@ 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: + def _finalize_cancel_outcome( + self, current: TrackableTask, signal: "CancelSignal | AbortSignal" + ) -> None: + """ + Set current's outcome for a cancellation, unless something else (e.g. + the worker thread finishing do_task() on its own) already got there first. + """ + with self._status_lock: + if current.outcome is None: + if signal.failure: + current.set_exception( + Exception(signal.reason or "Task failed for unknown reason") + ) + else: + current.set_result(None) + + def _apply_cancel(self, signal: "CancelSignal | AbortSignal") -> None: 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) + current = self._current 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)) + reason = signal.reason or "Task failed for unknown reason" + if current is not None: + self._ctx.run_engine.abort(reason) + self._finalize_cancel_outcome(current, signal) + add_span_attributes({"Task aborted": reason}) + LOGGER.error("Task failed: %s", reason) + if current is not None: + self._report_error(Exception(reason)) else: + if current is not None: + self._ctx.run_engine.stop() + self._finalize_cancel_outcome(current, signal) add_span_attributes( { "Task stopped": signal.reason @@ -643,9 +693,10 @@ def _report_status(self, current: TrackableTask | None) -> None: errors = self._errors warnings = self._warnings if current is not None: + task_complete = current.task_id in self._completed_tasks task_status = TaskStatus( task_id=current.task_id, - task_complete=current.is_complete, + task_complete=task_complete, task_failed=bool(current.errors) or isinstance(current.outcome, TaskError), result=current.outcome, @@ -654,7 +705,7 @@ def _report_status(self, current: TrackableTask | None) -> None: add_span_attributes( { "task_id": current.task_id, - "task_complete": current.is_complete, + "task_complete": task_complete, "task_failed": current.errors, } ) @@ -799,6 +850,12 @@ class ResumeSignal: pass +@dataclass +class AbortSignal: + failure: bool + reason: str + + @dataclass class CancelSignal: """ diff --git a/tests/unit_tests/service/test_interface.py b/tests/unit_tests/service/test_interface.py index ac756fe851..f9050851cd 100644 --- a/tests/unit_tests/service/test_interface.py +++ b/tests/unit_tests/service/test_interface.py @@ -386,7 +386,6 @@ def test_get_task_by_id( params={}, metadata=expected_metadata, ), - is_complete=False, is_pending=True, errors=[], ) @@ -416,7 +415,6 @@ def test_submit_task_inserts_metadata(context_mock: MagicMock): params={}, metadata=metadata, ), - is_complete=False, is_pending=True, errors=[], ) diff --git a/tests/unit_tests/service/test_rest_api.py b/tests/unit_tests/service/test_rest_api.py index dedd9d7ff3..fe73240601 100644 --- a/tests/unit_tests/service/test_rest_api.py +++ b/tests/unit_tests/service/test_rest_api.py @@ -43,7 +43,7 @@ WorkerTask, ) from blueapi.service.runner import WorkerDispatcher -from blueapi.worker.event import WorkerState +from blueapi.worker.event import TaskResult, WorkerState from blueapi.worker.task import Task from blueapi.worker.task_worker import TrackableTask @@ -371,7 +371,7 @@ def test_put_plan_fails_if_not_idle(mock_runner: Mock, client: TestClient) -> No # Set to non idle mock_runner.run.return_value = TrackableTask( - task=Task(name="none"), task_id=task_id_current, is_complete=False + task=Task(name="none"), task_id=task_id_current ) resp = client.put("/worker/task", json={"task_id": task_id_new}) @@ -386,7 +386,6 @@ def test_get_tasks(mock_runner: Mock, client: TestClient) -> None: TrackableTask( task_id="1", task=Task(name="first_task"), - is_complete=False, is_pending=True, ), ] @@ -433,7 +432,7 @@ def test_get_tasks_by_status(mock_runner: Mock, client: TestClient) -> None: TrackableTask( task_id="3", task=Task(name="third_task"), - is_complete=True, + outcome=TaskResult.from_result(42), is_pending=False, ), ] @@ -453,7 +452,7 @@ def test_get_tasks_by_status(mock_runner: Mock, client: TestClient) -> None: "params": {}, "metadata": {}, }, - "outcome": None, + "outcome": {"outcome": "success", "type": "int", "result": 42}, "task_id": "3", } ] @@ -534,7 +533,7 @@ def test_set_active_task_active_task_complete( mock_runner.run.return_value = TrackableTask( task_id="1", task=Task(name="a_completed_task"), - is_complete=True, + outcome=TaskResult.from_result(42), is_pending=False, ) @@ -553,7 +552,6 @@ def test_set_active_task_worker_already_running( mock_runner.run.return_value = TrackableTask( task_id="1", task=Task(name="a_running_task"), - is_complete=False, is_pending=False, ) diff --git a/tests/unit_tests/worker/test_task_worker.py b/tests/unit_tests/worker/test_task_worker.py index 7f9e8809de..fa34057755 100644 --- a/tests/unit_tests/worker/test_task_worker.py +++ b/tests/unit_tests/worker/test_task_worker.py @@ -3,7 +3,7 @@ import threading import time from collections.abc import Callable, Iterable -from concurrent.futures import Future +from concurrent.futures import Future, ThreadPoolExecutor from pathlib import Path from queue import Full from typing import Any, TypeVar @@ -13,6 +13,7 @@ import pydantic import pytest from bluesky.protocols import Movable, Readable, Status +from bluesky.run_engine import RunEngineResult from bluesky.utils import MsgGenerator from dodal.common import inject from dodal.common.types import UpdatingPathProvider @@ -38,6 +39,7 @@ WorkerState, ) from blueapi.worker.event import TaskError, TaskResult, TaskStatusEnum +from blueapi.worker.task_worker import CancelSignal _SIMPLE_TASK = Task(name="sleep", params={"time": 0.0}) _LONG_TASK = Task(name="sleep", params={"time": 1.0}) @@ -445,6 +447,14 @@ def raise_full(item): def test_metadata_passed_to_context(context: BlueskyContext): context.run_engine = Mock() context.run_engine.md = {} + context.run_engine.return_value = RunEngineResult( + run_start_uids=(), + plan_result=None, + exit_status="success", + interrupted=False, + reason="", + exception=None, + ) _TASK_WITH_METADATA.do_task(context) for metadata in [("foo", "bar"), ("baz", 0)]: assert metadata in context.run_engine.md.items() @@ -664,21 +674,18 @@ def callback(unused: Future[list[Any]], stream=stream, sub=sub): ], ) def test_get_tasks(worker: TaskWorker, status, expected_task_ids): + # A task that has started is tracked via _current, not _pending_tasks. + worker._current = TrackableTask( + task_id="task1", + task=Task(name="set_absolute", params={"movable": "fake_device", "value": 4.0}), + is_pending=False, + ) worker._pending_tasks = { - "task1": TrackableTask( - task_id="task1", - task=Task( - name="set_absolute", params={"movable": "fake_device", "value": 4.0} - ), - is_complete=False, - is_pending=False, - ), "task2": TrackableTask( task_id="task2", task=Task( name="set_absolute", params={"movable": "fake_device", "value": 4.0} ), - is_complete=False, is_pending=True, ), } @@ -688,7 +695,7 @@ def test_get_tasks(worker: TaskWorker, status, expected_task_ids): task=Task( name="set_absolute", params={"movable": "fake_device", "value": 4.0} ), - is_complete=True, + outcome=TaskResult.from_result(None), is_pending=False, ), } @@ -703,7 +710,10 @@ def test_submitting_completed_task_fails(worker: TaskWorker): with pytest.raises(ValueError): worker._submit_trackable_task( TrackableTask( - task_id="task1", task=_SIMPLE_TASK, is_complete=True, is_pending=False + task_id="task1", + task=_SIMPLE_TASK, + outcome=TaskResult.from_result(None), + is_pending=False, ) ) @@ -838,10 +848,6 @@ def test_cycle_without_otel_context(mock_logger: Mock, inert_worker: TaskWorker) inert_worker._cycle() assert inert_worker._current_task_otel_context is None - # Bad way to tell that this branch has been run, but I can't think of a better way - # Have to set these values to match output - task.is_complete = False - task.is_pending = True mock_logger.info.assert_called_with( "Task ran successfully - returned: %s", None, extra={"task_id": "0"} ) @@ -973,14 +979,20 @@ def test_pause_does_not_publish_error_event(worker: TaskWorker) -> None: assert all(e.errors == [] for e in events) -def test_paused_task_remains_in_pending(worker: TaskWorker) -> None: +def test_paused_task_tracked_via_current_not_pending(worker: TaskWorker) -> None: + # A task leaves _pending_tasks as soon as it starts - self._current is + # the sole source of truth for it from then on, paused or not. 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._pending_tasks assert task_id not in worker._completed_tasks + current = worker._current + assert current is not None + assert current.task_id == task_id + assert worker.get_task_by_id(task_id) is current def test_worker_state_is_paused(worker: TaskWorker) -> None: @@ -1029,7 +1041,7 @@ def test_cancel_active_task_abort(worker: TaskWorker) -> None: 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].errors == ["Task aborted"] assert events[-1].task_status is not None assert events[-1].task_status.task_failed assert task_id in worker._completed_tasks @@ -1137,10 +1149,9 @@ def test_cancel_active_task_does_not_overwrite_existing_outcome( assert current.outcome.message == "worker thread got there first" -def test_cancel_active_task_does_not_block_caller_when_paused( +def test_cancel_active_task_blocks_caller_until_paused_cleanup_finishes( worker: TaskWorker, ) -> None: - worker._ctx.register_plan(pausing_plan_with_slow_cleanup) task_id = worker.submit_task(_PAUSING_TASK_SLOW_CLEANUP) @@ -1150,50 +1161,39 @@ def test_cancel_active_task_does_not_block_caller_when_paused( 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(), + assert elapsed >= _ABORT_CLEANUP_DELAY, ( + f"cancel_active_task(failure=True) returned after only {elapsed:.2f}s, " + f"before the RunEngine's {_ABORT_CLEANUP_DELAY}s abort cleanup could " + "have finished." ) - 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 + assert worker._completed_tasks[task_id].outcome is not None + assert isinstance(worker._completed_tasks[task_id].outcome, TaskError) -def test_cancel_wins_race_with_concurrent_resume(worker: TaskWorker) -> None: - # A cancel queued right after a resume must still win, not be dropped in - # favour of the resume completing the task first. - worker._ctx.register_plan(pausing_plan) - task_id = worker.submit_task(_PAUSING_TASK) +def test_concurrent_resume_and_cancel_do_not_corrupt_state( + 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) - 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) + with ThreadPoolExecutor(1) as executor: + resume_future = executor.submit(worker.resume) + time.sleep(_ABORT_CLEANUP_DELAY / 4) + worker.cancel_active_task(failure=True, reason="changed my mind") + resume_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 + completed = worker._completed_tasks[task_id] + assert completed.outcome is not None -def test_second_cancel_while_paused_supersedes_first(worker: TaskWorker) -> None: - +def test_second_cancel_while_first_still_queued_supersedes_it( + worker: TaskWorker, +) -> None: worker._ctx.register_plan(pausing_plan) task_id = worker.submit_task(_PAUSING_TASK) @@ -1203,8 +1203,11 @@ def test_second_cancel_while_paused_supersedes_first(worker: TaskWorker) -> None 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") + first_signal = CancelSignal(failure=False, reason="first, graceful") + worker._pending_cancel = first_signal + worker._task_channel.put_nowait(first_signal) + second_signal = CancelSignal(failure=True, reason="second, urgent") + worker._pending_cancel = second_signal events = complete_future.result(timeout=5.0) assert events[-1].task_status is not None @@ -1261,7 +1264,7 @@ def test_resume_re_pauses_when_plan_pauses_again(worker: TaskWorker) -> None: current = worker._current assert current is not None assert current.outcome is None - assert task_id in worker._pending_tasks + assert task_id not in worker._pending_tasks assert task_id not in worker._completed_tasks From 6375d4cb6782885d539ee3de170f02a49b9af81d Mon Sep 17 00:00:00 2001 From: alexj9837 <52531949+Alexj9837@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:00:19 +0000 Subject: [PATCH 09/10] fix: cancel_active_task() waited on RunEngine state instead of the actual completion it needed, causing a race where callers woke up before _completed_tasks was updated. --- src/blueapi/worker/task_worker.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/src/blueapi/worker/task_worker.py b/src/blueapi/worker/task_worker.py index cea92ca325..6ba8a56361 100644 --- a/src/blueapi/worker/task_worker.py +++ b/src/blueapi/worker/task_worker.py @@ -113,12 +113,8 @@ class TaskWorker: _errors: list[str] _warnings: list[str] - # The queue is actually a channel between 2 threads - # most programming languages have a separate abstraction for this - # but Python reuses Queue - # So it's not used as a standard queue, - # but as a box in which to put the "current" task and nothing else - # So the calling thread can only ever submit one plan at a time. + # A channel to the worker thread, not really a queue - maxsize=1 makes it + # a single-item box, so only one plan can be in flight at a time. _task_channel: Queue # type: ignore _current: TrackableTask | None @@ -212,7 +208,7 @@ def cancel_active_task( with self._state_change: self._task_channel.put(AbortSignal(failure, reason)) if not self._state_change.wait_for( - lambda: self._state == WorkerState.IDLE, + lambda: current.task_id in self._completed_tasks, timeout=self._start_stop_timeout, ): raise TimeoutError( @@ -318,7 +314,6 @@ def submit_task(self, task: Task) -> str: task_id: str = str(uuid.uuid4()) add_span_attributes({"TaskId": task_id}) request_id = get_baggage("correlation_id") - # If request id is not a string, we do not pass it into a TrackableTask if not isinstance(request_id, str): LOGGER.warning(f"Invalid correlation id detected: {request_id}") request_id = None @@ -387,7 +382,6 @@ def stop(self) -> None: """ LOGGER.info("Attempting to stop worker") - # If the worker has not yet started there is nothing to do. if self._started.is_set(): self._task_channel.put(KillSignal()) else: From ce4ea4cddbe5f71211333ae4177bd51ff40a9d5b Mon Sep 17 00:00:00 2001 From: alexj9837 <52531949+Alexj9837@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:11:51 +0000 Subject: [PATCH 10/10] fix: cancel_active_task()'s immediate path returned before the worker thread finished writing to _completed_tasks, causingflaky failures , now waits on that instead of returning right after abort()/stop() --- src/blueapi/worker/task_worker.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/src/blueapi/worker/task_worker.py b/src/blueapi/worker/task_worker.py index 6ba8a56361..6fe6e64684 100644 --- a/src/blueapi/worker/task_worker.py +++ b/src/blueapi/worker/task_worker.py @@ -218,6 +218,9 @@ def cancel_active_task( else: # The worker thread's own do_task() may be finishing concurrently # on another thread - don't clobber whatever it already recorded. + # abort()/stop() must be called without holding _state_change: + # _on_state_change() needs that same lock, and abort()/stop() + # don't return until its state-change callbacks have run. try: if failure: with self._status_lock: @@ -230,7 +233,16 @@ def cancel_active_task( current.outcome = TaskResult.from_result(None) self._ctx.run_engine.stop() except TransitionError: - return current.task_id + pass + with self._state_change: + if not self._state_change.wait_for( + lambda: current.task_id in self._completed_tasks, + timeout=self._start_stop_timeout, + ): + raise TimeoutError( + "Worker did not finish cancelling within " + f"{self._start_stop_timeout} seconds" + ) return current.task_id @start_as_current_span(TRACER, "task_id")