初始项目

app/redis_lock.py

@@ -0,0 +1,102 @@
+import redis.asyncio as redis
+import asyncio
+import uuid
+from typing import Optional
+from app.config import settings
+
+
+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:
+        """获取分布式锁"""
+        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)
+        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()