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: