import os import time import uuid from contextlib import asynccontextmanager class DistributedLock: """分布式锁封装,优先使用 Redis,不可用时降级为文件锁""" def __init__(self, name: str, ttl_seconds: int = 600): self.name = name self.ttl = ttl_seconds self.token = str(uuid.uuid4()) self._use_redis = False self._redis = None self._file_path = f".lock_{self.name}" try: import redis # type: ignore from app.settings.config import settings self._redis = redis.Redis( host=getattr(settings, "REDIS_HOST", None) or "", port=getattr(settings, "REDIS_PORT", 6379), db=getattr(settings, "REDIS_DB", 0), password=getattr(settings, "REDIS_PASS", None) or None, socket_timeout=3, ) # 尝试 ping if self._redis.ping(): self._use_redis = True except Exception: self._use_redis = False async def acquire(self) -> bool: """获取锁,返回是否成功""" if self._use_redis and self._redis is not None: try: # NX+EX 设置锁,避免竞争 return bool(self._redis.set(f"lock:{self.name}", self.token, nx=True, ex=self.ttl)) except Exception: pass # 文件锁降级(单机安全) try: os.mkdir(self._file_path) return True except Exception: return False async def release(self) -> None: """释放锁""" if self._use_redis and self._redis is not None: try: # 简单释放;生产建议使用 Lua 脚本确保原子性 key = f"lock:{self.name}" val = self._redis.get(key) if val and val.decode() == self.token: self._redis.delete(key) except Exception: pass try: os.rmdir(self._file_path) except Exception: pass @asynccontextmanager async def context(self): """上下文管理:获取成功才进入""" acquired = await self.acquire() try: if acquired: yield True else: yield False finally: if acquired: await self.release()