fix(task-handler): ignore new work during shutdown

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

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