Task queue: bound watchdog reset path

The TaskQueueReader watchdog is a recovery path, not normal shutdown, but it used the same unbounded stop and channel close calls that caused shutdown hangs. A blocked reset can wedge the scheduler thread and prevent further recovery attempts.

Run the reset in a daemon helper with a bounded timeout, log and return if it cannot complete, and add a regression test for a blocking channel.close().

Author: xsedla1o

dp3/common/control.py

@@ -43,10 +43,7 @@ def as_message(self) -> str:
 class Control:
     """Class enabling remote control of the platform's internal events."""
 
-    def __init__(
-        self,
-        platform_config: PlatformConfig,
-    ) -> None:
+    def __init__(self, platform_config: PlatformConfig) -> None:
         self.log = logging.getLogger("Control")
         self.action_handlers: dict[ControlAction, Callable] = {}
 

dp3/task_processing/task_distributor.py

@@ -167,9 +167,6 @@ def _force_shutdown(self, alive_workers: int, reader_stopped: bool) -> None:
         os._exit(1)  # nuke entire process
 
     def _stop_task_queue_reader(self, timeout: float) -> bool:
-        if not self._task_queue_reader.running:
-            return True
-
         reader_stopped = False
 
         # TaskQueueReader.stop() is expected to honor its timeout, but run it in

dp3/task_processing/task_queue.py

@@ -118,17 +118,18 @@ def __init__(self, rabbit_config: dict | None = None) -> None:
     def __del__(self):
         self.disconnect()
 
-    def connect(self) -> None:
+    def connect(self, retry_while: Callable[[], bool] | None = None) -> bool:
         """Create a connection (or reconnect after error).
 
-        If connection can't be established, try it again indefinitely.
+        If connection can't be established, try it again indefinitely unless
+        ``retry_while`` is provided and returns false.
         """
         if self.connection:
             self.connection.close()
         self._connection_id += 1
 
         attempts = 0
-        while True:
+        while retry_while is None or retry_while():
             attempts += 1
             try:
                 self.connection = amqpstorm.Connection(**self.conn_params)
@@ -144,15 +145,20 @@ def connect(self) -> None:
                 channel.confirm_deliveries()
                 channel.basic.qos(PREFETCH_COUNT)
                 self.channel = channel
-                break
+                return True
             except amqpstorm.AMQPError as e:
                 sleep_time = RECONNECT_DELAYS[min(attempts, len(RECONNECT_DELAYS)) - 1]
                 self.log.error(
                     f"RabbitMQ connection error (will try to reconnect in {sleep_time}s): {e}"
                 )
-                time.sleep(sleep_time)
+                sleep_deadline = time.monotonic() + sleep_time
+                while time.monotonic() < sleep_deadline:
+                    if retry_while is not None and not retry_while():
+                        return False
+                    time.sleep(min(0.1, sleep_deadline - time.monotonic()))
             except KeyboardInterrupt:
-                break
+                return False
+        return False
 
     def disconnect(self) -> None:
         if self.connection:
@@ -407,6 +413,8 @@ def __init__(
 
         self._consuming_thread: threading.Thread | None = None
         self._processing_thread: threading.Thread | None = None
+        self._watchdog_recovery_thread: threading.Thread | None = None
+        self._stopping = False
 
         # Receive messages into 2 temporary queues
         # (max length should be equal to prefetch_count set in RabbitMQReader)
@@ -427,8 +435,9 @@ def start(self) -> None:
         if self.running:
             raise RuntimeError("Already running")
 
-        if not self.connection:
-            self.connect()
+        if not self.connection and not self.connect(retry_while=lambda: not self._stopping):
+            return
+        self._stopping = False
 
         self.log.info("Starting TaskQueueReader")
 
@@ -454,14 +463,19 @@ def stop(self, timeout: float | None = None) -> bool:
         Returns:
             Whether all internal reader threads stopped.
         """
+        self._stopping = True
+        deadline_ts = None if timeout is None else time.monotonic() + timeout
+        recovery_stopped = self._stop_watchdog_recovery_thread(_remaining_time(deadline_ts))
         if not self.running:
-            raise RuntimeError("Not running")
+            if recovery_stopped:
+                return True
+            self.log.error("TaskQueueReader watchdog recovery did not stop before timeout")
+            return False
 
         self.running = False
-        deadline_ts = None if timeout is None else time.monotonic() + timeout
         consuming_stopped = self._stop_consuming_thread(_remaining_time(deadline_ts))
         processing_stopped = self._stop_processing_thread(_remaining_time(deadline_ts))
-        stopped = consuming_stopped and processing_stopped
+        stopped = consuming_stopped and processing_stopped and recovery_stopped
         if stopped:
             self.log.info("TaskQueueReader stopped")
         else:
@@ -473,7 +487,7 @@ def reconnect(self) -> None:
         self.cache.clear()
         self.cache_pri.clear()
 
-        self.connect()
+        self.connect(retry_while=lambda: self.running and not self._stopping)
 
     def check(self) -> bool:
         """
@@ -604,32 +618,93 @@ def _msg_processing_thread_func(self):
                 self.log.exception("Error in user callback function. %s: %s", type(e), str(e))
                 self.log.error("Original message: %s", body)
 
-    def watchdog(self):
+    def watchdog(self) -> bool:
         """
-        Check whether both threads are running and perform a reset if not.
+        Check whether both threads are running and start recovery if not.
 
-        Register to be called periodically by scheduler.
+        Register to be called periodically by scheduler. RabbitMQ recovery remains
+        resilient and may retry indefinitely, but it runs in a daemon thread so it
+        cannot block scheduler execution or process shutdown.
         """
         proc = self._processing_thread is not None and self._processing_thread.is_alive()
         cons = self._consuming_thread is not None and self._consuming_thread.is_alive()
 
-        if not proc or not cons:
-            self.log.error(
-                "Dead threads detected, processing=%s, consuming=%s, restarting TaskQueueReader.",
-                "alive" if proc else "dead",
-                "alive" if cons else "dead",
-            )
-            self._stop_consuming_thread()
-            self._stop_processing_thread()
+        if proc and cons:
+            return True
+
+        self.log.error(
+            "Dead threads detected, processing=%s, consuming=%s, restarting TaskQueueReader.",
+            "alive" if proc else "dead",
+            "alive" if cons else "dead",
+        )
+        if self._watchdog_recovery_thread and self._watchdog_recovery_thread.is_alive():
+            self.log.warning("TaskQueueReader watchdog recovery is already running")
+            return False
 
-            if self.channel is not None:
-                self.channel.close()
-            self.channel = None
-            self.cache.clear()
-            self.cache_pri.clear()
+        self._watchdog_recovery_thread = threading.Thread(
+            target=self._recover_after_dead_threads,
+            name=f"TaskQueueReaderRecovery-{self.worker_index}",
+            daemon=True,
+        )
+        self._watchdog_recovery_thread.start()
+        return False
+
+    def _recover_after_dead_threads(self) -> None:
+        consuming_thread = self._consuming_thread
+        processing_thread = self._processing_thread
+        channel = self.channel
+        cache = self.cache
+        cache_pri = self.cache_pri
+        self.running = False
 
-            self.connect()
+        def recovery_can_continue() -> bool:
+            return not self._stopping
+
+        try:
+            if consuming_thread:
+                if consuming_thread.is_alive():
+                    with contextlib.suppress(amqpstorm.AMQPError):
+                        if channel is not None:
+                            channel.stop_consuming()
+                consuming_thread.join()
+
+            if processing_thread:
+                if processing_thread.is_alive():
+                    self.cache_full.set()  # break potential wait() for data
+                processing_thread.join()
+
+            if channel is not None:
+                channel.close()
+
+            if not recovery_can_continue():
+                return
+            if self._consuming_thread is consuming_thread:
+                self._consuming_thread = None
+            if self._processing_thread is processing_thread:
+                self._processing_thread = None
+            if self.channel is channel:
+                self.channel = None
+            if self.cache is cache:
+                self.cache.clear()
+            if self.cache_pri is cache_pri:
+                self.cache_pri.clear()
+
+            if not self.connect(retry_while=recovery_can_continue):
+                return
+            if not recovery_can_continue():
+                self.disconnect()
+                return
             self.start()
+        except Exception:
+            self.log.exception("Error while resetting TaskQueueReader")
+
+    def _stop_watchdog_recovery_thread(self, timeout: float | None = None) -> bool:
+        if self._watchdog_recovery_thread:
+            self._watchdog_recovery_thread.join(timeout=timeout)
+            if self._watchdog_recovery_thread.is_alive():
+                return False
+        self._watchdog_recovery_thread = None
+        return True
 
     def _stop_consuming_thread(self, timeout: float | None = None) -> bool:
         if self._consuming_thread:

dp3/worker.py

@@ -216,6 +216,11 @@ def main(app_name: str, config_dir: str, process_index: int, verbose: bool) -> i
         # Create instances of core components
         log.info(f"***** {app_name} worker {process_index} of {num_processes} start *****")
 
+        # Lock used to control when the program stops.
+        daemon_stop_lock = threading.Lock()
+        daemon_stop_lock.acquire()
+        clean_stop_requested = threading.Event()
+
         # EventCountLogger
         ecl = EventCountLogger(
             platform_config.config.get("event_logging.groups"),
@@ -253,11 +258,6 @@ def main(app_name: str, config_dir: str, process_index: int, verbose: bool) -> i
         Telemetry(db, platform_config, registrar)
         GarbageCollector(db, platform_config, registrar)
 
-        # Lock used to control when the program stops.
-        daemon_stop_lock = threading.Lock()
-        daemon_stop_lock.acquire()
-        clean_stop_requested = threading.Event()
-
         # Signal handler releasing the lock on SIGINT or SIGTERM.
         def sigint_handler(signum, frame):
             log.debug(