refactor: remove enable_lock attribute and enhance locking mechanism in Model class

This commit is contained in:
HibiKier
2025-12-23 23:23:58 +08:00
parent af75d7fc5a
commit 587396eb49
4 changed files with 51 additions and 21 deletions
-2
View File
@@ -42,8 +42,6 @@ class BanConsole(Model):
"""缓存类型""" """缓存类型"""
cache_key_field = ("user_id", "group_id") cache_key_field = ("user_id", "group_id")
"""缓存键字段""" """缓存键字段"""
enable_lock: ClassVar[list[DbLockType]] = [DbLockType.CREATE, DbLockType.UPSERT]
"""开启锁"""
lock_fields: ClassVar[dict[DbLockType, tuple[str, str]]] = { lock_fields: ClassVar[dict[DbLockType, tuple[str, str]]] = {
DbLockType.CREATE: ("user_id", "group_id"), DbLockType.CREATE: ("user_id", "group_id"),
DbLockType.UPSERT: ("user_id", "group_id"), DbLockType.UPSERT: ("user_id", "group_id"),
-2
View File
@@ -96,8 +96,6 @@ class GroupConsole(Model):
"""缓存类型""" """缓存类型"""
cache_key_field = ("group_id", "channel_id") cache_key_field = ("group_id", "channel_id")
"""缓存键字段""" """缓存键字段"""
enable_lock: ClassVar[list[DbLockType]] = [DbLockType.CREATE, DbLockType.UPSERT]
"""开启锁"""
lock_fields: ClassVar[dict[DbLockType, tuple[str, str]]] = { lock_fields: ClassVar[dict[DbLockType, tuple[str, str]]] = {
DbLockType.CREATE: ("group_id", "channel_id"), DbLockType.CREATE: ("group_id", "channel_id"),
DbLockType.UPSERT: ("group_id", "channel_id"), DbLockType.UPSERT: ("group_id", "channel_id"),
-1
View File
@@ -94,7 +94,6 @@ __all__ = [
async def cancel_pending_tasks(): async def cancel_pending_tasks():
"""仅取消 zhenxun 业务协程,避免干扰框架的关机流程。"""
loop = asyncio.get_running_loop() loop = asyncio.get_running_loop()
current = asyncio.current_task(loop=loop) current = asyncio.current_task(loop=loop)
pending = [] pending = []
+51 -16
View File
@@ -22,11 +22,13 @@ class Model(TortoiseModel):
增强的ORM基类,解决锁嵌套问题 增强的ORM基类,解决锁嵌套问题
""" """
# sem_data[cls_name][lock_type] 可以是 Semaphore(全局) # sem_data[cls][lock_type] 可以是 Semaphore(全局)
# 或 dict[key, Semaphore](按键) # 或 dict[key, Semaphore](按键)
sem_data: ClassVar[dict[str, dict[str, Any]]] = {} sem_data: ClassVar[dict[type["Model"], dict[DbLockType, Any]]] = {}
# 跟踪当前协程持有的锁 (lock_type, lock_key) # 跟踪当前协程持有的锁集合 {(cls, lock_type, lock_key), ...}
_current_locks: ClassVar[dict[int, tuple[DbLockType, Any | None]]] = {} _current_locks: ClassVar[
dict[int, set[tuple[type["Model"], DbLockType, Any | None]]]
] = {}
def __init_subclass__(cls, **kwargs): def __init_subclass__(cls, **kwargs):
super().__init_subclass__(**kwargs) super().__init_subclass__(**kwargs)
@@ -81,14 +83,26 @@ class Model(TortoiseModel):
@classmethod @classmethod
def get_semaphore(cls, lock_type: DbLockType, lock_key: Any | None = None): def get_semaphore(cls, lock_type: DbLockType, lock_key: Any | None = None):
enable_lock = getattr(cls, "enable_lock", None) """
if not enable_lock or lock_type not in enable_lock: 获取信号量
设计约定(弃用 enable_lock,仅通过 lock_fields 控制是否启用锁):
- 如果未配置 lock_fields,或其中不存在对应 lock_type,则不加锁
- 如果 lock_fields[lock_type] 配置了按字段的锁(如 tuple[str, ...]),
则调用处按字段值生成 lock_key,在此为不同 lock_key
分配不同信号量,实现「按键」互斥
- 如仅需全局锁,可在 lock_fields 中声明该 lock_type,
且在 _lock_context 传入 lock_key=None
"""
lock_fields: dict[DbLockType, Any] = getattr(cls, "lock_fields", {}) or {}
# 未在 lock_fields 中声明的 lock_type 不加锁
if lock_type not in lock_fields:
return None return None
cls_sem = cls.sem_data.setdefault(cls.__name__, {}) cls_sem = cls.sem_data.setdefault(cls, {})
# 是否配置了按字段的锁
lock_fields: dict[DbLockType, Any] = getattr(cls, "lock_fields", {}) or {} # 配置了按字段的锁并且提供了具体的 lock_key 时,使用「按键」锁
if lock_type in lock_fields and lock_key is not None: if lock_key is not None:
keyed = cls_sem.setdefault(lock_type, {}) keyed = cls_sem.setdefault(lock_type, {})
if not isinstance(keyed, dict): if not isinstance(keyed, dict):
# 兼容历史数据,重置为按键字典 # 兼容历史数据,重置为按键字典
@@ -109,7 +123,13 @@ class Model(TortoiseModel):
def _require_lock(cls, lock_type: DbLockType, lock_key: Any | None) -> bool: def _require_lock(cls, lock_type: DbLockType, lock_key: Any | None) -> bool:
"""检查是否需要真正加锁""" """检查是否需要真正加锁"""
task_id = id(asyncio.current_task()) task_id = id(asyncio.current_task())
return cls._current_locks.get(task_id) != (lock_type, lock_key) held = cls._current_locks.get(task_id)
if not held:
return True
# 同一协程内,如果已经持有完全相同的一把锁
# (同一模型 + 同一 lock_type + 同一 lock_key),视为重入,
# 不再重复加锁,避免自锁
return (cls, lock_type, lock_key) not in held
@classmethod @classmethod
@contextlib.asynccontextmanager @contextlib.asynccontextmanager
@@ -118,13 +138,28 @@ class Model(TortoiseModel):
task_id = id(asyncio.current_task()) task_id = id(asyncio.current_task())
need_lock = cls._require_lock(lock_type, lock_key) need_lock = cls._require_lock(lock_type, lock_key)
if need_lock and (sem := cls.get_semaphore(lock_type, lock_key)): if not need_lock:
cls._current_locks[task_id] = (lock_type, lock_key) # 已经持有这把锁,直接透传,支持可重入
yield
return
sem = cls.get_semaphore(lock_type, lock_key)
if not sem:
# 对于未启用锁的场景,直接继续执行
yield
return
lock_id = (cls, lock_type, lock_key)
held = cls._current_locks.setdefault(task_id, set())
held.add(lock_id)
try:
async with sem: async with sem:
yield yield
cls._current_locks.pop(task_id, None) finally:
else: # 安全移除当前锁记录
yield held.discard(lock_id)
if not held:
cls._current_locks.pop(task_id, None)
@classmethod @classmethod
async def create( async def create(