112 lines
3.1 KiB
Python
112 lines
3.1 KiB
Python
import redis.asyncio as redis
|
|
import asyncio
|
|
import uuid
|
|
from typing import Optional
|
|
from app.config import settings
|
|
from app.logger import get_redis_logger
|
|
|
|
logger = get_redis_logger()
|
|
|
|
|
|
class RedisLock:
|
|
def __init__(self, redis_client: redis.Redis, key: str, timeout: int = None):
|
|
self.redis_client = redis_client
|
|
self.key = f"lock:{key}"
|
|
self.timeout = timeout or settings.redis_lock_timeout
|
|
self.identifier = str(uuid.uuid4())
|
|
self.acquired = False
|
|
|
|
async def acquire(self) -> bool:
|
|
"""获取分布式锁"""
|
|
logger.debug(f"🔒 尝试获取Redis锁: {self.key}")
|
|
|
|
lua_script = """
|
|
if redis.call("GET", KEYS[1]) == false then
|
|
return redis.call("SETEX", KEYS[1], ARGV[1], ARGV[2])
|
|
else
|
|
return false
|
|
end
|
|
"""
|
|
|
|
result = await self.redis_client.eval(
|
|
lua_script,
|
|
1,
|
|
self.key,
|
|
self.timeout,
|
|
self.identifier
|
|
)
|
|
|
|
self.acquired = bool(result)
|
|
if self.acquired:
|
|
logger.debug(f"✅ Redis锁获取成功: {self.key}")
|
|
else:
|
|
logger.debug(f"❌ Redis锁获取失败: {self.key}")
|
|
|
|
return self.acquired
|
|
|
|
async def release(self) -> bool:
|
|
"""释放分布式锁"""
|
|
if not self.acquired:
|
|
return False
|
|
|
|
lua_script = """
|
|
if redis.call("GET", KEYS[1]) == ARGV[1] then
|
|
return redis.call("DEL", KEYS[1])
|
|
else
|
|
return 0
|
|
end
|
|
"""
|
|
|
|
result = await self.redis_client.eval(
|
|
lua_script,
|
|
1,
|
|
self.key,
|
|
self.identifier
|
|
)
|
|
|
|
self.acquired = False
|
|
return bool(result)
|
|
|
|
async def __aenter__(self):
|
|
"""异步上下文管理器入口"""
|
|
retries = 0
|
|
while retries < settings.redis_lock_max_retries:
|
|
if await self.acquire():
|
|
return self
|
|
await asyncio.sleep(settings.redis_lock_retry_delay)
|
|
retries += 1
|
|
|
|
raise TimeoutError(f"Failed to acquire lock {self.key} after {retries} retries")
|
|
|
|
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
|
"""异步上下文管理器出口"""
|
|
await self.release()
|
|
|
|
|
|
class RedisManager:
|
|
def __init__(self):
|
|
self.redis_client: Optional[redis.Redis] = None
|
|
|
|
async def connect(self):
|
|
"""连接Redis"""
|
|
self.redis_client = redis.from_url(
|
|
settings.redis_url,
|
|
encoding="utf-8",
|
|
decode_responses=True
|
|
)
|
|
await self.redis_client.ping()
|
|
|
|
async def disconnect(self):
|
|
"""断开Redis连接"""
|
|
if self.redis_client:
|
|
await self.redis_client.close()
|
|
|
|
def create_lock(self, key: str, timeout: int = None) -> RedisLock:
|
|
"""创建分布式锁"""
|
|
if not self.redis_client:
|
|
raise RuntimeError("Redis client not connected")
|
|
return RedisLock(self.redis_client, key, timeout)
|
|
|
|
|
|
# 全局Redis管理器实例
|
|
redis_manager = RedisManager() |