✨ Feature: 迁移至结构化并发框架 AnyIO (#3053)

nonebot/drivers/none.py

@@ -12,14 +12,18 @@ FrontMatter:
 """
 
 import signal
-import asyncio
-import threading
+from typing import Optional
 from typing_extensions import override
 
+import anyio
+from anyio.abc import TaskGroup
+from exceptiongroup import BaseExceptionGroup, catch
+
 from nonebot.log import logger
 from nonebot.consts import WINDOWS
 from nonebot.config import Env, Config
 from nonebot.drivers import Driver as BaseDriver
+from nonebot.utils import flatten_exception_group
 
 HANDLED_SIGNALS = (
     signal.SIGINT,  # Unix signal 2. Sent by Ctrl+C.
@@ -35,8 +39,8 @@ class Driver(BaseDriver):
     def __init__(self, env: Env, config: Config):
         super().__init__(env, config)
 
-        self.should_exit: asyncio.Event = asyncio.Event()
-        self.force_exit: bool = False
+        self.should_exit: anyio.Event = anyio.Event()
+        self.force_exit: anyio.Event = anyio.Event()
 
     @property
     @override
@@ -54,85 +58,98 @@ class Driver(BaseDriver):
     def run(self, *args, **kwargs):
         """启动 none driver"""
         super().run(*args, **kwargs)
-        loop = asyncio.get_event_loop()
-        loop.run_until_complete(self._serve())
+        anyio.run(self._serve)
 
     async def _serve(self):
-        self._install_signal_handlers()
-        await self._startup()
-        if self.should_exit.is_set():
-            return
-        await self._main_loop()
-        await self._shutdown()
+        async with anyio.create_task_group() as driver_tg:
+            driver_tg.start_soon(self._handle_signals)
+            driver_tg.start_soon(self._listen_force_exit, driver_tg)
+            driver_tg.start_soon(self._handle_lifespan, driver_tg)
 
-    async def _startup(self):
+    async def _handle_signals(self):
         try:
-            await self._lifespan.startup()
-        except Exception as e:
-            logger.opt(colors=True, exception=e).error(
-                "<r><bg #f8bbd0>Application startup failed. "
-                "Exiting.</bg #f8bbd0></r>"
-            )
-            self.should_exit.set()
-            return
-
-        logger.info("Application startup completed.")
-
-    async def _main_loop(self):
-        await self.should_exit.wait()
-
-    async def _shutdown(self):
-        logger.info("Shutting down")
-
-        logger.info("Waiting for application shutdown.")
-
-        try:
-            await self._lifespan.shutdown()
-        except Exception as e:
-            logger.opt(colors=True, exception=e).error(
-                "<r><bg #f8bbd0>Error when running shutdown function. "
-                "Ignored!</bg #f8bbd0></r>"
-            )
-
-        for task in asyncio.all_tasks():
-            if task is not asyncio.current_task() and not task.done():
-                task.cancel()
-        await asyncio.sleep(0.1)
-
-        tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
-        if tasks and not self.force_exit:
-            logger.info("Waiting for tasks to finish. (CTRL+C to force quit)")
-        while tasks and not self.force_exit:
-            await asyncio.sleep(0.1)
-            tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
-
-        for task in tasks:
-            task.cancel()
-
-        await asyncio.gather(*tasks, return_exceptions=True)
-
-        logger.info("Application shutdown complete.")
-        loop = asyncio.get_event_loop()
-        loop.stop()
-
-    def _install_signal_handlers(self) -> None:
-        if threading.current_thread() is not threading.main_thread():
-            # Signals can only be listened to from the main thread.
-            return
-
-        loop = asyncio.get_event_loop()
-
-        try:
-            for sig in HANDLED_SIGNALS:
-                loop.add_signal_handler(sig, self._handle_exit, sig, None)
+            with anyio.open_signal_receiver(*HANDLED_SIGNALS) as signal_receiver:
+                async for sig in signal_receiver:
+                    self.exit(force=self.should_exit.is_set())
         except NotImplementedError:
             # Windows
             for sig in HANDLED_SIGNALS:
-                signal.signal(sig, self._handle_exit)
+                signal.signal(sig, self._handle_legacy_signal)
 
-    def _handle_exit(self, sig, frame):
+    # backport for Windows signal handling
+    def _handle_legacy_signal(self, sig, frame):
         self.exit(force=self.should_exit.is_set())
 
+    async def _handle_lifespan(self, tg: TaskGroup):
+        try:
+            await self._startup()
+
+            if self.should_exit.is_set():
+                return
+
+            await self._listen_exit()
+
+            await self._shutdown()
+        finally:
+            tg.cancel_scope.cancel()
+
+    async def _startup(self):
+        def handle_exception(exc_group: BaseExceptionGroup[Exception]) -> None:
+            self.should_exit.set()
+
+            for exc in flatten_exception_group(exc_group):
+                logger.opt(colors=True, exception=exc).error(
+                    "<r><bg #f8bbd0>Error occurred while running startup hook."
+                    "</bg #f8bbd0></r>"
+                )
+            logger.error(
+                "<r><bg #f8bbd0>Application startup failed. "
+                "Exiting.</bg #f8bbd0></r>"
+            )
+
+        with catch({Exception: handle_exception}):
+            await self._lifespan.startup()
+
+        if not self.should_exit.is_set():
+            logger.info("Application startup completed.")
+
+    async def _listen_exit(self, tg: Optional[TaskGroup] = None):
+        await self.should_exit.wait()
+
+        if tg is not None:
+            tg.cancel_scope.cancel()
+
+    async def _shutdown(self):
+        logger.info("Shutting down")
+        logger.info("Waiting for application shutdown. (CTRL+C to force quit)")
+
+        error_occurred: bool = False
+
+        def handle_exception(exc_group: BaseExceptionGroup[Exception]) -> None:
+            nonlocal error_occurred
+
+            error_occurred = True
+
+            for exc in flatten_exception_group(exc_group):
+                logger.opt(colors=True, exception=exc).error(
+                    "<r><bg #f8bbd0>Error occurred while running shutdown hook."
+                    "</bg #f8bbd0></r>"
+                )
+            logger.error(
+                "<r><bg #f8bbd0>Application shutdown failed. "
+                "Exiting.</bg #f8bbd0></r>"
+            )
+
+        with catch({Exception: handle_exception}):
+            await self._lifespan.shutdown()
+
+        if not error_occurred:
+            logger.info("Application shutdown complete.")
+
+    async def _listen_force_exit(self, tg: TaskGroup):
+        await self.force_exit.wait()
+        tg.cancel_scope.cancel()
+
     def exit(self, force: bool = False):
         """退出 none driver
 
@@ -142,4 +159,4 @@ class Driver(BaseDriver):
         if not self.should_exit.is_set():
             self.should_exit.set()
         if force:
-            self.force_exit = True
+            self.force_exit.set()