fix(threads): a later stop waits again for a worker or observer an earlier stop left running (#302)

Review findings on #372:
- ImageGenerationHandler keeps the workers a stop signalled that had not
  ended; every later shutdown/close joins them again, so a second call no
  longer returns True while the first worker still runs.
- ImageMonitor.stop_monitoring keeps an observer or handler it left
  running and stops and joins it again on the next call.
- A closed handler takes no awaiting history path, and a re-check that
  raced close() is not kept due.
- Tests: second stops of a busy worker and a live observer, the closed
  handler's re-check paths, the handle test asserts the worker ended and
  stops before it asserts, and a no-new-thread check instead of set
  equality.

tests/test_image_monitor_singleton.py

@@ -350,6 +350,17 @@ class TestStartStop(unittest.TestCase):
             self.assertFalse(self.monitor.stop_monitoring(timeout=0.05))
         self.assertIn("watchdog observer still running", "\n".join(logs.output))
         self.assertFalse(self.monitor.is_monitoring)
+        fake.stay_alive = False  # it ends; tearDown's stop waits for it
+
+    def test_a_second_stop_monitoring_waits_for_the_observer(self):
+        fake = FakeObserver(stay_alive=True)
+        self._start_with(fake)
+        with self.assertLogs("prompt_manager.image_monitor", "WARNING"):
+            self.assertFalse(self.monitor.stop_monitoring(timeout=0.05))
+            self.assertFalse(self.monitor.stop_monitoring(timeout=0.05))
+        fake.stay_alive = False
+        self.assertTrue(self.monitor.stop_monitoring())
+        self.assertFalse(self.monitor._unfinished)
 
     def test_stop_monitoring_reports_an_observer_stop_that_raised(self):
         fake = FakeObserver(stop_error=OSError("descriptor gone"))
@@ -357,6 +368,7 @@ class TestStartStop(unittest.TestCase):
         with self.assertLogs("prompt_manager.image_monitor", "WARNING") as logs:
             self.assertFalse(self.monitor.stop_monitoring())
         self.assertIn("Error while stopping image monitoring", "\n".join(logs.output))
+        fake.stop_error = None  # tearDown's stop calls stop() again
 
     def test_stop_monitoring_without_an_observer_returns_true(self):
         self.assertTrue(self.monitor.stop_monitoring())
@@ -495,6 +507,27 @@ class TestHandlerScheduling(unittest.TestCase):
         worker.join(30)
         self.assertFalse(worker.is_alive())
 
+    def test_a_second_stop_waits_for_a_worker_still_finishing(self):
+        entered, release = threading.Event(), threading.Event()
+
+        def busy(path, prompt_snapshot=None):
+            entered.set()
+            release.wait(30)
+
+        self.handler.process_new_image = busy
+        self.handler.schedule(self._image("twice.png"))
+        worker = self.handler._worker
+        self.addCleanup(worker.join, 30)
+        self.addCleanup(release.set)
+        self.assertTrue(entered.wait(30), "the worker never took the image")
+
+        with self.assertLogs("prompt_manager.image_monitor", "WARNING"):
+            self.assertFalse(self.handler.close(timeout=0.05))
+            self.assertFalse(self.handler.close(timeout=0.05))
+        release.set()
+        self.assertTrue(self.handler.close(timeout=30))
+        self.assertFalse(worker.is_alive())
+
     def test_shutdown_on_the_worker_itself_only_signals(self):
         entered, result = threading.Event(), {}
 
@@ -530,10 +563,13 @@ class TestHandlerScheduling(unittest.TestCase):
             self.assertFalse(self.handler.schedule(path))
             # The image in hand can still ask for a history re-check and its
             # run can be sealed afterwards, on the history writer's thread.
+            self.handler._await_history(path)
+            self.assertEqual(self.handler.awaiting_history(), [])
             self._seal(self.handler, "late.png")
+            self.handler._schedule_recheck(path)
         self.assertIsNone(self.handler._worker)
         self.assertEqual(list(self.handler._history_due), [])
-        self.assertEqual(self._live_workers(), before)
+        self.assertEqual(self._live_workers() - before, set())
 
     def test_a_sealed_run_starts_a_worker_on_an_open_handler(self):
         self._image("open.png")
@@ -621,11 +657,14 @@ class TestCloseReleasesHandles(unittest.TestCase):
             Image.new("RGB", (8, 8), "red").save(path)
             try:
                 self.assertTrue(handler.schedule(path))
+                worker = handler._worker
                 self.assertTrue(done.wait(30), "the worker never processed the image")
             finally:
-                self.assertTrue(handler.close())
-                self.assertTrue(tracker.shutdown())
+                closed, shut = handler.close(), tracker.shutdown()
                 db.model.close_all()
+            self.assertTrue(closed)
+            self.assertTrue(shut)
+            self.assertFalse(worker.is_alive())
             assert_closed(self, folder)
 
 

tests/test_prompt_tracker.py

@@ -383,7 +383,7 @@ class TestTrackerShutdown(unittest.TestCase):
 
     def test_shutdown_on_the_cleanup_thread_only_signals(self):
         tracker = _real_tracker()
-        self.addCleanup(tracker.shutdown)
+        self.addCleanup(lambda: self.assertTrue(tracker.shutdown()))
         result = {}
         with patch.object(tracker, "cleanup_thread", threading.current_thread()):
             result["value"] = tracker.shutdown()

utils/image_monitor.py

@@ -150,6 +150,9 @@ class ImageGenerationHandler(FileSystemEventHandler):
         self._sealed_subscribed = False
         # Set by ``close``: no new work, no new worker, ever (#302).
         self._closed = False
+        # Workers a stop signalled that had not ended yet: every later stop
+        # waits for them again, so it never reports True too early.
+        self._finishing: List[threading.Thread] = []
 
     def on_created(self, event):
         """Handle filesystem creation events.
@@ -239,16 +242,23 @@ class ImageGenerationHandler(FileSystemEventHandler):
             worker, self._worker = self._worker, None
             work, self._work = self._work, queue.Queue()
             stop, self._stop = self._stop, threading.Event()
+            if worker is not None:
+                self._finishing.append(worker)
+            finishing = [w for w in self._finishing if w.is_alive()]
+            self._finishing = finishing
         stop.set()
-        if worker is None:
-            return True
-        if worker.is_alive():
+        if worker is not None and worker.is_alive():
             work.put(None)  # sentinel: wake a worker idle on an empty queue
-        if worker is threading.current_thread():
+        if threading.current_thread() in finishing:
             self.logger.debug("shutdown called on the image worker itself")
             return False
-        worker.join(timeout)
-        if worker.is_alive():
+        deadline = time.monotonic() + timeout
+        for thread in finishing:
+            thread.join(max(0.0, deadline - time.monotonic()))
+        with self._pending_lock:
+            self._finishing = [w for w in self._finishing if w.is_alive()]
+            still_running = bool(self._finishing)
+        if still_running:
             self.logger.warning(
                 f"image worker still running after {timeout} s; "
                 "it stops after the image in hand"
@@ -566,7 +576,12 @@ class ImageGenerationHandler(FileSystemEventHandler):
             self._awaiting.popitem(last=False)
 
     def _await_history(self, image_path: str) -> None:
-        """Re-check ``image_path`` once, history only, when its run is sealed."""
+        """Re-check ``image_path`` once, history only, when its run is sealed.
+
+        A closed handler re-checks nothing.
+        """
+        if self._closed:
+            return
         if not self._sealed_subscribed:
             subscribe_outputs_sealed(self._on_outputs_sealed)
             self._sealed_subscribed = True
@@ -624,6 +639,10 @@ class ImageGenerationHandler(FileSystemEventHandler):
         if self.schedule(image_path, HISTORY_ONLY):
             return
         with self._awaiting_lock:
+            # ``close`` sets the flag before ``shutdown`` clears this set, so
+            # a close that raced the refusal above is seen here.
+            if self._closed:
+                return
             self._history_due.pop(image_path, None)
             self._history_due[image_path] = None
             while len(self._history_due) > MAX_AWAITING_HISTORY:
@@ -922,6 +941,9 @@ class ImageMonitor:
         self.observer = None
         self.handler = None
         self.monitored_directories = []
+        # (observer, handler) pairs a stop left running: the next stop waits
+        # for them again instead of reporting True.
+        self._unfinished: List[Tuple[Any, Any]] = []
         self.logger = get_logger("prompt_manager.image_monitor")
         self.review_recorder: Optional[Callable[[str], str]] = None
         self.thumbnail_queue: Optional[Callable[[str], str]] = None
@@ -1055,41 +1077,63 @@ class ImageMonitor:
         Closes the handler first (its worker is joined and its dispatch
         returns at once), then stops the watchdog observer and joins it.
         watchdog's ``observer.stop()`` joins its emitters without a limit;
-        it counts against the deadline but is not bounded here. Safe to call
-        when monitoring never started and more than once.
+        it counts against the deadline but is not bounded here. A part an
+        earlier call left running is waited for again. Safe to call when
+        monitoring never started and more than once.
 
         Returns:
-            True when the worker and the observer have both ended.
+            True when every worker and observer it stopped has ended.
         """
         deadline = time.monotonic() + timeout
         observer, self.observer = self.observer, None
         handler, self.handler = self.handler, None
         self.monitored_directories = []
-        if observer is None:
+        pairs, self._unfinished = self._unfinished, []
+        if observer is not None:
+            pairs.append((observer, handler))
+        if not pairs:
             return True
 
         def remaining() -> float:
             return max(0.0, deadline - time.monotonic())
 
-        handler_ok = True
-        if handler is not None:
-            try:
-                handler_ok = handler.close(remaining())
-            except Exception as e:
-                handler_ok = False
-                self.logger.warning(f"Error while stopping the image worker: {e}")
-        stopped = True
+        for pair_observer, pair_handler in pairs:
+            handler_ok = self._close_handler(pair_handler, remaining())
+            observer_ok = self._stop_observer(pair_observer, remaining())
+            if not observer_ok:
+                self.logger.warning(
+                    f"watchdog observer still running after {timeout} s"
+                )
+            if not (handler_ok and observer_ok):
+                self._unfinished.append(
+                    (
+                        None if observer_ok else pair_observer,
+                        None if handler_ok else pair_handler,
+                    )
+                )
+        self.logger.debug("Image monitoring stopped")
+        return not self._unfinished
+
+    def _close_handler(self, handler, timeout: float) -> bool:
+        if handler is None:
+            return True
+        try:
+            return handler.close(timeout)
+        except Exception as e:
+            self.logger.warning(f"Error while stopping the image worker: {e}")
+            return False
+
+    def _stop_observer(self, observer, timeout: float) -> bool:
+        """``observer.stop()`` (safe to repeat), then a bounded join."""
+        if observer is None:
+            return True
         try:
             observer.stop()
-            observer.join(remaining())
+            observer.join(timeout)
         except Exception as e:
-            stopped = False
             self.logger.warning(f"Error while stopping image monitoring: {e}")
-        if observer.is_alive():
-            stopped = False
-            self.logger.warning(f"watchdog observer still running after {timeout} s")
-        self.logger.debug("Image monitoring stopped")
-        return handler_ok and stopped
+            return False
+        return not observer.is_alive()
 
     def detect_comfyui_output_dirs(self) -> list:
         """Auto-detect ComfyUI output directories.