From b80c1adb628edda4bb185ef84c6f3cacf313fa7f Mon Sep 17 00:00:00 2001 From: wakonig_k Date: Tue, 29 Sep 2026 10:26:57 +0200 Subject: [PATCH] fix(task-handler): warn on shutdown timeout and close cancellation races --- ophyd_devices/utils/psi_device_base_utils.py | 89 +++++++----- tests/test_psi_device_base.py | 28 ++++ tests/test_simulation.py | 32 +++++ tests/test_utils.py | 141 ++++++++++++++++++- 4 files changed, 251 insertions(+), 39 deletions(-) diff --git a/ophyd_devices/utils/psi_device_base_utils.py b/ophyd_devices/utils/psi_device_base_utils.py index f55c1be..10fe804 100644 --- a/ophyd_devices/utils/psi_device_base_utils.py +++ b/ophyd_devices/utils/psi_device_base_utils.py @@ -20,6 +20,7 @@ from ophyd.status import DeviceStatus as _DeviceStatus from ophyd.status import MoveStatus as _MoveStatus from ophyd.status import Status as _Status from ophyd.status import StatusBase as _StatusBase +from ophyd.status import WaitTimeoutError from ophyd.utils import StatusTimeoutError if TYPE_CHECKING: # pragma: no cover @@ -762,8 +763,8 @@ class TaskHandler: self._tasks = {} self._parent = parent self._lock = threading.RLock() - self._shutdown_lock = threading.RLock() self._shutting_down = False + self._cancelling_pending = set() self._entered_tasks = set() self._cancellable_tasks = set() self._cancellation_requested = set() @@ -808,6 +809,9 @@ class TaskHandler: if self._shutting_down: raise RuntimeError("Cannot start a task during shutdown.") thread = self._tasks[task_status.task_id][1] + if task_status.task_id in self._cancelling_pending: + logger.warning(f"Task with ID {task_status.task_id} is being cancelled.") + return if thread.ident is not None: logger.warning(f"Task with ID {task_status.task_id} was already started.") return @@ -871,26 +875,33 @@ class TaskHandler: def kill_task(self, task_status: TaskStatus) -> None: """Cancel a pending task or request cancellation of a running task.""" + task_id = task_status.task_id with self._lock: - task_info = self._tasks.get(task_status.task_id) - if task_info is None: + task_info = self._tasks.get(task_id) + if task_info is None or task_id in self._cancelling_pending: return thread = task_info[1] if thread.ident is not None: self._cancel_running_task_locked(task_status, thread) return - self._tasks.pop(task_status.task_id) + self._cancelling_pending.add(task_id) task_status.state = TaskState.KILLED - # Completing the status may invoke callbacks that reenter the handler. - task_status.set_exception(TaskKilledError(f"Task {task_status.task_id} was killed.")) + # Keep the pending task visible until its status is resolved. Callbacks + # run inside set_exception and may reenter the handler. + try: + task_status.set_exception(TaskKilledError(f"Task {task_id} was killed.")) + finally: + with self._lock: + self._tasks.pop(task_id, None) + self._cancelling_pending.discard(task_id) def _cancel_running_task_locked( self, task_status: TaskStatus, thread: threading.Thread ) -> None: """Request cancellation of a started worker while ``self._lock`` is held.""" task_id = task_status.task_id - if not thread.is_alive() or task_status.done or task_id in self._cancellation_requested: + if not thread.is_alive() or task_status.done: return # A started thread may be before its wrapper, inside the callable, or @@ -905,7 +916,6 @@ class TaskHandler: try: result = set_async_exc(ident, ctypes.py_object(TaskKilledError)) if result == 1: - self._cancellation_requested.add(task_id) return if result > 1: set_async_exc(ident, None) @@ -918,54 +928,63 @@ class TaskHandler: logger.warning(f"Exception raised while killing thread {ident}: {exc}") def shutdown(self, timeout: float | None = 5.0) -> None: - """Cancel tasks and wait for active threads to exit. + """Cancel tasks and wait briefly for active threads to exit. - Raises TimeoutError if an active task does not exit within ``timeout`` seconds. - The handler can be used again after shutdown. + Warn if a task remains active after ``timeout`` seconds. The handler can + be used again after shutdown. """ if timeout is not None and timeout < 0: raise ValueError("Shutdown timeout must be non-negative or None.") - with self._shutdown_lock: - tasks = self._begin_shutdown() - try: - deadline = None if timeout is None else time.monotonic() + timeout - for task_status, _ in tasks: + tasks = self._begin_shutdown() + if tasks is None: + return + try: + deadline = None if timeout is None else time.monotonic() + timeout + current_thread = threading.current_thread() + for task_status, thread in tasks: + if thread is not current_thread: self.kill_task(task_status) - self._wait_for_tasks(tasks, deadline) - finally: - with self._lock: - self._shutting_down = False + self._wait_for_tasks(tasks, deadline) + finally: + with self._lock: + self._shutting_down = False - def _begin_shutdown(self) -> tuple[tuple[TaskStatus, threading.Thread], ...]: + def _begin_shutdown(self) -> tuple[tuple[TaskStatus, threading.Thread], ...] | None: """Snapshot tasks and prevent new work while shutdown is in progress.""" with self._lock: if self._shutting_down: - raise RuntimeError("Task handler shutdown is already in progress.") + logger.warning("Task handler shutdown is already in progress.") + return None tasks = tuple(self._tasks.values()) - if any(thread is threading.current_thread() for _, thread in tasks): - raise RuntimeError("A task cannot shut down its own handler.") self._shutting_down = True return tasks def _wait_for_tasks( self, tasks: tuple[tuple[TaskStatus, threading.Thread], ...], deadline: float | None ) -> None: - """Wait for started workers and report any that survive the deadline.""" - # Worker cleanup needs self._lock, so joins must happen outside that lock. - for _, thread in tasks: - if thread.ident is None: + """Wait for started workers and pending status cancellations.""" + # Worker cleanup needs self._lock, so waits must happen outside that lock. + current_thread = threading.current_thread() + for status, thread in tasks: + if thread is current_thread: continue remaining = None if deadline is None else max(0, deadline - time.monotonic()) - thread.join(remaining) + if thread.ident is None: + try: + status.exception(timeout=remaining) + except WaitTimeoutError: + pass + else: + thread.join(remaining) alive = [status.task_id for status, thread in tasks if thread.is_alive()] - if not alive: - return - with self._lock: - self._cancellation_requested.difference_update( - task_id for task_id in alive if task_id in self._cancellable_tasks + pending = [ + status.task_id for status, thread in tasks if thread.ident is None and not status.done + ] + if alive or pending: + logger.warning( + f"Task handler shutdown left running threads {alive} and pending statuses {pending}." ) - raise TimeoutError(f"Task threads did not stop during shutdown: {', '.join(alive)}") class FileHandler: diff --git a/tests/test_psi_device_base.py b/tests/test_psi_device_base.py index de95402..0abefcf 100644 --- a/tests/test_psi_device_base.py +++ b/tests/test_psi_device_base.py @@ -113,6 +113,34 @@ def test_psi_device_base_destroy_with_pending_task(device): assert isinstance(status.exception(), TaskKilledError) +def test_psi_device_base_destroy_after_task_shutdown_timeout(device): + """A blocked task must not prevent base-device destruction.""" + started = threading.Event() + release = threading.Event() + + def blocked_task(): + started.set() + release.wait() + + status = device.task_handler.submit_task(blocked_task) + worker = device.task_handler._tasks[status.task_id][1] + assert started.wait(1) + real_shutdown = device.task_handler.shutdown + try: + with mock.patch.object( + device.task_handler, "shutdown", side_effect=lambda: real_shutdown(timeout=0.01) + ): + with mock.patch( + "ophyd_devices.utils.psi_device_base_utils.set_async_exc", return_value=1 + ): + device.destroy() + assert device.destroyed + finally: + release.set() + worker.join(timeout=1) + assert not worker.is_alive() + + def test_psi_device_base_wait_for_signals(device_positioner): """Test wait_for_signals method""" device: SimPositionerDevice = device_positioner diff --git a/tests/test_simulation.py b/tests/test_simulation.py index f3ea908..3f49f32 100644 --- a/tests/test_simulation.py +++ b/tests/test_simulation.py @@ -463,6 +463,38 @@ def test_linear_traj(initial_position, final_position, max_velocity, acceleratio assert trajectory.ended +def test_sim_camera_stop_after_task_shutdown_timeout(camera): + """A blocked task must not prevent the camera from completing stop cleanup.""" + started = threading.Event() + release = threading.Event() + + def blocked_task(): + started.set() + release.wait() + + status = camera.task_handler.submit_task(blocked_task) + worker = camera.task_handler._tasks[status.task_id][1] + assert started.wait(1) + real_shutdown = camera.task_handler.shutdown + try: + with mock.patch.object( + camera.task_handler, "shutdown", side_effect=lambda: real_shutdown(timeout=0.01) + ): + with mock.patch( + "ophyd_devices.utils.psi_device_base_utils.set_async_exc", return_value=1 + ): + with mock.patch( + "ophyd_devices.utils.psi_device_base_utils.logger.warning" + ) as warning: + camera.stop() + assert camera.stopped + assert any(status.task_id in str(call) for call in warning.call_args_list) + finally: + release.set() + worker.join(timeout=1) + assert not worker.is_alive() + + @pytest.mark.parametrize("proxy_active", [True, False]) def test_sim_camera_proxies(camera, proxy_active): """Test mocking compute_method with framework class""" diff --git a/tests/test_utils.py b/tests/test_utils.py index d08f718..c401933 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -320,8 +320,63 @@ def test_utils_task_handler_shutdown_waits_for_running_task(task_handler): assert status.state == TaskState.COMPLETED +def test_utils_task_handler_shutdown_from_completion_callback(task_handler): + """A callback reentering shutdown must not block the worker being joined.""" + started = threading.Event() + release = threading.Event() + cancellation_requested = threading.Event() + callback_done = threading.Event() + shutdown_done = threading.Event() + outcome = {} + + def task(): + started.set() + release.wait() + + status = task_handler.submit_task(task) + worker = task_handler._tasks[status.task_id][1] + assert started.wait(1) + + def callback(_): + task_handler.shutdown(timeout=0.1) + callback_done.set() + + status.add_callback(callback) + + def stop_handler(): + try: + task_handler.shutdown(timeout=2) + except Exception as exc: # pylint: disable=broad-except + outcome["exception"] = exc + finally: + shutdown_done.set() + + def record_cancellation(*_): + cancellation_requested.set() + return 1 + + stopper = threading.Thread(target=stop_handler, daemon=True) + with mock.patch( + "ophyd_devices.utils.psi_device_base_utils.set_async_exc", side_effect=record_cancellation + ): + stopper.start() + try: + assert cancellation_requested.wait(1) + release.set() + assert callback_done.wait(1) + assert shutdown_done.wait(1) + finally: + release.set() + stopper.join(timeout=3) + worker.join(timeout=1) + + assert "exception" not in outcome + assert not worker.is_alive() + assert status.done + + def test_utils_task_handler_shutdown_timeout(task_handler): - """A blocked worker causes an explicit timeout instead of a successful return.""" + """A blocked worker produces a warning without preventing shutdown.""" started = threading.Event() release = threading.Event() @@ -334,9 +389,10 @@ def test_utils_task_handler_shutdown_timeout(task_handler): assert started.wait(1) try: with mock.patch("ophyd_devices.utils.psi_device_base_utils.set_async_exc", return_value=1): - with pytest.raises(TimeoutError, match="Task threads did not stop"): + with mock.patch("ophyd_devices.utils.psi_device_base_utils.logger.warning") as warning: task_handler.shutdown(timeout=0.01) assert worker.is_alive() + assert any(status.task_id in str(call) for call in warning.call_args_list) finally: release.set() worker.join(timeout=1) @@ -362,8 +418,7 @@ def test_utils_task_handler_can_retry_cancellation_after_timeout(task_handler): worker = task_handler._tasks[status.task_id][1] assert started.wait(1) try: - with pytest.raises(TimeoutError, match="Task threads did not stop"): - task_handler.shutdown(timeout=0.01) + task_handler.shutdown(timeout=0.01) first_wait.set() assert first_caught.wait(1) task_handler.kill_task(status) @@ -377,6 +432,84 @@ def test_utils_task_handler_can_retry_cancellation_after_timeout(task_handler): assert isinstance(status.exception(), TaskKilledError) +def test_utils_task_handler_can_retry_direct_cancellation(task_handler): + """A second kill request reaches a task that caught the first one.""" + first_wait = threading.Event() + second_wait = threading.Event() + started = threading.Event() + first_caught = threading.Event() + + def task(): + started.set() + try: + first_wait.wait() + except TaskKilledError: + first_caught.set() + second_wait.wait() + + status = task_handler.submit_task(task) + worker = task_handler._tasks[status.task_id][1] + assert started.wait(1) + try: + task_handler.kill_task(status) + first_wait.set() + assert first_caught.wait(1) + task_handler.kill_task(status) + finally: + second_wait.set() + worker.join(timeout=1) + + assert not worker.is_alive() + assert status.done + assert status.state == TaskState.KILLED + assert isinstance(status.exception(), TaskKilledError) + + +def test_utils_task_handler_shutdown_waits_for_pending_cancellation(task_handler): + """Shutdown observes a queued status that another thread is cancelling.""" + completion_started = threading.Event() + finish_completion = threading.Event() + shutdown_done = threading.Event() + outcome = {} + ran = threading.Event() + status = task_handler.submit_task(ran.set, run=False) + original_set_exception = status.set_exception + + def delayed_set_exception(exc): + completion_started.set() + finish_completion.wait() + original_set_exception(exc) + + def stop_handler(): + try: + task_handler.shutdown(timeout=1) + except Exception as exc: # pylint: disable=broad-except + outcome["exception"] = exc + finally: + shutdown_done.set() + + with mock.patch.object(status, "set_exception", side_effect=delayed_set_exception): + canceller = threading.Thread(target=task_handler.kill_task, args=(status,), daemon=True) + canceller.start() + assert completion_started.wait(1) + task_handler.start_task(status) + assert not ran.is_set() + stopper = threading.Thread(target=stop_handler, daemon=True) + stopper.start() + try: + assert not shutdown_done.wait(0.05) + finally: + finish_completion.set() + canceller.join(timeout=1) + stopper.join(timeout=2) + + assert shutdown_done.is_set() + assert "exception" not in outcome + assert status.done + assert status.state == TaskState.KILLED + assert not ran.is_set() + + def test_utils_task_handler_shutdown_rejects_callback_submission(task_handler): """A completion callback cannot add new work during shutdown.""" ran = threading.Event()