ElapsingDreams/LiteyukiBot
1
0
Fork 0
forked from bot/app
Code Pull Requests Activity

⬇️

Browse Source
This commit is contained in:
远野千束
2024-08-29 13:50:12 +08:00
parent 3a3ef4d6ae
commit cb3ee4b72f
93 changed files with 1632 additions and 9271 deletions
Download Patch File Download Diff File Expand all files Collapse all files

79
liteyuki/comm/storage.py
View File

@ -2,13 +2,13 @@
"""
共享内存模块。类似于redis，但是更加轻量级并且线程安全
"""
import asyncio
import threading
from typing import Any, Coroutine, Optional, TypeAlias, Callable
from typing import Any, Callable, Optional
from liteyuki.comm import channel
from liteyuki.comm.channel import Channel, ON_RECEIVE_FUNC, ASYNC_ON_RECEIVE_FUNC
from liteyuki.utils import IS_MAIN_PROCESS, is_coroutine_callable, run_coroutine, run_coroutine_in_thread
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 = {}
@ -31,24 +31,13 @@ def _get_lock(key) -> threading.Lock:
raise RuntimeError("Cannot get lock in sub process.")
class Subscriber:
def __init__(self):
self._subscribers = {}
def receive(self) -> Any:
pass
def unsubscribe(self) -> None:
pass
class KeyValueStore:
def __init__(self):
self._store = {}
self.active_chan = Channel[tuple[str, Optional[dict[str, Any]]]](_id="shared_memory-active")
self.passive_chan = Channel[tuple[str, Optional[dict[str, Any]]]](_id="shared_memory-passive")
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]](_id="shared_memory-publish")
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
@ -170,7 +159,7 @@ class KeyValueStore:
"publish",
{
"channel": channel_,
"data" : data
"data" : data
}
)
)
@ -184,12 +173,8 @@ class KeyValueStore:
Returns:
装饰器
"""
if IS_MAIN_PROCESS and not self.is_main_receive_loop_running:
threading.Thread(target=self._start_receive_loop, daemon=True).start()
shared_memory.is_main_receive_loop_running = True
elif not IS_MAIN_PROCESS and not self.is_sub_receive_loop_running:
threading.Thread(target=self._start_receive_loop, daemon=True).start()
shared_memory.is_sub_receive_loop_running = True
if not IS_MAIN_PROCESS:
raise RuntimeError("Cannot subscribe in sub process.")
def decorator(func: ON_RECEIVE_FUNC) -> ON_RECEIVE_FUNC:
async def wrapper(data: Any):
@ -211,38 +196,29 @@ class KeyValueStore:
return decorator
@staticmethod
def run_subscriber_receive_funcs(channel_: str, data: Any):
async def run_subscriber_receive_funcs(channel_: str, data: Any):
"""
运行订阅者接收函数
Args:
channel_: 频道
data: 数据
"""
if IS_MAIN_PROCESS:
if channel_ in _on_main_subscriber_receive_funcs and _on_main_subscriber_receive_funcs[channel_]:
run_coroutine_in_thread(*[func(data) for func in _on_main_subscriber_receive_funcs[channel_]])
else:
if channel_ in _on_sub_subscriber_receive_funcs and _on_sub_subscriber_receive_funcs[channel_]:
run_coroutine_in_thread(*[func(data) for func in _on_sub_subscriber_receive_funcs[channel_]])
[asyncio.create_task(func(data)) for func in _on_main_subscriber_receive_funcs[channel_]]
def _start_receive_loop(self):
async def start_receive_loop(self):
"""
启动发布订阅接收器循环，在主进程中运行，若有子进程订阅则推送给子进程
"""
if IS_MAIN_PROCESS:
while True:
data = self.active_chan.receive()
if data[0] == "publish":
# 运行主进程订阅函数
self.run_subscriber_receive_funcs(data[1]["channel"], data[1]["data"])
# 推送给子进程
self.publish_channel.send(data)
else:
while True:
data = self.publish_channel.receive()
if data[0] == "publish":
# 运行子进程订阅函数
self.run_subscriber_receive_funcs(data[1]["channel"], data[1]["data"])
if not IS_MAIN_PROCESS:
raise RuntimeError("Cannot start receive loop in sub process.")
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:
@ -262,7 +238,6 @@ 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"]
@ -289,14 +264,6 @@ if IS_MAIN_PROCESS:
recv_chan = data[1]["recv_chan"]
recv_chan.send(shared_memory.get_all())
else:
# 子进程在入口函数中对shared_memory进行初始化
@channel.publish_channel.on_receive()
def on_publish(data: tuple[str, Any]):
channel_, data = data
shared_memory.run_subscriber_receive_funcs(channel_, data)
_ref_count = 0 # import 引用计数, 防止获取空指针
if not IS_MAIN_PROCESS:
if (shared_memory is None) and _ref_count > 1: