fix(task-handler): warn on shutdown timeout and close cancellation races

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:

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

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"""

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()