perf: lock user creation to reduce timeout under load

This commit is contained in:
HibiKier
2025-12-23 17:11:35 +08:00
parent cb9c5834df
commit ed23ad319a
+18 -3
View File
@@ -1,3 +1,6 @@
import asyncio
from typing import ClassVar
from tortoise import fields from tortoise import fields
from tortoise.exceptions import IntegrityError from tortoise.exceptions import IntegrityError
@@ -37,20 +40,32 @@ class UserConsole(Model):
"""缓存类型""" """缓存类型"""
cache_key_field = "user_id" cache_key_field = "user_id"
"""缓存键字段""" """缓存键字段"""
_create_locks: ClassVar[dict[str, asyncio.Lock]] = {}
@classmethod @classmethod
async def get_user(cls, user_id: str, platform: str | None = None) -> "UserConsole": async def get_user(cls, user_id: str, platform: str | None = None) -> "UserConsole":
if user := await cls.get_or_none(user_id=user_id): if user := await cls.get_or_none(user_id=user_id):
return user return user
lock = cls._create_locks.setdefault(user_id, asyncio.Lock())
async with lock:
# 双重检查,避免锁等待后重复创建
if user := await cls.get_or_none(user_id=user_id):
return user
for _ in range(3):
try: try:
user = await cls.create( uid = await cls.get_new_uid()
return await cls.create(
user_id=user_id, user_id=user_id,
uid=await cls.get_new_uid(), uid=uid,
platform=platform, platform=platform,
) )
return user
except IntegrityError: except IntegrityError:
# 可能是 uid 或 user_id 竞争,短暂等待再重试
await asyncio.sleep(0.05)
# 多次重试仍失败,按 user_id 再查一遍兜底
return await cls.get(user_id=user_id) return await cls.get(user_id=user_id)
@classmethod @classmethod