From 9b07d41f86e75a96175ee268d43fddd7f55d6555 Mon Sep 17 00:00:00 2001
From: snowy <snowykami@outlook.com>
Date: Thu, 15 Aug 2024 17:32:02 +0800
Subject: [PATCH] =?UTF-8?q?:zap:=20=E5=9C=A8=E7=BB=93=E6=9D=9F=E8=BF=9B?=
 =?UTF-8?q?=E7=A8=8B=E6=97=B6=E6=97=A0=E6=B3=95=E6=9D=80=E6=AD=BB=E8=BF=9B?=
 =?UTF-8?q?=E7=A8=8B=E7=9A=84=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 liteyuki/bot/__init__.py                      | 17 +++++-
 liteyuki/core/manager.py                      | 57 ++++++++++---------
 .../nonebot_launcher/__init__.py              |  4 +-
 3 files changed, 50 insertions(+), 28 deletions(-)

diff --git a/liteyuki/bot/__init__.py b/liteyuki/bot/__init__.py
index 35904a7b..c0684f0e 100644
--- a/liteyuki/bot/__init__.py
+++ b/liteyuki/bot/__init__.py
@@ -38,7 +38,6 @@ class LiteyukiBot:
 
         self.loop = asyncio.new_event_loop()
         asyncio.set_event_loop(self.loop)
-        self.loop_thread = threading.Thread(target=self.loop.run_forever, daemon=True)
         self.stop_event = threading.Event()
         self.call_restart_count = 0
 
@@ -51,6 +50,22 @@ class LiteyukiBot:
         self.lifespan.before_start() # 启动前钩子
         self.process_manager.start_all()
         self.lifespan.after_start() # 启动后钩子
+        self.keep_alive()
+
+    def keep_alive(self):
+        """
+        保持轻雪运行
+        Returns:
+
+        """
+        try:
+            while not self.stop_event.is_set():
+                time.sleep(1)
+        except KeyboardInterrupt:
+            logger.info("Liteyuki is stopping...")
+            self.stop()
+        finally:
+            self.lifespan.after_shutdown()
 
     def restart(self, delay: int = 0):
         """
diff --git a/liteyuki/core/manager.py b/liteyuki/core/manager.py
index b2df2958..4e3ffef3 100644
--- a/liteyuki/core/manager.py
+++ b/liteyuki/core/manager.py
@@ -8,14 +8,18 @@ Copyright (C) 2020-2024 LiteyukiStudio. All Rights Reserved
 @File    : manager.py
 @Software: PyCharm
 """
-import asyncio
+
+import atexit
 import threading
+import signal
 from multiprocessing import Process
-from typing import TYPE_CHECKING
+from typing import Any, Callable, Optional, Protocol, TYPE_CHECKING, TypeAlias
 
 from liteyuki.comm import Channel, get_channel, set_channels
 from liteyuki.log import logger
 
+TARGET_FUNC: TypeAlias = Callable[[Channel, Channel, ...], Any]
+
 if TYPE_CHECKING:
     from liteyuki.bot import LiteyukiBot
 
@@ -36,6 +40,10 @@ class ProcessManager:
         self.targets: dict[str, tuple[callable, tuple, dict]] = {}
         self.processes: dict[str, Process] = {}
 
+        atexit.register(self.terminate_all)
+        signal.signal(signal.SIGINT, self._handle_exit)
+        signal.signal(signal.SIGTERM, self._handle_exit)
+
     def start(self, name: str):
         """
         开启后自动监控进程，并添加到进程字典中
@@ -61,40 +69,37 @@ class ProcessManager:
 
         def _start_monitor():
             while True:
-                try:
-                    data = chan_active.receive()
-                    if data == 0:
-                        # 停止
-                        logger.info(f"Stopping process {name}")
-                        self.bot.lifespan.before_process_shutdown()
-                        self.terminate(name)
-                        break
-                    elif data == 1:
-                        # 重启
-                        logger.info(f"Restarting process {name}")
-                        self.bot.lifespan.before_process_shutdown()
-                        self.bot.lifespan.before_process_restart()
-                        self.terminate(name)
-                        _start_process()
-                        continue
-                    else:
-                        logger.warning("Unknown data received, ignored.")
-                except KeyboardInterrupt:
+                data = chan_active.receive()
+                if data == 0:
+                    # 停止
                     logger.info(f"Stopping process {name}")
                     self.bot.lifespan.before_process_shutdown()
-                    self.terminate_all()
+                    self.terminate(name)
                     break
-
-        threading.Thread(target=_start_monitor).start()
+                elif data == 1:
+                    # 重启
+                    logger.info(f"Restarting process {name}")
+                    self.bot.lifespan.before_process_shutdown()
+                    self.bot.lifespan.before_process_restart()
+                    self.terminate(name)
+                    _start_process()
+                    continue
+                else:
+                    logger.warning("Unknown data received, ignored.")
 
     def start_all(self):
         """
         启动所有进程
         """
         for name in self.targets:
-            self.start(name)
+            threading.Thread(target=self.start, args=(name,), daemon=True).start()
 
-    def add_target(self, name: str, target, args: tuple = (), kwargs=None):
+    def _handle_exit(self, signum, frame):
+        logger.info("Received signal, stopping all processes.")
+        self.terminate_all()
+        exit(0)
+
+    def add_target(self, name: str, target: TARGET_FUNC, args: tuple = (), kwargs=None):
         """
         添加进程
         Args:
diff --git a/src/liteyuki_plugins/nonebot_launcher/__init__.py b/src/liteyuki_plugins/nonebot_launcher/__init__.py
index 2be0a808..62e146a0 100644
--- a/src/liteyuki_plugins/nonebot_launcher/__init__.py
+++ b/src/liteyuki_plugins/nonebot_launcher/__init__.py
@@ -8,6 +8,8 @@ Copyright (C) 2020-2024 LiteyukiStudio. All Rights Reserved
 @File    : __init__.py.py
 @Software: PyCharm
 """
+import asyncio
+import time
 
 import nonebot
 
@@ -21,7 +23,7 @@ __plugin_meta__ = PluginMetadata(
 )
 
 
-def nb_run(chan_active: "Channel", chan_passive: "Channel", **kwargs):
+def nb_run(chan_active: "Channel", chan_passive: "Channel", *args, **kwargs):
     """
     初始化NoneBot并运行在子进程
     Args: