75 lines
2.4 KiB
Python
75 lines
2.4 KiB
Python
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() |