diff --git a/tests/test_image_monitor_singleton.py b/tests/test_image_monitor_singleton.py index a89321d..0b90716 100644 --- a/tests/test_image_monitor_singleton.py +++ b/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) diff --git a/tests/test_prompt_tracker.py b/tests/test_prompt_tracker.py index f7c0a1e..35a5c9c 100644 --- a/tests/test_prompt_tracker.py +++ b/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() diff --git a/utils/image_monitor.py b/utils/image_monitor.py index 7e6d8c6..d887186 100644 --- a/utils/image_monitor.py +++ b/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.