fix: 重构 SessionLockManager 以支持每个事件循环的独立锁管理，并添加单元测试

Author: whatevertogo

astrbot/core/utils/session_lock.py

@@ -1,29 +1,54 @@
 import asyncio
+import threading
+import weakref
 from collections import defaultdict
 from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+
+
+@dataclass
+class _LoopLockState:
+    access_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
+    locks: dict[str, asyncio.Lock] = field(default_factory=dict)
+    lock_count: dict[str, int] = field(default_factory=lambda: defaultdict(int))
 
 
 class SessionLockManager:
     def __init__(self) -> None:
-        self._locks: dict[str, asyncio.Lock] = defaultdict(asyncio.Lock)
-        self._lock_count: dict[str, int] = defaultdict(int)
-        self._access_lock = asyncio.Lock()
+        self._state_guard = threading.Lock()
+        self._loop_states: weakref.WeakKeyDictionary[
+            asyncio.AbstractEventLoop, _LoopLockState
+        ] = weakref.WeakKeyDictionary()
+
+    def _get_loop_state(self) -> _LoopLockState:
+        loop = asyncio.get_running_loop()
+        with self._state_guard:
+            state = self._loop_states.get(loop)
+            if state is None:
+                state = _LoopLockState()
+                self._loop_states[loop] = state
+            return state
 
     @asynccontextmanager
     async def acquire_lock(self, session_id: str):
-        async with self._access_lock:
-            lock = self._locks[session_id]
-            self._lock_count[session_id] += 1
+        state = self._get_loop_state()
+
+        async with state.access_lock:
+            lock = state.locks.get(session_id)
+            if lock is None:
+                lock = asyncio.Lock()
+                state.locks[session_id] = lock
+            state.lock_count[session_id] += 1
 
         try:
             async with lock:
                 yield
         finally:
-            async with self._access_lock:
-                self._lock_count[session_id] -= 1
-                if self._lock_count[session_id] == 0:
-                    self._locks.pop(session_id, None)
-                    self._lock_count.pop(session_id, None)
+            async with state.access_lock:
+                state.lock_count[session_id] -= 1
+                if state.lock_count[session_id] == 0:
+                    state.locks.pop(session_id, None)
+                    state.lock_count.pop(session_id, None)
 
 
 session_lock_manager = SessionLockManager()

autogen-tests

@@ -0,0 +1 @@
+Subproject commit 13e144e5476a76ca0d76bf4f07a6401d133a03ed

tests/unit/test_session_lock.py

@@ -0,0 +1,32 @@
+import asyncio
+import threading
+
+from astrbot.core.utils.session_lock import SessionLockManager
+
+
+def _run_worker(manager: SessionLockManager, errors: list[BaseException]) -> None:
+    async def worker():
+        for _ in range(200):
+            async with manager.acquire_lock("shared-session"):
+                await asyncio.sleep(0)
+
+    try:
+        asyncio.run(worker())
+    except BaseException as exc:  # noqa: BLE001
+        errors.append(exc)
+
+
+def test_session_lock_manager_isolated_per_event_loop():
+    manager = SessionLockManager()
+    errors: list[BaseException] = []
+
+    threads = [
+        threading.Thread(target=_run_worker, args=(manager, errors)),
+        threading.Thread(target=_run_worker, args=(manager, errors)),
+    ]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join()
+
+    assert not errors