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

liteyuki/comm/__init__.py

@@ -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

liteyuki/comm/channel.py

@@ -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

liteyuki/comm/shared_memory.py

@@ -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

liteyuki/comm/storage.py

@@ -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