1
0
forked from bot/app

新增线程安全共享内存储存器

This commit is contained in:
远野千束 2024-08-16 21:38:22 +08:00
parent adc9b76688
commit 222250bc41
15 changed files with 246 additions and 66 deletions

View File

@ -7,7 +7,6 @@ from liteyuki.bot import (
from liteyuki.comm import (
Channel,
chan,
Event
)
@ -17,7 +16,19 @@ from liteyuki.plugin import (
)
from liteyuki.log import (
logger,
init_log
init_log,
logger
)
__all__ = [
"LiteyukiBot",
"get_bot",
"get_config",
"get_config_with_compat",
"Channel",
"Event",
"load_plugin",
"load_plugins",
"init_log",
"logger"
]

View File

@ -47,9 +47,9 @@ class LiteyukiBot:
"""
启动逻辑
"""
self.lifespan.before_start() # 启动前钩子
self.lifespan.before_start() # 启动前钩子
self.process_manager.start_all()
self.lifespan.after_start() # 启动后钩子
self.lifespan.after_start() # 启动后钩子
self.keep_alive()
def keep_alive(self):
@ -98,10 +98,9 @@ class LiteyukiBot:
Args:
name: 进程名称, 默认为None, 所有进程
Returns:
"""
self.loop.create_task(self.lifespan.before_process_shutdown()) # 重启前钩子
self.loop.create_task(self.lifespan.before_process_shutdown()) # 停止前钩子
self.lifespan.before_process_shutdown() # 重启前钩子
self.lifespan.before_process_shutdown() # 停止前钩子
if name is not None:
chan_active = get_channel(f"{name}-active")

View File

@ -11,20 +11,31 @@ Copyright (C) 2020-2024 LiteyukiStudio. All Rights Reserved
"""
from liteyuki.comm.channel import (
Channel,
chan,
get_channel,
set_channel,
set_channels,
get_channels
get_channels,
active_channel,
passive_channel
)
from liteyuki.comm.event import Event
__all__ = [
"Channel",
"chan",
"Event",
"get_channel",
"set_channel",
"set_channels",
"get_channels"
"get_channels",
"active_channel",
"passive_channel"
]
from liteyuki.utils import IS_MAIN_PROCESS
# 第一次引用必定为赋值
_ref_count = 0
if not IS_MAIN_PROCESS:
if (active_channel is None or passive_channel is None) and _ref_count > 0:
raise RuntimeError("Error: Channel not initialized in sub process")
_ref_count += 1

View File

@ -17,7 +17,7 @@ from multiprocessing import Pipe
from typing import Any, Optional, Callable, Awaitable, List, TypeAlias
from uuid import uuid4
from liteyuki.utils import is_coroutine_callable, run_coroutine
from liteyuki.utils import IS_MAIN_PROCESS, is_coroutine_callable, run_coroutine
SYNC_ON_RECEIVE_FUNC: TypeAlias = Callable[[Any], Any]
ASYNC_ON_RECEIVE_FUNC: TypeAlias = Callable[[Any], Awaitable[Any]]
@ -27,11 +27,13 @@ SYNC_FILTER_FUNC: TypeAlias = Callable[[Any], bool]
ASYNC_FILTER_FUNC: TypeAlias = Callable[[Any], Awaitable[bool]]
FILTER_FUNC: TypeAlias = SYNC_FILTER_FUNC | ASYNC_FILTER_FUNC
_channel: dict[str, "Channel"] = {}
_callback_funcs: dict[str, ON_RECEIVE_FUNC] = {}
"""子进程可用的主动和被动通道"""
active_channel: Optional["Channel"] = None
passive_channel: Optional["Channel"] = None
class Channel:
"""
@ -40,8 +42,6 @@ class Channel:
"""
def __init__(self, _id: str):
# self.main_send_conn, self.sub_receive_conn = Pipe()
# self.sub_send_conn, self.main_receive_conn = Pipe()
self.conn_send, self.conn_recv = Pipe()
self._closed = False
self._on_main_receive_funcs: list[str] = []
@ -102,12 +102,16 @@ class Channel:
async def wrapper(data: Any) -> Any:
if filter_func is not None:
if is_coroutine_callable(filter_func):
if not await filter_func(data):
if not (await filter_func(data)):
return
else:
if not filter_func(data):
return
return await func(data)
if is_coroutine_callable(func):
return await func(data)
else:
return func(data)
function_id = str(uuid4())
_callback_funcs[function_id] = wrapper
@ -164,10 +168,6 @@ class Channel:
return self.receive()
"""默认通道实例,可直接从模块导入使用"""
chan = Channel("default")
def set_channel(name: str, channel: Channel):
"""
设置通道实例
@ -175,6 +175,9 @@ def set_channel(name: str, channel: Channel):
name: 通道名称
channel: 通道实例
"""
if not IS_MAIN_PROCESS:
raise RuntimeError(f"Function {__name__} should only be called in the main process.")
if not isinstance(channel, Channel):
raise TypeError(f"channel must be an instance of Channel, {type(channel)} found")
_channel[name] = channel
@ -186,6 +189,9 @@ 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)
@ -197,6 +203,9 @@ def get_channel(name: str) -> Optional[Channel]:
name: 通道名称
Returns:
"""
if not IS_MAIN_PROCESS:
raise RuntimeError(f"Function {__name__} should only be called in the main process.")
return _channel.get(name, None)
@ -205,4 +214,7 @@ 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

View File

@ -1,14 +0,0 @@
# -*- coding: utf-8 -*-
"""
共享内存模块类似于redis但是更加轻量级
"""
memory_database = {}
def set_memory(key: str, value: any) -> None:
pass
def get_mem_data(key: str) -> any:
pass

114
liteyuki/comm/storage.py Normal file
View File

@ -0,0 +1,114 @@
# -*- coding: utf-8 -*-
"""
共享内存模块类似于redis但是更加轻量级并且线程安全
"""
import threading
from typing import Any, Optional
from liteyuki.utils import IS_MAIN_PROCESS
from liteyuki.comm.channel import Channel
if IS_MAIN_PROCESS:
_locks = {}
def _get_lock(key):
if IS_MAIN_PROCESS:
if key not in _locks:
_locks[key] = threading.Lock()
return _locks[key]
else:
raise RuntimeError("Cannot get lock in sub process.")
class KeyValueStore:
def __init__(self):
self._store = {}
self.active_chan = Channel(_id="shared_memory-active")
self.passive_chan = Channel(_id="shared_memory-passive")
def set(self, key: str, value: any) -> None:
if IS_MAIN_PROCESS:
lock = _get_lock(key)
with lock:
self._store[key] = value
else:
# 向主进程发送请求拿取
self.passive_chan.send(("set", key, value))
def get(self, key: str, default: Optional[any] = None) -> any:
if IS_MAIN_PROCESS:
lock = _get_lock(key)
with lock:
return self._store.get(key, default)
else:
self.passive_chan.send(("get", key, default))
return self.active_chan.receive()
def delete(self, key: str) -> None:
if IS_MAIN_PROCESS:
lock = _get_lock(key)
with lock:
if key in self._store:
del self._store[key]
del _locks[key]
else:
# 向主进程发送请求删除
self.passive_chan.send(("delete", key))
def get_all(self) -> dict[str, any]:
if IS_MAIN_PROCESS:
return self._store
else:
self.passive_chan.send(("get_all",))
return self.active_chan.receive()
class GlobalKeyValueStore:
_instance = None
_lock = threading.Lock()
@classmethod
def get_instance(cls):
if IS_MAIN_PROCESS:
if cls._instance is None:
with cls._lock:
if cls._instance is None:
cls._instance = KeyValueStore()
return cls._instance
else:
raise RuntimeError("Cannot get instance in sub process.")
shared_memory: Optional[KeyValueStore] = None
# 全局单例访问点
if IS_MAIN_PROCESS:
shared_memory = GlobalKeyValueStore.get_instance()
@shared_memory.passive_chan.on_receive(lambda d: d[0] == "get")
def on_get(d):
print(shared_memory.get_all())
shared_memory.active_chan.send(shared_memory.get(d[1], d[2]))
print("发送数据:", shared_memory.get(d[1], d[2]))
@shared_memory.passive_chan.on_receive(lambda d: d[0] == "set")
def on_set(d):
shared_memory.set(d[1], d[2])
@shared_memory.passive_chan.on_receive(lambda d: d[0] == "delete")
def on_delete(d):
shared_memory.delete(d[1])
else:
shared_memory = None
_ref_count = 0 # 引用计数
if not IS_MAIN_PROCESS:
if (shared_memory is None) and _ref_count > 1:
raise RuntimeError("Shared memory not initialized.")
_ref_count += 1

View File

@ -16,7 +16,7 @@ import toml
import yaml
from pydantic import BaseModel
from liteyuki import logger
from liteyuki.log import logger
_SUPPORTED_CONFIG_FORMATS = (".yaml", ".yml", ".json", ".toml")

View File

@ -10,18 +10,21 @@ Copyright (C) 2020-2024 LiteyukiStudio. All Rights Reserved
"""
import atexit
import threading
import signal
import threading
from multiprocessing import Process
from typing import Any, Callable, Optional, Protocol, TYPE_CHECKING, TypeAlias
from typing import Any, Callable, TYPE_CHECKING, TypeAlias
from liteyuki.comm import Channel, get_channel, set_channels
from liteyuki.comm.storage import shared_memory
from liteyuki.log import logger
from liteyuki.utils import IS_MAIN_PROCESS
TARGET_FUNC: TypeAlias = Callable[[Channel, Channel, ...], Any]
TARGET_FUNC: TypeAlias = Callable[..., Any]
if TYPE_CHECKING:
from liteyuki.bot import LiteyukiBot
from liteyuki.comm.storage import KeyValueStore
TIMEOUT = 10
@ -30,9 +33,29 @@ __all__ = [
]
# Update the delivery_channel_wrapper function to return the top-level wrapper
def _delivery_channel_wrapper(func: TARGET_FUNC, chan_active: Channel, chan_passive: Channel, sm: "KeyValueStore", *args, **kwargs):
"""
子进程入口函数
"""
# 给子进程设置通道
if IS_MAIN_PROCESS:
raise RuntimeError("Function should only be called in a sub process.")
from liteyuki.comm import channel
channel.active_channel = chan_active
channel.passive_channel = chan_passive
# 给子进程创建共享内存实例
from liteyuki.comm import storage
storage.shared_memory = sm
func(*args, **kwargs)
class ProcessManager:
"""
在主进程中被调用
进程管理器
"""
def __init__(self, bot: "LiteyukiBot"):
@ -61,7 +84,6 @@ class ProcessManager:
process = Process(target=self.targets[name][0], args=self.targets[name][1],
kwargs=self.targets[name][2])
self.processes[name] = process
process.start()
# 启动进程并监听信号
@ -114,9 +136,9 @@ class ProcessManager:
kwargs = {}
chan_active = Channel(_id=f"{name}-active")
chan_passive = Channel(_id=f"{name}-passive")
kwargs["chan_active"] = chan_active
kwargs["chan_passive"] = chan_passive
self.targets[name] = (target, args, kwargs)
self.targets[name] = (_delivery_channel_wrapper, (target, chan_active, chan_passive, shared_memory, *args), kwargs)
# 主进程通道
set_channels(
{
f"{name}-active" : chan_active,
@ -124,7 +146,7 @@ class ProcessManager:
}
)
def join(self):
def join_all(self):
for name, process in self.targets:
process.join()

View File

@ -8,7 +8,7 @@ Copyright (C) 2020-2024 LiteyukiStudio. All Rights Reserved
@File : __init__.py.py
@Software: PyCharm
"""
from liteyuki import get_config_with_compat, load_plugin
from liteyuki import get_config, load_plugin, get_bot
from liteyuki.plugin import PluginMetadata, load_plugins
__plugin_meta__ = PluginMetadata(
@ -18,9 +18,17 @@ __plugin_meta__ = PluginMetadata(
description="插件加载器,用于加载轻雪原生插件"
)
load_plugins("src/liteyuki_plugins")
for plugin in get_config_with_compat("liteyuki.plugins", ("plugins", ), []):
load_plugin(plugin)
for plugin_dir in get_config_with_compat("liteyuki.plugin_dirs", ("plugins_dirs", ), []):
load_plugins(plugin_dir)
def default_plugins_loader():
"""
默认插件加载器应在初始化时调用
"""
load_plugins("src/liteyuki_plugins")
for plugin in get_config("liteyuki.plugins", []):
load_plugin(plugin)
for plugin_dir in get_config("liteyuki.plugin_dirs", []):
load_plugins(plugin_dir)
default_plugins_loader()

View File

@ -12,6 +12,7 @@ from liteyuki.log import logger
IS_MAIN_PROCESS = multiprocessing.current_process().name == "MainProcess"
def is_coroutine_callable(call: Callable[..., Any]) -> bool:
"""
判断是否为协程可调用对象
@ -87,5 +88,6 @@ def async_wrapper(func: Callable[..., Any]) -> Callable[..., Coroutine]:
async def wrapper(*args, **kwargs):
return func(*args, **kwargs)
wrapper.__signature__ = inspect.signature(func)
return wrapper

View File

@ -8,7 +8,7 @@ from src.utils.base.data_manager import InstalledPlugin, plugin_db
from src.utils.base.resource import load_resources
from src.utils.message.tools import check_for_package
from liteyuki import get_bot, chan
from liteyuki import get_bot
from nonebot_plugin_apscheduler import scheduler

View File

@ -8,12 +8,9 @@ Copyright (C) 2020-2024 LiteyukiStudio. All Rights Reserved
@File : __init__.py.py
@Software: PyCharm
"""
import asyncio
import time
import nonebot
from liteyuki.comm import Channel, set_channel
from liteyuki.core import IS_MAIN_PROCESS
from liteyuki.plugin import PluginMetadata
from .nb_utils import adapter_manager, driver_manager
@ -23,20 +20,15 @@ __plugin_meta__ = PluginMetadata(
)
def nb_run(chan_active: "Channel", chan_passive: "Channel", *args, **kwargs):
def nb_run(*args, **kwargs):
"""
初始化NoneBot并运行在子进程
Args:
chan_active:
chan_passive:
**kwargs:
Returns:
"""
# 给子进程传递通道对象
set_channel("nonebot-active", chan_active)
set_channel("nonebot-passive", chan_passive)
kwargs.update(kwargs.get("nonebot", {})) # nonebot配置优先
nonebot.init(**kwargs)

View File

@ -7,4 +7,5 @@ Copyright (C) 2020-2024 LiteyukiStudio. All Rights Reserved
@Email : snowykami@outlook.com
@File : __init__.py
@Software: PyCharm
"""
"""
from .after_start import *

View File

@ -11,13 +11,15 @@ Copyright (C) 2020-2024 LiteyukiStudio. All Rights Reserved
import time
from liteyuki import get_bot
from liteyuki.comm.storage import shared_memory
liteyuki = get_bot()
@liteyuki.on_after_start
@liteyuki.on_before_start
def save_startup_timestamp():
"""
储存启动的时间戳
"""
startup_timestamp = time.time()
shared_memory.set("startup_timestamp", startup_timestamp)

View File

@ -0,0 +1,20 @@
# -*- coding: utf-8 -*-
"""
Copyright (C) 2020-2024 LiteyukiStudio. All Rights Reserved
@Time : 2024/8/16 下午8:30
@Author : snowykami
@Email : snowykami@outlook.com
@File : ts_ly_comm.py
@Software: PyCharm
"""
from nonebot.plugin import PluginMetadata
from liteyuki.comm.storage import shared_memory
__plugin_meta__ = PluginMetadata(
name="轻雪通信测试",
description="用于测试轻雪插件通信",
usage="不面向用户",
)
print("共享内存数据:", shared_memory.get("startup_timestamp", default=None))