Merge pull request #220 from ezmsg-org/fix/219-shutdown-pending

Fix: Multiple shutdown issues and enhance shutdown UX

Author: griffinmilsap

src/ezmsg/core/backend.py

@@ -222,6 +222,15 @@ def graph_address(self) -> AddressType | None:
             return self._graph_context.graph_address
         return self._graph_address
 
+    @property
+    def strict_shutdown(self) -> bool:
+        value = os.environ.get("EZMSG_STRICT_SHUTDOWN", "")
+        return value.lower() in ("1", "true", "yes", "on")
+
+    @strict_shutdown.setter
+    def strict_shutdown(self, value: bool) -> None:
+        os.environ["EZMSG_STRICT_SHUTDOWN"] = "1" if value else "0"
+
     @property
     def graph_server_spawned(self) -> bool:
         return self._graph_server_spawned

src/ezmsg/core/backendprocess.py

@@ -2,15 +2,18 @@
 import concurrent.futures
 import logging
 import inspect
+import os
 import time
 import traceback
 import threading
+import weakref
 
 from abc import abstractmethod
 from collections import defaultdict
 from collections.abc import Callable, Coroutine, Generator, Sequence
 from functools import wraps, partial
-from copy import deepcopy
+from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures.thread import _worker
 from multiprocessing import Process
 from multiprocessing.synchronize import Event as EventType
 from multiprocessing.synchronize import Barrier as BarrierType
@@ -20,7 +23,6 @@
 
 from .stream import Stream, InputStream, OutputStream
 from .unit import Unit, TIMEIT_ATTR, SUBSCRIBES_ATTR
-from .messagechannel import LeakyQueue
 
 from .graphcontext import GraphContext
 from .pubclient import Publisher
@@ -29,6 +31,36 @@
 
 logger = logging.getLogger("ezmsg")
 
+STRICT_SHUTDOWN_ENV = "EZMSG_STRICT_SHUTDOWN"
+
+
+def _strict_shutdown_enabled() -> bool:
+    value = os.environ.get(STRICT_SHUTDOWN_ENV, "")
+    return value.lower() in ("1", "true", "yes", "on")
+
+
+class _DaemonThreadPoolExecutor(ThreadPoolExecutor):
+    def _adjust_thread_count(self) -> None:
+        if self._broken:
+            return
+        num_threads = len(self._threads)
+        if num_threads >= self._max_workers:
+            return
+        thread_name = f"{self._thread_name_prefix or 'ThreadPool'}_{num_threads}"
+        thread = threading.Thread(
+            name=thread_name,
+            target=_worker,
+            args=(
+                weakref.ref(self),
+                self._work_queue,
+                self._initializer,
+                self._initargs,
+            ),
+        )
+        thread.daemon = True
+        thread.start()
+        self._threads.add(thread)
+
 
 class Complete(Exception):
     """
@@ -228,9 +260,14 @@ async def setup_state():
                             ),
                             loop=loop,
                         ).result()
+
+        except asyncio.CancelledError:
+            pass
+
         except Exception:
             self.start_barrier.abort()
-            logger.error(f"{traceback.format_exc()}")
+            # logger.error(f"{traceback.format_exc()}")
+            raise
 
         try:
             logger.debug("Waiting at start barrier!")
@@ -260,9 +297,9 @@ async def coro_wrapper(coro):
                     fn(unit)
                 except NormalTermination:
                     self.term_ev.set()
-                except Exception:
-                    logger.error(f"Exception in Main: {unit.address}")
-                    logger.error(traceback.format_exc())
+                # except Exception:
+                #     logger.error(f"Exception in Main: {unit.address}")
+                #     logger.error(traceback.format_exc())
 
             while True:
                 try:
@@ -352,6 +389,7 @@ def task_wrapper(
         self, unit: Unit, task: Callable
     ) -> Callable[..., Coroutine[Any, Any, None]]:
         task_address = f"{unit.address}:{task.__name__}"
+        strict_shutdown = _strict_shutdown_enabled()
 
         async def publish(stream: Stream, obj: Any) -> None:
             if stream.address in self.pubs:
@@ -400,9 +438,15 @@ async def wrapped_task(msg: Any = None) -> None:
                 self.term_ev.set()
                 raise
 
+            except asyncio.CancelledError:
+                # Normal during shutdown; propagate without logging.
+                raise
+
             except Exception:
                 logger.error(f"Exception in Task: {task_address}")
                 logger.error(traceback.format_exc())
+                if strict_shutdown:
+                    raise
 
         return wrapped_task
 
@@ -491,6 +535,24 @@ def new_threaded_event_loop(
     :rtype: Generator[asyncio.AbstractEventLoop, None, None]
     """
     loop = asyncio.new_event_loop()
+    strict_shutdown = _strict_shutdown_enabled()
+    shutdown_suppress = threading.Event()
+    suppressed_shutdown_errors = {"count": 0}
+    suppressed_lock = threading.Lock()
+    if not strict_shutdown:
+        loop.set_default_executor(
+            _DaemonThreadPoolExecutor(thread_name_prefix="EZMSG")
+        )
+        def _loop_exception_handler(
+            loop_obj: asyncio.AbstractEventLoop, context: dict
+        ) -> None:
+            if shutdown_suppress.is_set():
+                with suppressed_lock:
+                    suppressed_shutdown_errors["count"] += 1
+                return
+            loop_obj.default_exception_handler(context)
+
+        loop.set_exception_handler(_loop_exception_handler)
     thread = threading.Thread(target=run_loop, name="TaskThread", args=(loop,))
     thread.start()
 
@@ -502,6 +564,49 @@ def new_threaded_event_loop(
             logger.debug("Waiting at event...")
             # ev.wait()
         logger.debug("Stopping and closing task thread")
+
+        if not strict_shutdown:
+            shutdown_suppress.set()
+            # Cancel and await remaining tasks before stopping the loop.
+            async def _cancel_remaining() -> int:
+                tasks = [
+                    t
+                    for t in asyncio.all_tasks()
+                    if t is not asyncio.current_task() and not t.done()
+                ]
+                for t in tasks:
+                    t.cancel()
+                if tasks:
+                    await asyncio.wait(tasks)
+                return len(tasks)
+
+            cancelled_count = 0
+            forced_interrupt = False
+            fut = asyncio.run_coroutine_threadsafe(_cancel_remaining(), loop)
+            try:
+                cancelled_count = fut.result()
+            except KeyboardInterrupt:
+                forced_interrupt = True
+                fut.cancel()
+            except Exception:
+                cancelled_count = 0
+
+            suppressed_count = suppressed_shutdown_errors["count"]
+            if cancelled_count or suppressed_count or forced_interrupt:
+                if forced_interrupt and not cancelled_count and not suppressed_count:
+                    logger.warning(
+                        "Shutdown interrupted; tasks may still be running. "
+                        "Re-run with EZMSG_STRICT_SHUTDOWN=1 to debug tasks with poor shutdown behavior."
+                    )
+                else:
+                    logger.warning(
+                        "Shutdown suppressed %d error(s) and cancelled %d task(s). "
+                        "Shutdown was NOT clean; re-run with EZMSG_STRICT_SHUTDOWN=1 "
+                        "to debug tasks with poor shutdown behavior.",
+                        suppressed_count,
+                        cancelled_count,
+                    )
+
         loop.call_soon_threadsafe(loop.stop)
         thread.join()
         loop.close()

src/ezmsg/core/netprotocol.py

@@ -224,7 +224,10 @@ async def read_str(reader: asyncio.StreamReader) -> str:
 
 
 async def close_stream_writer(writer: asyncio.StreamWriter):
-    writer.close()
+    try:
+        writer.close()
+    except RuntimeError:
+        return  # Event loop is closed, transport is already gone
     # ConnectionResetError can be raised on wait_closed.
     # See: https://github.com/python/cpython/issues/83037
     try:

tests/clean_shutdown_examples_runner.py

@@ -0,0 +1,150 @@
+import asyncio
+import os
+import time
+import multiprocessing as mp
+import signal
+import sys
+import threading
+
+import ezmsg.core as ez
+
+
+class CountSettings(ez.Settings):
+    num_msgs: int = 3
+
+
+class CountPublisher(ez.Unit):
+    SETTINGS = CountSettings
+    OUTPUT = ez.OutputStream(int)
+
+    @ez.publisher(OUTPUT)
+    async def publish(self):
+        for i in range(self.SETTINGS.num_msgs):
+            yield self.OUTPUT, i
+            await asyncio.sleep(0.01)
+        raise ez.Complete
+
+
+class CompleteSubscriber(ez.Unit):
+    SETTINGS = CountSettings
+    INPUT = ez.InputStream(int)
+
+    async def initialize(self) -> None:
+        self._count = 0
+
+    @ez.subscriber(INPUT)
+    async def on_message(self, _msg: int) -> None:
+        self._count += 1
+        if self._count >= self.SETTINGS.num_msgs:
+            raise ez.Complete
+
+
+class NormalTerminationSubscriber(ez.Unit):
+    SETTINGS = CountSettings
+    INPUT = ez.InputStream(int)
+
+    async def initialize(self) -> None:
+        self._count = 0
+
+    @ez.subscriber(INPUT)
+    async def on_message(self, _msg: int) -> None:
+        self._count += 1
+        if self._count >= self.SETTINGS.num_msgs:
+            raise ez.NormalTermination
+
+
+class NormalTerminationSubscriberWithThread(ez.Unit):
+    SETTINGS = CountSettings
+    INPUT = ez.InputStream(int)
+
+    async def initialize(self) -> None:
+        self._count = 0
+        self._stop_thread = False
+
+    @ez.thread
+    def background(self) -> None:
+        while not self._stop_thread:
+            time.sleep(0.05)
+
+    @ez.subscriber(INPUT)
+    async def on_message(self, _msg: int) -> None:
+        self._count += 1
+        if self._count >= self.SETTINGS.num_msgs:
+            self._stop_thread = True
+            raise ez.NormalTermination
+
+
+class InfiniteTask(ez.Unit):
+    @ez.task
+    async def run(self) -> None:
+        while True:
+            await asyncio.sleep(0.2)
+
+
+class BaseSystem(ez.Collection):
+    PUB = CountPublisher()
+    SUB = CompleteSubscriber()
+
+    def configure(self) -> None:
+        self.PUB.apply_settings(CountSettings())
+        self.SUB.apply_settings(CountSettings())
+
+    def network(self) -> ez.NetworkDefinition:
+        return ((self.PUB.OUTPUT, self.SUB.INPUT),)
+
+    def process_components(self):
+        return (self.PUB, self.SUB)
+
+
+class CompleteSystem(BaseSystem):
+    SUB = CompleteSubscriber()
+
+
+class NormalTerminationSystem(BaseSystem):
+    SUB = NormalTerminationSubscriber()
+
+
+class NormalTerminationThreadSystem(BaseSystem):
+    SUB = NormalTerminationSubscriberWithThread()
+
+
+class InfiniteSystem(ez.Collection):
+    TASK = InfiniteTask()
+
+    def process_components(self):
+        return (self.TASK,)
+
+
+SYSTEMS = {
+    "complete": CompleteSystem,
+    "normalterm": NormalTerminationSystem,
+    "normalterm_thread": NormalTerminationThreadSystem,
+    "infinite": InfiniteSystem,
+}
+
+
+def main() -> None:
+    if os.environ.get("EZMSG_INBAND_SIGINT"):
+        def _listen() -> None:
+            for line in sys.stdin:
+                if line.strip().upper() == "SIGINT":
+                    signal.raise_signal(signal.SIGINT)
+
+        threading.Thread(target=_listen, daemon=True).start()
+
+    start_method = os.environ.get("EZMSG_MP_START")
+    if start_method:
+        mp.set_start_method(start_method, force=True)
+    case = os.environ.get("EZMSG_SHUTDOWN_EXAMPLE")
+    if case not in SYSTEMS:
+        raise SystemExit(
+            "EZMSG_SHUTDOWN_EXAMPLE must be one of: "
+            + ", ".join(sorted(SYSTEMS))
+        )
+    system = SYSTEMS[case]()
+    print("READY", flush=True)
+    ez.run(SYSTEM=system)
+
+
+if __name__ == "__main__":
+    main()