支持通道通过通道在进程中传递

This commit is contained in:
远野千束 2024-08-17 23:46:43 +08:00
parent aa9cae7008
commit 8d78e643e0
4 changed files with 81 additions and 34 deletions

View File

@ -9,7 +9,8 @@ import time
from typing import Any, Optional from typing import Any, Optional
from liteyuki.bot.lifespan import (LIFESPAN_FUNC, Lifespan) from liteyuki.bot.lifespan import (LIFESPAN_FUNC, Lifespan)
from liteyuki.comm import get_channel from liteyuki.comm.channel import get_channel
from liteyuki.comm.storage import shared_memory
from liteyuki.utils import IS_MAIN_PROCESS from liteyuki.utils import IS_MAIN_PROCESS
from liteyuki.core.manager import ProcessManager from liteyuki.core.manager import ProcessManager
from liteyuki.log import init_log, logger from liteyuki.log import init_log, logger
@ -24,28 +25,40 @@ __all__ = [
class LiteyukiBot: class LiteyukiBot:
def __init__(self, *args, **kwargs): def __init__(self, *args, **kwargs) -> None:
"""
初始化轻雪实例
Args:
*args:
**kwargs: 配置
"""
"""常规操作"""
print_logo() print_logo()
global _BOT_INSTANCE global _BOT_INSTANCE
_BOT_INSTANCE = self # 引用 _BOT_INSTANCE = self # 引用
"""配置"""
self.config: dict[str, Any] = kwargs self.config: dict[str, Any] = kwargs
"""初始化"""
self.init(**self.config) # 初始化 self.init(**self.config) # 初始化
logger.info("Liteyuki is initializing...") logger.info("Liteyuki is initializing...")
"""生命周期管理"""
self.lifespan = Lifespan() self.lifespan = Lifespan()
self.process_manager: ProcessManager = ProcessManager(lifespan=self.lifespan)
self.process_manager: ProcessManager = ProcessManager(bot=self) """事件循环"""
self.loop = asyncio.new_event_loop() self.loop = asyncio.new_event_loop()
asyncio.set_event_loop(self.loop) asyncio.set_event_loop(self.loop)
self.stop_event = threading.Event() self.stop_event = threading.Event()
self.call_restart_count = 0 self.call_restart_count = 0
"""插件加载"""
load_plugins("liteyuki/plugins") # 加载轻雪插件 load_plugins("liteyuki/plugins") # 加载轻雪插件
"""信号处理"""
signal.signal(signal.SIGINT, self._handle_exit) signal.signal(signal.SIGINT, self._handle_exit)
signal.signal(signal.SIGTERM, self._handle_exit) signal.signal(signal.SIGTERM, self._handle_exit)
atexit.register(self.process_manager.terminate_all) # 注册退出时的函数 atexit.register(self.process_manager.terminate_all) # 注册退出时的函数
@ -229,15 +242,17 @@ class LiteyukiBot:
return self.lifespan.on_after_nonebot_init(func) return self.lifespan.on_after_nonebot_init(func)
_BOT_INSTANCE: Optional[LiteyukiBot] = None _BOT_INSTANCE: LiteyukiBot
def get_bot() -> Optional[LiteyukiBot]: def get_bot() -> LiteyukiBot:
""" """
获取轻雪实例 获取轻雪实例
Returns: Returns:
LiteyukiBot: 当前的轻雪实例 LiteyukiBot: 当前的轻雪实例
""" """
if IS_MAIN_PROCESS: if IS_MAIN_PROCESS:
if _BOT_INSTANCE is None: if _BOT_INSTANCE is None:
raise RuntimeError("Liteyuki instance not initialized.") raise RuntimeError("Liteyuki instance not initialized.")
@ -246,7 +261,7 @@ def get_bot() -> Optional[LiteyukiBot]:
raise RuntimeError("Can't get bot instance in sub process.") raise RuntimeError("Can't get bot instance in sub process.")
def get_config(key: str = None, default: Any = None) -> Any: def get_config(key: str, default: Any = None) -> Any:
""" """
获取配置 获取配置
Args: Args:
@ -256,8 +271,6 @@ def get_config(key: str = None, default: Any = None) -> Any:
Returns: Returns:
Any: 配置值 Any: 配置值
""" """
if key is None:
return get_bot().config
return get_bot().config.get(key, default) return get_bot().config.get(key, default)

View File

@ -15,6 +15,7 @@ from multiprocessing import Pipe
from typing import Any, Callable, Coroutine, Generic, Optional, TypeAlias, TypeVar, get_args from typing import Any, Callable, Coroutine, Generic, Optional, TypeAlias, TypeVar, get_args
from liteyuki.utils import IS_MAIN_PROCESS, is_coroutine_callable, run_coroutine from liteyuki.utils import IS_MAIN_PROCESS, is_coroutine_callable, run_coroutine
from liteyuki.log import logger
T = TypeVar("T") T = TypeVar("T")
@ -217,6 +218,10 @@ class Channel(Generic[T]):
def __next__(self) -> Any: def __next__(self) -> Any:
return self.receive() return self.receive()
def __del__(self):
self.close()
logger.debug(f"Channel {self.name} deleted.")
"""子进程可用的主动和被动通道""" """子进程可用的主动和被动通道"""
active_channel: Optional["Channel"] = None active_channel: Optional["Channel"] = None
@ -232,8 +237,20 @@ if IS_MAIN_PROCESS:
@channel_deliver_passive_channel.on_receive(filter_func=lambda data: data[0] == "set_channel") @channel_deliver_passive_channel.on_receive(filter_func=lambda data: data[0] == "set_channel")
def on_set_channel(data: tuple[str, dict[str, Any]]): def on_set_channel(data: tuple[str, dict[str, Any]]):
name, channel, temp_channel = data[1]["name"], data[1]["channel"], _channel[data[0]] name, channel = data[1]["name"], data[1]["channel"]
temp_channel.send(set_channel(name, channel)) set_channel(name, channel)
@channel_deliver_passive_channel.on_receive(filter_func=lambda data: data[0] == "get_channel")
def on_get_channel(data: tuple[str, dict[str, Any]]):
name, recv_chan = data[1]["name"], data[1]["recv_chan"]
recv_chan.send(get_channel(name))
@channel_deliver_passive_channel.on_receive(filter_func=lambda data: data[0] == "get_channels")
def on_get_channels(data: tuple[str, dict[str, Any]]):
recv_chan = data[1]["recv_chan"]
recv_chan.send(get_channels())
def set_channel(name: str, channel: Channel): def set_channel(name: str, channel: Channel):
@ -266,9 +283,6 @@ def set_channels(channels: dict[str, Channel]):
Args: Args:
channels: 通道名称 channels: 通道名称
""" """
if not IS_MAIN_PROCESS:
raise RuntimeError(f"Function {__name__} should only be called in the main process.")
for name, channel in channels.items(): for name, channel in channels.items():
set_channel(name, channel) set_channel(name, channel)
@ -280,18 +294,38 @@ def get_channel(name: str) -> Channel:
name: 通道名称 name: 通道名称
Returns: Returns:
""" """
if not IS_MAIN_PROCESS: if IS_MAIN_PROCESS:
raise RuntimeError(f"Function {__name__} should only be called in the main process.")
return _channel[name] return _channel[name]
else:
recv_chan = Channel[Channel[Any]]("recv_chan")
channel_deliver_passive_channel.send(
(
"get_channel",
{
"name" : name,
"recv_chan": recv_chan
}
)
)
return recv_chan.receive()
def get_channels() -> dict[str, Channel]: def get_channels() -> dict[str, Channel]:
""" """
获取通道实例 获取通道实例
Returns: Returns:
""" """
if not IS_MAIN_PROCESS: if IS_MAIN_PROCESS:
raise RuntimeError(f"Function {__name__} should only be called in the main process.")
return _channel return _channel
else:
recv_chan = Channel[dict[str, Channel[Any]]]("recv_chan")
channel_deliver_passive_channel.send(
(
"get_channels",
{
"recv_chan": recv_chan
}
)
)
return recv_chan.receive()

View File

@ -19,7 +19,7 @@ from liteyuki.log import logger
from liteyuki.utils import IS_MAIN_PROCESS from liteyuki.utils import IS_MAIN_PROCESS
if TYPE_CHECKING: if TYPE_CHECKING:
from liteyuki.bot import LiteyukiBot from liteyuki.bot.lifespan import Lifespan
from liteyuki.comm.storage import KeyValueStore from liteyuki.comm.storage import KeyValueStore
if IS_MAIN_PROCESS: if IS_MAIN_PROCESS:
@ -76,8 +76,8 @@ class ProcessManager:
进程管理器 进程管理器
""" """
def __init__(self, bot: "LiteyukiBot"): def __init__(self, lifespan: "Lifespan"):
self.bot = bot self.lifespan = lifespan
self.targets: dict[str, tuple[Callable, tuple, dict]] = {} self.targets: dict[str, tuple[Callable, tuple, dict]] = {}
self.processes: dict[str, Process] = {} self.processes: dict[str, Process] = {}
@ -108,14 +108,14 @@ class ProcessManager:
if data == 0: if data == 0:
# 停止 # 停止
logger.info(f"Stopping process {name}") logger.info(f"Stopping process {name}")
self.bot.lifespan.before_process_shutdown() self.lifespan.before_process_shutdown()
self.terminate(name) self.terminate(name)
break break
elif data == 1: elif data == 1:
# 重启 # 重启
logger.info(f"Restarting process {name}") logger.info(f"Restarting process {name}")
self.bot.lifespan.before_process_shutdown() self.lifespan.before_process_shutdown()
self.bot.lifespan.before_process_restart() self.lifespan.before_process_restart()
self.terminate(name) self.terminate(name)
_start_process() _start_process()
continue continue

View File

@ -1,9 +1,8 @@
import multiprocessing
from nonebot.plugin import PluginMetadata from nonebot.plugin import PluginMetadata
from liteyuki.comm import get_channel
from .rt_guide import * from liteyuki import get_bot
from .crt_matchers import * from .crt_matchers import * # type: ignore
from .rt_guide import * # type: ignore
__plugin_meta__ = PluginMetadata( __plugin_meta__ = PluginMetadata(
@ -18,3 +17,4 @@ __plugin_meta__ = PluginMetadata(
"default_enable": True, "default_enable": True,
} }
) )
print(get_bot())