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

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 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.core.manager import ProcessManager
from liteyuki.log import init_log, logger
@ -24,28 +25,40 @@ __all__ = [
class LiteyukiBot:
def __init__(self, *args, **kwargs):
def __init__(self, *args, **kwargs) -> None:
"""
初始化轻雪实例
Args:
*args:
**kwargs: 配置
"""
"""常规操作"""
print_logo()
global _BOT_INSTANCE
_BOT_INSTANCE = self # 引用
"""配置"""
self.config: dict[str, Any] = kwargs
"""初始化"""
self.init(**self.config) # 初始化
logger.info("Liteyuki is initializing...")
"""生命周期管理"""
self.lifespan = Lifespan()
self.process_manager: ProcessManager = ProcessManager(lifespan=self.lifespan)
self.process_manager: ProcessManager = ProcessManager(bot=self)
"""事件循环"""
self.loop = asyncio.new_event_loop()
asyncio.set_event_loop(self.loop)
self.stop_event = threading.Event()
self.call_restart_count = 0
"""插件加载"""
load_plugins("liteyuki/plugins") # 加载轻雪插件
"""信号处理"""
signal.signal(signal.SIGINT, self._handle_exit)
signal.signal(signal.SIGTERM, self._handle_exit)
atexit.register(self.process_manager.terminate_all) # 注册退出时的函数
@ -229,15 +242,17 @@ class LiteyukiBot:
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:
LiteyukiBot: 当前的轻雪实例
"""
if IS_MAIN_PROCESS:
if _BOT_INSTANCE is None:
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.")
def get_config(key: str = None, default: Any = None) -> Any:
def get_config(key: str, default: Any = None) -> Any:
"""
获取配置
Args:
@ -256,8 +271,6 @@ def get_config(key: str = None, default: Any = None) -> Any:
Returns:
Any: 配置值
"""
if key is None:
return get_bot().config
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 liteyuki.utils import IS_MAIN_PROCESS, is_coroutine_callable, run_coroutine
from liteyuki.log import logger
T = TypeVar("T")
@ -217,6 +218,10 @@ class Channel(Generic[T]):
def __next__(self) -> Any:
return self.receive()
def __del__(self):
self.close()
logger.debug(f"Channel {self.name} deleted.")
"""子进程可用的主动和被动通道"""
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")
def on_set_channel(data: tuple[str, dict[str, Any]]):
name, channel, temp_channel = data[1]["name"], data[1]["channel"], _channel[data[0]]
temp_channel.send(set_channel(name, channel))
name, channel = data[1]["name"], data[1]["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):
@ -266,9 +283,6 @@ def set_channels(channels: dict[str, Channel]):
Args:
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():
set_channel(name, channel)
@ -280,10 +294,21 @@ def get_channel(name: str) -> Channel:
name: 通道名称
Returns:
"""
if not IS_MAIN_PROCESS:
raise RuntimeError(f"Function {__name__} should only be called in the main process.")
if IS_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]:
@ -291,7 +316,16 @@ def get_channels() -> dict[str, Channel]:
获取通道实例
Returns:
"""
if not IS_MAIN_PROCESS:
raise RuntimeError(f"Function {__name__} should only be called in the main process.")
return _channel
if IS_MAIN_PROCESS:
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
if TYPE_CHECKING:
from liteyuki.bot import LiteyukiBot
from liteyuki.bot.lifespan import Lifespan
from liteyuki.comm.storage import KeyValueStore
if IS_MAIN_PROCESS:
@ -76,8 +76,8 @@ class ProcessManager:
进程管理器
"""
def __init__(self, bot: "LiteyukiBot"):
self.bot = bot
def __init__(self, lifespan: "Lifespan"):
self.lifespan = lifespan
self.targets: dict[str, tuple[Callable, tuple, dict]] = {}
self.processes: dict[str, Process] = {}
@ -108,14 +108,14 @@ class ProcessManager:
if data == 0:
# 停止
logger.info(f"Stopping process {name}")
self.bot.lifespan.before_process_shutdown()
self.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.lifespan.before_process_shutdown()
self.lifespan.before_process_restart()
self.terminate(name)
_start_process()
continue

View File

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