140ddc70e6
* 优化分段消息发送逻辑，为分段消息添加消息队列 * 删除了不必要的代码 * style: code quality * 将消息队列机制重构为会话锁机制 * perf: narrow the lock scope * refactor: replace get_lock with async context manager for session locks * refactor: optimize session lock management with defaultdict --------- Co-authored-by: Soulter <905617992@qq.com> Co-authored-by: Raven95676 <Raven95676@gmail.com>
30 lines
913 B
Python
30 lines
913 B
Python
import asyncio
|
|
from collections import defaultdict
|
|
from contextlib import asynccontextmanager
|
|
|
|
|
|
class SessionLockManager:
|
|
def __init__(self):
|
|
self._locks: dict[str, asyncio.Lock] = defaultdict(asyncio.Lock)
|
|
self._lock_count: dict[str, int] = defaultdict(int)
|
|
self._access_lock = asyncio.Lock()
|
|
|
|
@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
|
|
|
|
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)
|
|
|
|
|
|
session_lock_manager = SessionLockManager()
|