mirror of
https://github.com/TriM-Organization/LiteyukiBot-TriM.git
synced 2025-09-17 17:46:22 +00:00
🔀手动Merge轻雪主仓库a77f97f
This commit is contained in:
@ -1,30 +1,38 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Copyright (C) 2020-2024 LiteyukiStudio. All Rights Reserved
|
||||
|
||||
@Time : 2024/7/26 下午10:36
|
||||
@Author : snowykami
|
||||
@Email : snowykami@outlook.com
|
||||
@File : __init__.py
|
||||
@Software: PyCharm
|
||||
该模块用于轻雪主进程和Nonebot子进程之间的通信
|
||||
依赖关系
|
||||
event -> _
|
||||
storage -> channel_
|
||||
rpc -> channel_, storage
|
||||
"""
|
||||
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("无法在子进程中初始化 Channel")
|
||||
_ref_count += 1
|
||||
|
||||
@ -1,219 +1,305 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Copyright (C) 2020-2024 LiteyukiStudio. All Rights Reserved
|
||||
|
||||
@Time : 2024/7/26 下午11:21
|
||||
@Author : snowykami
|
||||
@Email : snowykami@outlook.com
|
||||
@File : channel.py
|
||||
@Software: PyCharm
|
||||
|
||||
本模块定义了一个通用的通道类，用于进程间通信
|
||||
"""
|
||||
import functools
|
||||
import multiprocessing
|
||||
import threading
|
||||
import asyncio
|
||||
from multiprocessing import Pipe
|
||||
from typing import Any, Optional, Callable, Awaitable, List, TypeAlias
|
||||
from uuid import uuid4
|
||||
from typing import Any, Callable, Coroutine, Generic, Optional, TypeAlias, TypeVar, get_args
|
||||
|
||||
from liteyuki.utils import is_coroutine_callable, run_coroutine
|
||||
from liteyuki.log import logger
|
||||
from liteyuki.utils import IS_MAIN_PROCESS, is_coroutine_callable
|
||||
|
||||
SYNC_ON_RECEIVE_FUNC: TypeAlias = Callable[[Any], Any]
|
||||
ASYNC_ON_RECEIVE_FUNC: TypeAlias = Callable[[Any], Awaitable[Any]]
|
||||
ON_RECEIVE_FUNC: TypeAlias = SYNC_ON_RECEIVE_FUNC | ASYNC_ON_RECEIVE_FUNC
|
||||
T = TypeVar("T")
|
||||
|
||||
SYNC_FILTER_FUNC: TypeAlias = Callable[[Any], bool]
|
||||
ASYNC_FILTER_FUNC: TypeAlias = Callable[[Any], Awaitable[bool]]
|
||||
FILTER_FUNC: TypeAlias = SYNC_FILTER_FUNC | ASYNC_FILTER_FUNC
|
||||
SYNC_ON_RECEIVE_FUNC: TypeAlias = Callable[[T], Any] # 同步接收函数
|
||||
ASYNC_ON_RECEIVE_FUNC: TypeAlias = Callable[[T], Coroutine[Any, Any, Any]] # 异步接收函数
|
||||
ON_RECEIVE_FUNC: TypeAlias = SYNC_ON_RECEIVE_FUNC | ASYNC_ON_RECEIVE_FUNC # 接收函数
|
||||
|
||||
IS_MAIN_PROCESS = multiprocessing.current_process().name == "MainProcess"
|
||||
SYNC_FILTER_FUNC: TypeAlias = Callable[[T], bool] # 同步过滤函数
|
||||
ASYNC_FILTER_FUNC: TypeAlias = Callable[[T], Coroutine[Any, Any, bool]] # 异步过滤函数
|
||||
FILTER_FUNC: TypeAlias = SYNC_FILTER_FUNC | ASYNC_FILTER_FUNC # 过滤函数
|
||||
|
||||
_func_id: int = 0
|
||||
_channel: dict[str, "Channel"] = {}
|
||||
_callback_funcs: dict[str, ON_RECEIVE_FUNC] = {}
|
||||
_callback_funcs: dict[int, ON_RECEIVE_FUNC] = {}
|
||||
|
||||
|
||||
class Channel:
|
||||
class Channel(Generic[T]):
|
||||
"""
|
||||
通道类，用于进程间通信，进程内不可用，仅限主进程和子进程之间通信
|
||||
通道类，可以在进程间和进程内通信，双向但同时只能有一个发送者和一个接收者
|
||||
有两种接收工作方式，但是只能选择一种，主动接收和被动接收，主动接收使用 `receive` 方法，被动接收使用 `on_receive` 装饰器
|
||||
"""
|
||||
|
||||
def __init__(self, _id: str):
|
||||
self.main_send_conn, self.sub_receive_conn = Pipe()
|
||||
self.sub_send_conn, self.main_receive_conn = Pipe()
|
||||
self._closed = False
|
||||
self._on_main_receive_funcs: list[str] = []
|
||||
self._on_sub_receive_funcs: list[str] = []
|
||||
self.name: str = _id
|
||||
def __init__(self, name: str, type_check: Optional[bool] = None):
|
||||
"""
|
||||
初始化通道
|
||||
Args:
|
||||
name: 通道ID
|
||||
type_check: 是否开启类型检查, 若为空，则传入泛型默认开启，否则默认关闭
|
||||
"""
|
||||
|
||||
self.is_main_receive_loop_running = False
|
||||
self.is_sub_receive_loop_running = False
|
||||
self.conn_send, self.conn_recv = Pipe()
|
||||
self._conn_send_inner, self._conn_recv_inner = Pipe() # 内部通道，用于子进程通信
|
||||
self._closed = False
|
||||
self._on_main_receive_func_ids: list[int] = []
|
||||
self._on_sub_receive_func_ids: list[int] = []
|
||||
self.name: str = name
|
||||
|
||||
self.is_receive_loop_running = False
|
||||
|
||||
if type_check is None:
|
||||
# 若传入泛型则默认开启类型检查
|
||||
type_check = self._get_generic_type() is not None
|
||||
|
||||
elif type_check:
|
||||
if self._get_generic_type() is None:
|
||||
raise TypeError("Type hint 是强制类型检查之所必须")
|
||||
self.type_check = type_check
|
||||
if name in _channel:
|
||||
raise ValueError(f"Channel {name} 已存在")
|
||||
|
||||
if IS_MAIN_PROCESS:
|
||||
if name in _channel:
|
||||
raise ValueError(f"Channel {name} 已存在")
|
||||
_channel[name] = self
|
||||
logger.debug(f"Channel {name} 已在主进程中初始化")
|
||||
else:
|
||||
logger.debug(f"Channel {name} 已初始化于子进程中，之后应于主进程中手动设置为妙")
|
||||
|
||||
def _get_generic_type(self) -> Optional[type]:
|
||||
"""
|
||||
获取通道传递泛型类型
|
||||
Returns:
|
||||
Optional[type]: 泛型类型
|
||||
"""
|
||||
if hasattr(self, '__orig_class__'):
|
||||
return get_args(self.__orig_class__)[0]
|
||||
return None
|
||||
|
||||
def _validate_structure(self, data: Any, structure: type) -> bool:
|
||||
"""
|
||||
验证数据结构
|
||||
Args:
|
||||
data: 数据
|
||||
structure: 结构
|
||||
Returns:
|
||||
bool: 是否通过验证
|
||||
"""
|
||||
if isinstance(structure, type):
|
||||
return isinstance(data, structure)
|
||||
elif isinstance(structure, tuple):
|
||||
if not isinstance(data, tuple) or len(data) != len(structure):
|
||||
return False
|
||||
return all(self._validate_structure(d, s) for d, s in zip(data, structure))
|
||||
elif isinstance(structure, list):
|
||||
if not isinstance(data, list):
|
||||
return False
|
||||
return all(self._validate_structure(d, structure[0]) for d in data)
|
||||
elif isinstance(structure, dict):
|
||||
if not isinstance(data, dict):
|
||||
return False
|
||||
return all(k in data and self._validate_structure(data[k], structure[k]) for k in structure)
|
||||
return False
|
||||
|
||||
def __str__(self):
|
||||
return f"Channel({self.name})"
|
||||
|
||||
def send(self, data: Any):
|
||||
def send(self, data: T):
|
||||
"""
|
||||
发送数据
|
||||
发送数据，发送函数为同步函数，没有异步的必要
|
||||
Args:
|
||||
data: 数据
|
||||
data (T): 数据
|
||||
"""
|
||||
if self._closed:
|
||||
raise RuntimeError("无法发送至已关闭的通道中")
|
||||
if IS_MAIN_PROCESS:
|
||||
print("主进程发送数据：", data)
|
||||
self.main_send_conn.send(data)
|
||||
else:
|
||||
print("子进程发送数据：", data)
|
||||
self.sub_send_conn.send(data)
|
||||
if self.type_check:
|
||||
_type = self._get_generic_type()
|
||||
if _type is not None and not self._validate_structure(data, _type):
|
||||
raise TypeError(f"该数据必须为 {_type} 实例，而非 {type(data)}")
|
||||
|
||||
def receive(self) -> Any:
|
||||
if self._closed:
|
||||
raise RuntimeError("数据无法向已关闭的 Channel 中发送")
|
||||
self.conn_send.send(data)
|
||||
|
||||
def receive(self) -> T:
|
||||
"""
|
||||
接收数据
|
||||
Args:
|
||||
同步接收数据，会阻塞线程
|
||||
Returns:
|
||||
T: 数据
|
||||
"""
|
||||
if self._closed:
|
||||
raise RuntimeError("无法从已关闭的通道中接收")
|
||||
raise RuntimeError("无法在已关闭的 Channel 中接取数据")
|
||||
|
||||
while True:
|
||||
# 判断receiver是否为None或者receiver是否等于接收者，是则接收数据，否则不动数据
|
||||
if IS_MAIN_PROCESS:
|
||||
data = self.main_receive_conn.recv()
|
||||
print("主进程接收数据：", data)
|
||||
else:
|
||||
data = self.sub_receive_conn.recv()
|
||||
print("子进程接收数据：", data)
|
||||
|
||||
data = self.conn_recv.recv()
|
||||
return data
|
||||
|
||||
def close(self):
|
||||
async def async_receive(self) -> T:
|
||||
"""
|
||||
关闭通道
|
||||
异步接收数据，会挂起等待
|
||||
Returns:
|
||||
T: 数据
|
||||
"""
|
||||
self._closed = True
|
||||
self.sub_receive_conn.close()
|
||||
self.main_send_conn.close()
|
||||
self.sub_send_conn.close()
|
||||
self.main_receive_conn.close()
|
||||
loop = asyncio.get_running_loop()
|
||||
data = await loop.run_in_executor(None, self.receive)
|
||||
return data
|
||||
|
||||
def on_receive(self, filter_func: Optional[FILTER_FUNC] = None) -> Callable[[ON_RECEIVE_FUNC], ON_RECEIVE_FUNC]:
|
||||
def on_receive(self, filter_func: Optional[FILTER_FUNC] = None) -> Callable[[Callable[[T], Any]], Callable[[T], Any]]:
|
||||
"""
|
||||
接收数据并执行函数
|
||||
Args:
|
||||
filter_func: 过滤函数，为None则不过滤
|
||||
filter_func ([`Optional`](https%3A//docs.python.org/3/library/typing.html#typing.Optional)[[`FILTER_FUNC`](#var-FILTER_FUNC)], optional): 过滤函数. Defaults to None.
|
||||
Returns:
|
||||
装饰器，装饰一个函数在接收到数据后执行
|
||||
Callable[[Callable[[T], Any]], Callable[[T], Any]]: 装饰器
|
||||
"""
|
||||
if (not self.is_sub_receive_loop_running) and not IS_MAIN_PROCESS:
|
||||
threading.Thread(target=self._start_sub_receive_loop).start()
|
||||
if not IS_MAIN_PROCESS:
|
||||
raise RuntimeError("on_receive 仅可用于主进程内")
|
||||
|
||||
if (not self.is_main_receive_loop_running) and IS_MAIN_PROCESS:
|
||||
threading.Thread(target=self._start_main_receive_loop).start()
|
||||
def decorator(func: Callable[[T], Any]) -> Callable[[T], Any]:
|
||||
global _func_id
|
||||
|
||||
def decorator(func: ON_RECEIVE_FUNC) -> ON_RECEIVE_FUNC:
|
||||
async def wrapper(data: Any) -> Any:
|
||||
async def wrapper(data: T) -> 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)): # type: ignore
|
||||
return
|
||||
else:
|
||||
if not filter_func(data):
|
||||
return
|
||||
return await func(data)
|
||||
|
||||
function_id = str(uuid4())
|
||||
_callback_funcs[function_id] = wrapper
|
||||
if is_coroutine_callable(func):
|
||||
return await func(data)
|
||||
else:
|
||||
return func(data)
|
||||
|
||||
_callback_funcs[_func_id] = wrapper
|
||||
if IS_MAIN_PROCESS:
|
||||
self._on_main_receive_funcs.append(function_id)
|
||||
self._on_main_receive_func_ids.append(_func_id)
|
||||
else:
|
||||
self._on_sub_receive_funcs.append(function_id)
|
||||
self._on_sub_receive_func_ids.append(_func_id)
|
||||
_func_id += 1
|
||||
return func
|
||||
|
||||
return decorator
|
||||
|
||||
def _run_on_main_receive_funcs(self, data: Any):
|
||||
async def _run_on_receive_funcs(self, data: Any):
|
||||
"""
|
||||
运行接收函数
|
||||
Args:
|
||||
data: 数据
|
||||
"""
|
||||
for func_id in self._on_main_receive_funcs:
|
||||
func = _callback_funcs[func_id]
|
||||
run_coroutine(func(data))
|
||||
|
||||
def _run_on_sub_receive_funcs(self, data: Any):
|
||||
"""
|
||||
运行接收函数
|
||||
Args:
|
||||
data: 数据
|
||||
"""
|
||||
for func_id in self._on_sub_receive_funcs:
|
||||
func = _callback_funcs[func_id]
|
||||
run_coroutine(func(data))
|
||||
|
||||
def _start_main_receive_loop(self):
|
||||
"""
|
||||
开始接收数据
|
||||
"""
|
||||
self.is_main_receive_loop_running = True
|
||||
while not self._closed:
|
||||
data = self.main_receive_conn.recv()
|
||||
self._run_on_main_receive_funcs(data)
|
||||
|
||||
def _start_sub_receive_loop(self):
|
||||
"""
|
||||
开始接收数据
|
||||
"""
|
||||
self.is_sub_receive_loop_running = True
|
||||
while not self._closed:
|
||||
data = self.sub_receive_conn.recv()
|
||||
self._run_on_sub_receive_funcs(data)
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self) -> Any:
|
||||
return self.receive()
|
||||
if IS_MAIN_PROCESS:
|
||||
[asyncio.create_task(_callback_funcs[func_id](data)) for func_id in self._on_main_receive_func_ids]
|
||||
else:
|
||||
[asyncio.create_task(_callback_funcs[func_id](data)) for func_id in self._on_sub_receive_func_ids]
|
||||
|
||||
|
||||
"""默认通道实例，可直接从模块导入使用"""
|
||||
chan = Channel("default")
|
||||
"""子进程可用的主动和被动通道"""
|
||||
active_channel: Channel = Channel(name="active_channel") # 主动通道
|
||||
passive_channel: Channel = Channel(name="passive_channel") # 被动通道
|
||||
publish_channel: Channel[tuple[str, dict[str, Any]]] = Channel(name="publish_channel") # 发布通道
|
||||
"""通道传递通道，主进程创建单例，子进程初始化时实例化"""
|
||||
channel_deliver_active_channel: Channel[Channel[Any]] # 主动通道传递通道
|
||||
channel_deliver_passive_channel: Channel[tuple[str, dict[str, Any]]] # 被动通道传递通道
|
||||
|
||||
if IS_MAIN_PROCESS:
|
||||
channel_deliver_active_channel = Channel(name="channel_deliver_active_channel") # 主动通道传递通道
|
||||
channel_deliver_passive_channel = Channel(name="channel_deliver_passive_channel") # 被动通道传递通道
|
||||
|
||||
|
||||
def set_channel(name: str, channel: 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]]):
|
||||
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"):
|
||||
"""
|
||||
设置通道实例
|
||||
Args:
|
||||
name: 通道名称
|
||||
channel: 通道实例
|
||||
name ([`str`](https%3A//docs.python.org/3/library/stdtypes.html#str)): 通道名称
|
||||
channel ([`Channel`](#class-channel-generic-t)): 通道实例
|
||||
"""
|
||||
_channel[name] = channel
|
||||
if not isinstance(channel, Channel):
|
||||
raise TypeError(f"channel_ 必须为 Channel 实例，而非 {type(channel)}")
|
||||
|
||||
if IS_MAIN_PROCESS:
|
||||
if name in _channel:
|
||||
raise ValueError(f"Channel {name} 已存在")
|
||||
_channel[name] = channel
|
||||
else:
|
||||
# 请求主进程设置通道
|
||||
channel_deliver_passive_channel.send(
|
||||
(
|
||||
"set_channel", {
|
||||
"name" : name,
|
||||
"channel_": channel,
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def set_channels(channels: dict[str, Channel]):
|
||||
def set_channels(channels: dict[str, "Channel"]):
|
||||
"""
|
||||
设置通道实例
|
||||
Args:
|
||||
channels: 通道名称
|
||||
channels ([`dict`](https%3A//docs.python.org/3/library/stdtypes.html#dict)[[`str`](https%3A//docs.python.org/3/library/stdtypes.html#str), [`Channel`](#class-channel-generic-t)]): 通道实例
|
||||
"""
|
||||
for name, channel in channels.items():
|
||||
_channel[name] = channel
|
||||
set_channel(name, channel)
|
||||
|
||||
|
||||
def get_channel(name: str) -> Optional[Channel]:
|
||||
def get_channel(name: str) -> "Channel":
|
||||
"""
|
||||
获取通道实例
|
||||
Args:
|
||||
name: 通道名称
|
||||
name ([`str`](https%3A//docs.python.org/3/library/stdtypes.html#str)): 通道名称
|
||||
Returns:
|
||||
[`Channel`](#class-channel-generic-t): 通道实例
|
||||
"""
|
||||
return _channel.get(name, None)
|
||||
if IS_MAIN_PROCESS:
|
||||
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:
|
||||
[`dict`](https%3A//docs.python.org/3/library/stdtypes.html#dict)[[`str`](https%3A//docs.python.org/3/library/stdtypes.html#str), [`Channel`](#class-channel-generic-t)]: 通道实例
|
||||
"""
|
||||
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()
|
||||
|
||||
@ -1,12 +1,6 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Copyright (C) 2020-2024 LiteyukiStudio. All Rights Reserved
|
||||
|
||||
@Time : 2024/7/26 下午10:47
|
||||
@Author : snowykami
|
||||
@Email : snowykami@outlook.com
|
||||
@File : event.py
|
||||
@Software: PyCharm
|
||||
本模块用于轻雪主进程和子进程之间的通信的事件类
|
||||
"""
|
||||
from typing import Any
|
||||
|
||||
|
||||
26
liteyuki/comm/rpc.py
Normal file
26
liteyuki/comm/rpc.py
Normal file
@ -0,0 +1,26 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
本模块用于实现RPC(基于IPC)通信
|
||||
"""
|
||||
|
||||
from typing import TypeAlias, Callable, Any
|
||||
|
||||
from liteyuki.comm.channel import Channel
|
||||
|
||||
ON_CALLING_FUNC: TypeAlias = Callable[[tuple, dict], Any]
|
||||
|
||||
|
||||
class RPC:
|
||||
"""
|
||||
RPC类
|
||||
"""
|
||||
|
||||
def __init__(self, on_calling: ON_CALLING_FUNC) -> None:
|
||||
self.on_calling = on_calling
|
||||
|
||||
def call(self, args: tuple, kwargs: dict) -> Any:
|
||||
"""
|
||||
调用
|
||||
"""
|
||||
# 获取self.calling函数名
|
||||
return self.on_calling(args, kwargs)
|
||||
48
liteyuki/comm/socks_channel.py
Normal file
48
liteyuki/comm/socks_channel.py
Normal file
@ -0,0 +1,48 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
基于socket的通道
|
||||
"""
|
||||
|
||||
|
||||
class SocksChannel:
|
||||
"""
|
||||
通道类，可以在进程间和进程内通信，双向但同时只能有一个发送者和一个接收者
|
||||
有两种接收工作方式，但是只能选择一种，主动接收和被动接收，主动接收使用 `receive` 方法，被动接收使用 `on_receive` 装饰器
|
||||
"""
|
||||
|
||||
def __init__(self, name: str):
|
||||
"""
|
||||
初始化通道
|
||||
Args:
|
||||
name: 通道ID
|
||||
"""
|
||||
|
||||
self._name = name
|
||||
self._conn_send = None
|
||||
self._conn_recv = None
|
||||
self._closed = False
|
||||
|
||||
def send(self, data):
|
||||
"""
|
||||
发送数据
|
||||
Args:
|
||||
data: 数据
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
def receive(self):
|
||||
"""
|
||||
接收数据
|
||||
Returns:
|
||||
data: 数据
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
def close(self):
|
||||
"""
|
||||
关闭通道
|
||||
"""
|
||||
|
||||
pass
|
||||
247
liteyuki/comm/storage.py
Normal file
247
liteyuki/comm/storage.py
Normal file
@ -0,0 +1,247 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
共享内存模块。类似于redis，但是更加轻量级并且线程安全
|
||||
"""
|
||||
import asyncio
|
||||
import threading
|
||||
from typing import Any, Callable, Optional
|
||||
|
||||
from liteyuki.comm import channel
|
||||
from liteyuki.comm.channel import ASYNC_ON_RECEIVE_FUNC, Channel, ON_RECEIVE_FUNC
|
||||
from liteyuki.utils import (
|
||||
IS_MAIN_PROCESS,
|
||||
is_coroutine_callable,
|
||||
run_coroutine_in_thread,
|
||||
)
|
||||
|
||||
if IS_MAIN_PROCESS:
|
||||
_locks = {}
|
||||
|
||||
_on_main_subscriber_receive_funcs: dict[str, list[ASYNC_ON_RECEIVE_FUNC]] = {} # type: ignore
|
||||
"""主进程订阅者接收函数"""
|
||||
_on_sub_subscriber_receive_funcs: dict[str, list[ASYNC_ON_RECEIVE_FUNC]] = {} # type: ignore
|
||||
"""子进程订阅者接收函数"""
|
||||
|
||||
|
||||
def _get_lock(key) -> threading.Lock:
|
||||
"""
|
||||
获取锁
|
||||
"""
|
||||
if IS_MAIN_PROCESS:
|
||||
if key not in _locks:
|
||||
_locks[key] = threading.Lock()
|
||||
return _locks[key]
|
||||
else:
|
||||
raise RuntimeError("无法在子进程中获取线程锁")
|
||||
|
||||
|
||||
class KeyValueStore:
|
||||
def __init__(self):
|
||||
self._store = {}
|
||||
self.active_chan = Channel[tuple[str, Optional[dict[str, Any]]]](
|
||||
name="shared_memory-active"
|
||||
)
|
||||
self.passive_chan = Channel[tuple[str, Optional[dict[str, Any]]]](
|
||||
name="shared_memory-passive"
|
||||
)
|
||||
|
||||
self.publish_channel = Channel[tuple[str, Any]](name="shared_memory-publish")
|
||||
|
||||
self.is_main_receive_loop_running = False
|
||||
self.is_sub_receive_loop_running = False
|
||||
|
||||
def set(self, key: str, value: Any) -> None:
|
||||
"""
|
||||
设置键值对
|
||||
Args:
|
||||
key: 键
|
||||
value: 值
|
||||
|
||||
"""
|
||||
if IS_MAIN_PROCESS:
|
||||
lock = _get_lock(key)
|
||||
with lock:
|
||||
self._store[key] = value
|
||||
else:
|
||||
# 向主进程发送请求拿取
|
||||
self.passive_chan.send(("set", {"key": key, "value": value}))
|
||||
|
||||
def get(self, key: str, default: Optional[Any] = None) -> Optional[Any]:
|
||||
"""
|
||||
获取键值对
|
||||
Args:
|
||||
key: 键
|
||||
default: 默认值
|
||||
|
||||
Returns:
|
||||
Any: 值
|
||||
"""
|
||||
if IS_MAIN_PROCESS:
|
||||
lock = _get_lock(key)
|
||||
with lock:
|
||||
return self._store.get(key, default)
|
||||
else:
|
||||
recv_chan = Channel[Optional[Any]]("recv_chan")
|
||||
self.passive_chan.send(
|
||||
("get", {"key": key, "default": default, "recv_chan": recv_chan})
|
||||
)
|
||||
return recv_chan.receive()
|
||||
|
||||
def delete(self, key: str, ignore_key_error: bool = True) -> None:
|
||||
"""
|
||||
删除键值对
|
||||
Args:
|
||||
key: 键
|
||||
ignore_key_error: 是否忽略键不存在的错误
|
||||
|
||||
Returns:
|
||||
"""
|
||||
if IS_MAIN_PROCESS:
|
||||
lock = _get_lock(key)
|
||||
with lock:
|
||||
if key in self._store:
|
||||
try:
|
||||
del self._store[key]
|
||||
del _locks[key]
|
||||
except KeyError as e:
|
||||
if not ignore_key_error:
|
||||
raise e
|
||||
else:
|
||||
# 向主进程发送请求删除
|
||||
self.passive_chan.send(("delete", {"key": key}))
|
||||
|
||||
def get_all(self) -> dict[str, Any]:
|
||||
"""
|
||||
获取所有键值对
|
||||
Returns:
|
||||
dict[str, Any]: 键值对
|
||||
"""
|
||||
if IS_MAIN_PROCESS:
|
||||
return self._store
|
||||
else:
|
||||
recv_chan = Channel[dict[str, Any]]("recv_chan")
|
||||
self.passive_chan.send(("get_all", {"recv_chan": recv_chan}))
|
||||
return recv_chan.receive()
|
||||
|
||||
def publish(self, channel_: str, data: Any) -> None:
|
||||
"""
|
||||
发布消息
|
||||
Args:
|
||||
channel_: 频道
|
||||
data: 数据
|
||||
|
||||
Returns:
|
||||
"""
|
||||
self.active_chan.send(("publish", {"channel": channel_, "data": data}))
|
||||
|
||||
def on_subscriber_receive(
|
||||
self, channel_: str
|
||||
) -> Callable[[ON_RECEIVE_FUNC], ON_RECEIVE_FUNC]:
|
||||
"""
|
||||
订阅者接收消息时的回调
|
||||
Args:
|
||||
channel_: 频道
|
||||
|
||||
Returns:
|
||||
装饰器
|
||||
"""
|
||||
if not IS_MAIN_PROCESS:
|
||||
raise RuntimeError("无法订阅一子线程消息")
|
||||
|
||||
def decorator(func: ON_RECEIVE_FUNC) -> ON_RECEIVE_FUNC:
|
||||
async def wrapper(data: Any):
|
||||
if is_coroutine_callable(func):
|
||||
await func(data)
|
||||
else:
|
||||
func(data)
|
||||
|
||||
if IS_MAIN_PROCESS:
|
||||
if channel_ not in _on_main_subscriber_receive_funcs:
|
||||
_on_main_subscriber_receive_funcs[channel_] = []
|
||||
_on_main_subscriber_receive_funcs[channel_].append(wrapper)
|
||||
else:
|
||||
if channel_ not in _on_sub_subscriber_receive_funcs:
|
||||
_on_sub_subscriber_receive_funcs[channel_] = []
|
||||
_on_sub_subscriber_receive_funcs[channel_].append(wrapper)
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
@staticmethod
|
||||
async def run_subscriber_receive_funcs(channel_: str, data: Any):
|
||||
"""
|
||||
运行订阅者接收函数
|
||||
Args:
|
||||
channel_: 频道
|
||||
data: 数据
|
||||
"""
|
||||
[
|
||||
asyncio.create_task(func(data))
|
||||
for func in _on_main_subscriber_receive_funcs[channel_]
|
||||
]
|
||||
|
||||
async def start_receive_loop(self):
|
||||
"""
|
||||
启动发布订阅接收器循环，在主进程中运行，若有子进程订阅则推送给子进程
|
||||
"""
|
||||
|
||||
if not IS_MAIN_PROCESS:
|
||||
raise RuntimeError("无法在子进程中启用订阅接收器循环")
|
||||
while True:
|
||||
data = await self.active_chan.async_receive()
|
||||
if data[0] == "publish":
|
||||
# 运行主进程订阅函数
|
||||
await self.run_subscriber_receive_funcs(
|
||||
data[1]["channel"], data[1]["data"]
|
||||
)
|
||||
# 推送给子进程
|
||||
self.publish_channel.send(data)
|
||||
|
||||
|
||||
class GlobalKeyValueStore:
|
||||
_instance = None
|
||||
_lock = threading.Lock()
|
||||
|
||||
@classmethod
|
||||
def get_instance(cls):
|
||||
if cls._instance is None:
|
||||
with cls._lock:
|
||||
if cls._instance is None:
|
||||
cls._instance = KeyValueStore()
|
||||
return cls._instance
|
||||
|
||||
|
||||
shared_memory: KeyValueStore = GlobalKeyValueStore.get_instance() # 共享内存对象
|
||||
|
||||
# 全局单例访问点
|
||||
if IS_MAIN_PROCESS:
|
||||
|
||||
@shared_memory.passive_chan.on_receive(lambda d: d[0] == "get")
|
||||
def on_get(data: tuple[str, dict[str, Any]]):
|
||||
key = data[1]["key"]
|
||||
default = data[1]["default"]
|
||||
recv_chan = data[1]["recv_chan"]
|
||||
recv_chan.send(shared_memory.get(key, default))
|
||||
|
||||
@shared_memory.passive_chan.on_receive(lambda d: d[0] == "set")
|
||||
def on_set(data: tuple[str, dict[str, Any]]):
|
||||
key = data[1]["key"]
|
||||
value = data[1]["value"]
|
||||
shared_memory.set(key, value)
|
||||
|
||||
@shared_memory.passive_chan.on_receive(lambda d: d[0] == "delete")
|
||||
def on_delete(data: tuple[str, dict[str, Any]]):
|
||||
key = data[1]["key"]
|
||||
shared_memory.delete(key)
|
||||
|
||||
@shared_memory.passive_chan.on_receive(lambda d: d[0] == "get_all")
|
||||
def on_get_all(data: tuple[str, dict[str, Any]]):
|
||||
recv_chan = data[1]["recv_chan"]
|
||||
recv_chan.send(shared_memory.get_all())
|
||||
|
||||
|
||||
_ref_count = 0 # import 引用计数, 防止获取空指针
|
||||
if not IS_MAIN_PROCESS:
|
||||
if (shared_memory is None) and _ref_count > 1:
|
||||
raise RuntimeError("共享内存未初始化")
|
||||
_ref_count += 1
|
||||
Reference in New Issue
Block a user