From 11497bcd4d3c8bccc627258cdc63cfe7b9e782bf Mon Sep 17 00:00:00 2001
From: wakonig_k <klaus.wakonig@psi.ch>
Date: Tue, 29 Sep 2026 10:06:51 +0200
Subject: [PATCH] fix(task-handler): handle shutdown and cancellation races

---
 ophyd_devices/utils/psi_device_base_utils.py | 153 ++++++++++---
 tests/test_psi_device_base.py                |  15 ++
 tests/test_utils.py                          | 219 +++++++++++++++++++
 3 files changed, 358 insertions(+), 29 deletions(-)

diff --git a/ophyd_devices/utils/psi_device_base_utils.py b/ophyd_devices/utils/psi_device_base_utils.py
index 0cd20bf..f55c1be 100644
--- a/ophyd_devices/utils/psi_device_base_utils.py
+++ b/ophyd_devices/utils/psi_device_base_utils.py
@@ -5,6 +5,7 @@ from __future__ import annotations
 import ctypes
 import operator
 import threading
+import time
 import traceback
 import uuid
 from enum import Enum
@@ -761,6 +762,11 @@ class TaskHandler:
         self._tasks = {}
         self._parent = parent
         self._lock = threading.RLock()
+        self._shutdown_lock = threading.RLock()
+        self._shutting_down = False
+        self._entered_tasks = set()
+        self._cancellable_tasks = set()
+        self._cancellation_requested = set()
 
     def submit_task(
         self,
@@ -784,9 +790,12 @@ class TaskHandler:
             name=f"task {task_status.task_id}",
             daemon=True,
         )
-        self._tasks.update({task_status.task_id: (task_status, thread)})
-        if run is True:
-            self.start_task(task_status)
+        with self._lock:
+            if self._shutting_down:
+                raise RuntimeError("Cannot submit a task during shutdown.")
+            self._tasks[task_status.task_id] = (task_status, thread)
+            if run is True:
+                self.start_task(task_status)
         return task_status
 
     def start_task(self, task_status: TaskStatus) -> None:
@@ -795,19 +804,31 @@ class TaskHandler:
         Args:
             task_status: The task status object.
         """
-        thread = self._tasks[task_status.task_id][1]
-        if thread.is_alive():
-            logger.warning(f"Task with ID {task_status.task_id} is already running.")
-            return
-        task_status.state = TaskState.RUNNING
-        thread.start()
+        with self._lock:
+            if self._shutting_down:
+                raise RuntimeError("Cannot start a task during shutdown.")
+            thread = self._tasks[task_status.task_id][1]
+            if thread.ident is not None:
+                logger.warning(f"Task with ID {task_status.task_id} was already started.")
+                return
+            task_status.state = TaskState.RUNNING
+            thread.start()
 
     def _wrap_task(
         self, task: Callable, task_args: tuple, task_kwargs: dict, task_status: TaskStatus
     ):
         """Wrap the task in a function"""
         try:
-            task(*task_args, **task_kwargs)
+            with self._lock:
+                self._entered_tasks.add(task_status.task_id)
+                if task_status.task_id in self._cancellation_requested:
+                    raise TaskKilledError()
+                self._cancellable_tasks.add(task_status.task_id)
+            try:
+                task(*task_args, **task_kwargs)
+            finally:
+                with self._lock:
+                    self._cancellable_tasks.discard(task_status.task_id)
         except TimeoutError as exc:
             content = traceback.format_exc()
             logger.warning(
@@ -820,8 +841,7 @@ class TaskHandler:
             task_status.set_exception(exc)
         except TaskKilledError as exc:
             exc = exc.__class__(
-                f"Task {task_status.task_id} was killed. ThreadID:"
-                f" {self._tasks[task_status.task_id][1].ident}"
+                f"Task {task_status.task_id} was killed. ThreadID: {threading.get_ident()}"
             )
             content = traceback.format_exc()
             logger.warning(
@@ -845,32 +865,107 @@ class TaskHandler:
         finally:
             with self._lock:
                 self._tasks.pop(task_status.task_id, None)
+                self._entered_tasks.discard(task_status.task_id)
+                self._cancellable_tasks.discard(task_status.task_id)
+                self._cancellation_requested.discard(task_status.task_id)
 
     def kill_task(self, task_status: TaskStatus) -> None:
-        """Kill the thread
+        """Cancel a pending task or request cancellation of a running task."""
+        with self._lock:
+            task_info = self._tasks.get(task_status.task_id)
+            if task_info is None:
+                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)
+            task_status.state = TaskState.KILLED
 
-        task_status: The task status object.
-        """
-        thread = self._tasks[task_status.task_id][1]
-        exception_cls = TaskKilledError
+        # Completing the status may invoke callbacks that reenter the handler.
+        task_status.set_exception(TaskKilledError(f"Task {task_status.task_id} was killed."))
+
+    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:
+            return
+
+        # A started thread may be before its wrapper, inside the callable, or
+        # completing its status. Only the callable can be interrupted safely.
+        if task_id not in self._entered_tasks:
+            self._cancellation_requested.add(task_id)
+            return
+        if task_id not in self._cancellable_tasks:
+            return
 
         ident = ctypes.c_long(thread.ident)
-        exc = ctypes.py_object(exception_cls)
         try:
-            res = set_async_exc(ident, exc)
-            if res == 0:
-                raise ValueError("Invalid thread ID")
-            if res > 1:
+            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)
-                logger.warning(f"Exception raise while kille Thread {ident}; return value: {res}")
-        except Exception as e:  # pylint: disable=broad-except
-            logger.warning(f"Exception raised while killing thread {ident}: {e}")
+                logger.warning(
+                    f"Could not cancel task {task_id}: exception was raised in {result} threads."
+                )
+                return
+            logger.warning(f"Could not cancel task {task_id}: invalid thread ID.")
+        except Exception as exc:  # pylint: disable=broad-except
+            logger.warning(f"Exception raised while killing thread {ident}: {exc}")
 
-    def shutdown(self):
-        """Shutdown all tasks of task handler"""
+    def shutdown(self, timeout: float | None = 5.0) -> None:
+        """Cancel tasks and wait 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.
+        """
+        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:
+                    self.kill_task(task_status)
+                self._wait_for_tasks(tasks, deadline)
+            finally:
+                with self._lock:
+                    self._shutting_down = False
+
+    def _begin_shutdown(self) -> tuple[tuple[TaskStatus, threading.Thread], ...]:
+        """Snapshot tasks and prevent new work while shutdown is in progress."""
         with self._lock:
-            for info in self._tasks.values():
-                self.kill_task(info[0])
+            if self._shutting_down:
+                raise RuntimeError("Task handler shutdown is already in progress.")
+            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:
+                continue
+            remaining = None if deadline is None else max(0, deadline - time.monotonic())
+            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
+            )
+        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 7b91e4c..de95402 100644
--- a/tests/test_psi_device_base.py
+++ b/tests/test_psi_device_base.py
@@ -16,6 +16,7 @@ from ophyd_devices.sim.sim_camera import SimCamera
 from ophyd_devices.sim.sim_positioner import SimPositioner
 from ophyd_devices.tests.utils import get_mock_scan_info
 from ophyd_devices.utils.bec_signals import FileEventSignal, PreviewSignal, ProgressSignal
+from ophyd_devices.utils.psi_device_base_utils import TaskKilledError, TaskState
 
 # pylint: disable=redefined-outer-name
 # pylint: disable=protected-access
@@ -98,6 +99,20 @@ def device():
     yield SimDevice(name="device", prefix="test:")
 
 
+def test_psi_device_base_destroy_with_pending_task(device):
+    """A queued task must not prevent device destruction."""
+    ran = threading.Event()
+    status = device.task_handler.submit_task(ran.set, run=False)
+
+    device.destroy()
+
+    assert device.destroyed
+    assert not ran.is_set()
+    assert status.done
+    assert status.state == TaskState.KILLED
+    assert isinstance(status.exception(), TaskKilledError)
+
+
 def test_psi_device_base_wait_for_signals(device_positioner):
     """Test wait_for_signals method"""
     device: SimPositionerDevice = device_positioner
diff --git a/tests/test_utils.py b/tests/test_utils.py
index a25bb35..d08f718 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -234,6 +234,225 @@ def test_utils_task_handler_shutdown(task_handler):
     assert status1.exception().__class__ == TaskKilledError
 
 
+def test_utils_task_handler_kill_pending_task(task_handler):
+    """A task that has not started can be cancelled without running it."""
+    ran = threading.Event()
+    status = task_handler.submit_task(ran.set, run=False)
+
+    task_handler.kill_task(status)
+
+    assert not ran.is_set()
+    assert status.done
+    assert status.state == TaskState.KILLED
+    assert isinstance(status.exception(), TaskKilledError)
+    assert status.task_id not in task_handler._tasks
+    task_handler.kill_task(status)
+
+
+def test_utils_task_handler_shutdown_pending_task_and_reuse(task_handler):
+    """Shutdown cancels queued tasks and leaves the handler ready for later work."""
+    ran = threading.Event()
+    status = task_handler.submit_task(ran.set, run=False)
+
+    task_handler.shutdown()
+
+    assert not ran.is_set()
+    assert status.done
+    assert status.state == TaskState.KILLED
+    assert isinstance(status.exception(), TaskKilledError)
+    assert status.task_id not in task_handler._tasks
+    task_handler.shutdown()
+
+    next_status = task_handler.submit_task(ran.set)
+    next_status.wait(timeout=1)
+    assert ran.is_set()
+    assert next_status.state == TaskState.COMPLETED
+
+
+def test_utils_task_handler_shutdown_waits_for_running_task(task_handler):
+    """A successful shutdown returns after its running worker has exited."""
+    started = threading.Event()
+    release = threading.Event()
+    cancellation_requested = 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 stop_handler():
+        try:
+            task_handler.shutdown(timeout=1)
+            outcome["snapshot"] = (
+                worker.is_alive(),
+                status.done,
+                status.task_id in task_handler._tasks,
+            )
+        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)
+            assert not shutdown_done.wait(0.05)
+        finally:
+            release.set()
+            assert shutdown_done.wait(2)
+            stopper.join(timeout=1)
+
+    assert "exception" not in outcome
+    assert outcome["snapshot"] == (False, True, False)
+    assert status.state == TaskState.COMPLETED
+
+
+def test_utils_task_handler_shutdown_timeout(task_handler):
+    """A blocked worker causes an explicit timeout instead of a successful return."""
+    started = threading.Event()
+    release = threading.Event()
+
+    def task():
+        started.set()
+        release.wait()
+
+    status = task_handler.submit_task(task)
+    worker = task_handler._tasks[status.task_id][1]
+    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"):
+                task_handler.shutdown(timeout=0.01)
+        assert worker.is_alive()
+    finally:
+        release.set()
+        worker.join(timeout=1)
+    assert not worker.is_alive()
+
+
+def test_utils_task_handler_can_retry_cancellation_after_timeout(task_handler):
+    """A task that catches one cancellation can be cancelled again."""
+    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:
+        with pytest.raises(TimeoutError, match="Task threads did not stop"):
+            task_handler.shutdown(timeout=0.01)
+        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_rejects_callback_submission(task_handler):
+    """A completion callback cannot add new work during shutdown."""
+    ran = threading.Event()
+    rejected = threading.Event()
+    status = task_handler.submit_task(lambda: None, run=False)
+
+    def submit_from_callback(_):
+        try:
+            task_handler.submit_task(ran.set)
+        except RuntimeError:
+            rejected.set()
+
+    status.add_callback(submit_from_callback)
+    task_handler.shutdown()
+
+    assert rejected.wait(1)
+    assert not ran.is_set()
+    assert not task_handler._tasks
+
+
+def test_utils_task_handler_cancel_before_wrapper_starts(task_handler):
+    """Cancellation between thread start and wrapper entry still resolves the status."""
+    wrapper_entered = threading.Event()
+    enter_wrapper = threading.Event()
+    ran = threading.Event()
+    original_wrapper = task_handler._wrap_task
+
+    def delayed_wrapper(*args):
+        wrapper_entered.set()
+        enter_wrapper.wait()
+        original_wrapper(*args)
+
+    with mock.patch.object(task_handler, "_wrap_task", side_effect=delayed_wrapper):
+        status = task_handler.submit_task(ran.set)
+    worker = task_handler._tasks[status.task_id][1]
+    try:
+        assert wrapper_entered.wait(1)
+        task_handler.kill_task(status)
+    finally:
+        enter_wrapper.set()
+        worker.join(timeout=1)
+
+    assert not worker.is_alive()
+    assert not ran.is_set()
+    assert status.done
+    assert status.state == TaskState.KILLED
+    assert isinstance(status.exception(), TaskKilledError)
+
+
+def test_utils_task_handler_does_not_interrupt_status_completion(task_handler):
+    """Cancellation after the callable returns must not strand its status."""
+    completion_started = threading.Event()
+    finish_completion = threading.Event()
+    status = task_handler.submit_task(lambda: None, run=False)
+    original_set_finished = status.set_finished
+
+    def delayed_set_finished():
+        completion_started.set()
+        finish_completion.wait()
+        original_set_finished()
+
+    with mock.patch.object(status, "set_finished", side_effect=delayed_set_finished):
+        task_handler.start_task(status)
+        worker = task_handler._tasks[status.task_id][1]
+        try:
+            assert completion_started.wait(1)
+            task_handler.kill_task(status)
+        finally:
+            finish_completion.set()
+            worker.join(timeout=1)
+
+    assert not worker.is_alive()
+    assert status.done
+    assert status.state == TaskState.COMPLETED
+
+
 ##########################################
 #########  Test PSI cusomt signals  ######
 ##########################################