fix: proper return codes

Author: griffinmilsap

src/ezmsg/core/backend.py

@@ -4,6 +4,7 @@
 import enum
 import logging
 import os
+import signal
 from threading import BrokenBarrierError
 from multiprocessing import Event, Barrier
 from multiprocessing.synchronize import Event as EventType
@@ -23,6 +24,7 @@
 from .backendprocess import (
     BackendProcess,
     DefaultBackendProcess,
+    ShutdownSummary,
     new_threaded_event_loop,
 )
 
@@ -170,6 +172,7 @@ class GraphRunner:
     _graph_context: GraphContext | None
     _loop: asyncio.AbstractEventLoop | None
     _loop_cm: object | None
+    _loop_shutdown_summary: ShutdownSummary | None
     _main_process: BackendProcess | None
     _spawned_processes: list[BackendProcess]
     _start_participant: bool
@@ -208,6 +211,7 @@ def __init__(
         self._graph_context = None
         self._loop = None
         self._loop_cm = None
+        self._loop_shutdown_summary = None
         self._main_process = None
         self._spawned_processes = []
         self._start_participant = False
@@ -320,7 +324,10 @@ def _initialize(self, force_single_process: bool, wait_for_ready: bool) -> bool:
         if self._execution_context is None:
             return False
 
-        self._loop_cm = new_threaded_event_loop()
+        self._loop_shutdown_summary = ShutdownSummary()
+        self._loop_cm = new_threaded_event_loop(
+            shutdown_summary=self._loop_shutdown_summary
+        )
         self._loop = self._loop_cm.__enter__()
 
         try:
@@ -389,12 +396,15 @@ def _run_main_process(self) -> None:
         self._main_process = self._execution_context.processes[0]
         self._start_processes(self._execution_context.processes[1:])
 
+        interrupts = 0
+        forced_sigint = False
         try:
             self._main_process.process(self._loop)
             self._join_spawned_processes()
             logger.info("All processes exited normally")
 
         except KeyboardInterrupt:
+            interrupts += 1
             logger.info(
                 "Attempting graceful shutdown, interrupt again to force quit..."
             )
@@ -404,17 +414,71 @@ def _run_main_process(self) -> None:
                 self._join_spawned_processes()
 
             except KeyboardInterrupt:
+                interrupts += 1
+                forced_sigint = True
                 logger.warning("Interrupt intercepted, force quitting")
                 self._execution_context.start_barrier.abort()
                 self._execution_context.stop_barrier.abort()
                 for proc in self._spawned_processes:
                     proc.terminate()
 
         finally:
-            self._join_spawned_processes()
-            self._cleanup()
+            while True:
+                try:
+                    self._join_spawned_processes()
+                    self._cleanup()
+                    break
+                except KeyboardInterrupt:
+                    interrupts += 1
+                    if interrupts >= 2:
+                        forced_sigint = True
+                        logger.warning("Interrupt intercepted, force quitting")
+                        if self._execution_context is not None:
+                            self._execution_context.start_barrier.abort()
+                            self._execution_context.stop_barrier.abort()
+                        for proc in self._spawned_processes:
+                            proc.terminate()
+                        self._cleanup()
+                        break
+                    logger.info(
+                        "Interrupt received during cleanup; attempting graceful shutdown..."
+                    )
+                    if self._execution_context is not None:
+                        self._execution_context.term_ev.set()
             self._started = False
             self._stopped = True
+            if interrupts and not forced_sigint and self._shutdown_was_unclean():
+                forced_sigint = True
+            if forced_sigint:
+                self._exit_with_sigint()
+
+    def _shutdown_was_unclean(self) -> bool:
+        main_shutdown_errors = bool(
+            self._main_process is not None
+            and getattr(self._main_process, "_shutdown_errors", False)
+        )
+        summary = self._loop_shutdown_summary
+        loop_unclean = bool(summary is not None and summary.unclean)
+        return main_shutdown_errors or loop_unclean
+
+    def _exit_with_sigint(self) -> None:
+        prev_handler = None
+        try:
+            prev_handler = signal.getsignal(signal.SIGINT)
+            signal.signal(signal.SIGINT, signal.SIG_DFL)
+            signal.raise_signal(signal.SIGINT)
+        except Exception:
+            code = 0xC000013A if os.name == "nt" else 130
+            raise SystemExit(code)
+        finally:
+            if prev_handler is not None:
+                try:
+                    signal.signal(signal.SIGINT, prev_handler)
+                except Exception:
+                    pass
+
+        code = 0xC000013A if os.name == "nt" else 130
+        raise SystemExit(code)
 
     def _cleanup(self) -> None:
         if self._cleanup_done:

src/ezmsg/core/backendprocess.py

@@ -9,6 +9,7 @@
 import weakref
 
 from abc import abstractmethod
+from dataclasses import dataclass
 from collections import defaultdict
 from collections.abc import Callable, Coroutine, Generator, Sequence
 from functools import wraps, partial
@@ -39,7 +40,28 @@ def _strict_shutdown_enabled() -> bool:
     return value.lower() in ("1", "true", "yes", "on")
 
 
+@dataclass
+class ShutdownSummary:
+    cancelled_tasks: int = 0
+    executor_active: int = 0
+    suppressed_errors: int = 0
+    forced_interrupt: bool = False
+
+    @property
+    def unclean(self) -> bool:
+        return bool(
+            self.executor_active
+            or self.suppressed_errors
+            or self.forced_interrupt
+        )
+
+
 class _DaemonThreadPoolExecutor(ThreadPoolExecutor):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self._active_count = 0
+        self._active_lock = threading.Lock()
+
     def _adjust_thread_count(self) -> None:
         if self._broken:
             return
@@ -61,6 +83,22 @@ def _adjust_thread_count(self) -> None:
         thread.start()
         self._threads.add(thread)
 
+    def submit(self, fn, /, *args, **kwargs):
+        fut = super().submit(fn, *args, **kwargs)
+        with self._active_lock:
+            self._active_count += 1
+
+        def _decrement(_):
+            with self._active_lock:
+                self._active_count -= 1
+
+        fut.add_done_callback(_decrement)
+        return fut
+
+    def active_count(self) -> int:
+        with self._active_lock:
+            return self._active_count
+
 
 class Complete(Exception):
     """
@@ -178,11 +216,13 @@ class DefaultBackendProcess(BackendProcess):
     """
 
     pubs: dict[str, Publisher]
+    _shutdown_errors: bool
 
     def process(self, loop: asyncio.AbstractEventLoop) -> None:
         main_func = None
         context = GraphContext(self.graph_address)
         coro_callables: dict[str, Callable[[], Coroutine[Any, Any, None]]] = dict()
+        self._shutdown_errors = False
 
         try:
             self.pubs = dict()
@@ -445,6 +485,8 @@ async def wrapped_task(msg: Any = None) -> None:
             except Exception:
                 logger.error(f"Exception in Task: {task_address}")
                 logger.error(traceback.format_exc())
+                if self.term_ev.is_set():
+                    self._shutdown_errors = True
                 if strict_shutdown:
                     raise
 
@@ -522,6 +564,7 @@ def run_loop(loop: asyncio.AbstractEventLoop):
 @contextmanager
 def new_threaded_event_loop(
     ev: threading.Event | None = None,
+    shutdown_summary: ShutdownSummary | None = None,
 ) -> Generator[asyncio.AbstractEventLoop, None, None]:
     """
     Create a new asyncio event loop running in a separate thread.
@@ -531,6 +574,8 @@ def new_threaded_event_loop(
 
     :param ev: Optional event to signal when the loop is ready.
     :type ev: threading.Event | None
+    :param shutdown_summary: Optional shutdown summary object to populate.
+    :type shutdown_summary: ShutdownSummary | None
     :return: Context manager yielding the event loop.
     :rtype: Generator[asyncio.AbstractEventLoop, None, None]
     """
@@ -539,10 +584,10 @@ def new_threaded_event_loop(
     shutdown_suppress = threading.Event()
     suppressed_shutdown_errors = {"count": 0}
     suppressed_lock = threading.Lock()
+    executor = None
     if not strict_shutdown:
-        loop.set_default_executor(
-            _DaemonThreadPoolExecutor(thread_name_prefix="EZMSG")
-        )
+        executor = _DaemonThreadPoolExecutor(thread_name_prefix="EZMSG")
+        loop.set_default_executor(executor)
         def _loop_exception_handler(
             loop_obj: asyncio.AbstractEventLoop, context: dict
         ) -> None:
@@ -607,6 +652,14 @@ async def _cancel_remaining() -> int:
                         cancelled_count,
                     )
 
+            if shutdown_summary is not None:
+                shutdown_summary.cancelled_tasks = cancelled_count
+                shutdown_summary.executor_active = (
+                    executor.active_count() if executor is not None else 0
+                )
+                shutdown_summary.suppressed_errors = suppressed_count
+                shutdown_summary.forced_interrupt = forced_interrupt
+
         loop.call_soon_threadsafe(loop.stop)
         thread.join()
         loop.close()

tests/clean_shutdown_examples_runner.py

@@ -142,8 +142,29 @@ def _listen() -> None:
             + ", ".join(sorted(SYSTEMS))
         )
     system = SYSTEMS[case]()
-    print("READY", flush=True)
-    ez.run(SYSTEM=system)
+    runner = ez.GraphRunner(SYSTEM=system)
+    ready_emitted = threading.Event()
+    done = threading.Event()
+
+    def _emit_ready() -> None:
+        if not ready_emitted.is_set():
+            print("READY", flush=True)
+            ready_emitted.set()
+
+    def _watch_ready() -> None:
+        while not done.is_set():
+            if runner.running:
+                _emit_ready()
+                return
+            time.sleep(0.01)
+        _emit_ready()
+
+    threading.Thread(target=_watch_ready, daemon=True).start()
+    try:
+        runner.run_blocking()
+    finally:
+        done.set()
+        _emit_ready()
 
 
 if __name__ == "__main__":

tests/shutdown_runner.py

@@ -2,6 +2,7 @@
 import os
 import socket
 import threading
+import time
 import signal
 import sys
 
@@ -73,8 +74,28 @@ def _listen() -> None:
             "EZMSG_SHUTDOWN_TEST must be one of: " + ", ".join(sorted(UNITS))
         )
     runner = ez.GraphRunner(SYSTEM=UNITS[target]())
-    print("READY", flush=True)
-    runner.run_blocking()
+    ready_emitted = threading.Event()
+    done = threading.Event()
+
+    def _emit_ready() -> None:
+        if not ready_emitted.is_set():
+            print("READY", flush=True)
+            ready_emitted.set()
+
+    def _watch_ready() -> None:
+        while not done.is_set():
+            if runner.running:
+                _emit_ready()
+                return
+            time.sleep(0.01)
+        _emit_ready()
+
+    threading.Thread(target=_watch_ready, daemon=True).start()
+    try:
+        runner.run_blocking()
+    finally:
+        done.set()
+        _emit_ready()
 
 
 if __name__ == "__main__":