diff --git a/ophyd_devices/utils/psi_device_base_utils.py b/ophyd_devices/utils/psi_device_base_utils.py index 10fe804..fdf5bb5 100644 --- a/ophyd_devices/utils/psi_device_base_utils.py +++ b/ophyd_devices/utils/psi_device_base_utils.py @@ -778,6 +778,8 @@ class TaskHandler: ) -> TaskStatus: """Submit a task to the task handler. + During shutdown, return a cancelled status and log a warning. + Args: task: The task to run. run: Whether to run the task immediately. @@ -792,23 +794,32 @@ class TaskHandler: daemon=True, ) 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) + if not self._shutting_down: + self._tasks[task_status.task_id] = (task_status, thread) + if run is True: + self.start_task(task_status) + return task_status + + logger.warning(f"Task with ID {task_status.task_id} was ignored during shutdown.") + task_status.state = TaskState.KILLED + task_status.set_exception(TaskKilledError(f"Task {task_status.task_id} was killed.")) return task_status def start_task(self, task_status: TaskStatus) -> None: - """Start a task, + """Start a pending task, or warn and ignore it during shutdown. Args: task_status: The task status object. """ with self._lock: if self._shutting_down: - raise RuntimeError("Cannot start a task during shutdown.") - thread = self._tasks[task_status.task_id][1] + logger.warning(f"Task with ID {task_status.task_id} was ignored during shutdown.") + return + task_info = self._tasks.get(task_status.task_id) + if task_info is None: + logger.warning(f"Task with ID {task_status.task_id} is no longer pending.") + return + thread = task_info[1] if task_status.task_id in self._cancelling_pending: logger.warning(f"Task with ID {task_status.task_id} is being cancelled.") return diff --git a/tests/test_utils.py b/tests/test_utils.py index c401933..c1cf6c9 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -510,26 +510,53 @@ def test_utils_task_handler_shutdown_waits_for_pending_cancellation(task_handler 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.""" +def test_utils_task_handler_shutdown_ignores_callback_submission(task_handler): + """A submission during shutdown returns a cancelled status without running.""" ran = threading.Event() - rejected = threading.Event() + submitted = [] status = task_handler.submit_task(lambda: None, run=False) def submit_from_callback(_): - try: - task_handler.submit_task(ran.set) - except RuntimeError: - rejected.set() + submitted.append(task_handler.submit_task(ran.set)) status.add_callback(submit_from_callback) - task_handler.shutdown() + with mock.patch("ophyd_devices.utils.psi_device_base_utils.logger.warning") as warning: + task_handler.shutdown() - assert rejected.wait(1) + assert len(submitted) == 1 + assert submitted[0].done + assert submitted[0].state == TaskState.KILLED + assert isinstance(submitted[0].exception(), TaskKilledError) + assert any("ignored during shutdown" in str(call) for call in warning.call_args_list) assert not ran.is_set() assert not task_handler._tasks +def test_utils_task_handler_shutdown_ignores_callback_start(task_handler): + """A start request during shutdown leaves the queued task cancelled.""" + ran = threading.Event() + first = task_handler.submit_task(lambda: None, run=False) + second = task_handler.submit_task(ran.set, run=False) + started_from_callback = threading.Event() + + def start_from_callback(_): + task_handler.start_task(second) + started_from_callback.set() + + first.add_callback(start_from_callback) + with mock.patch("ophyd_devices.utils.psi_device_base_utils.logger.warning") as warning: + task_handler.shutdown() + + assert started_from_callback.is_set() + assert first.done and second.done + assert second.state == TaskState.KILLED + assert isinstance(second.exception(), TaskKilledError) + assert any("ignored during shutdown" in str(call) for call in warning.call_args_list) + assert not ran.is_set() + assert not task_handler._tasks + task_handler.start_task(second) + + 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()