""" Redis Client - AWOOOI 分散式狀態儲存 ===================================== Phase 6.1.1: Multi-Sig Redis 遷移 Features: - 非同步連線池 (Connection Pool) - Lifespan 管理 (啟動/關閉) - 分散式鎖 (Distributed Lock) - 環境變數驅動 (禁止硬編碼 IP) 統帥鐵律: - 所有 Redis 操作必須使用此模組 - 禁止在其他地方直接建立 Redis 連線 """ import asyncio from collections.abc import AsyncGenerator from contextlib import asynccontextmanager import redis.asyncio as redis import structlog from src.core.config import settings logger = structlog.get_logger(__name__) # ============================================================================= # Connection Pool (API 用 - 短超時) # ============================================================================= _redis_pool: redis.Redis | None = None # ============================================================================= # Worker 專屬連線池 (長超時，用於 XREADGROUP 阻塞操作) # ============================================================================= _worker_redis_pool: redis.Redis | None = None async def init_redis_pool() -> redis.Redis: """ 初始化 Redis 連線池 統帥鐵律: 連線池在 Lifespan 啟動時建立 """ global _redis_pool if _redis_pool is not None: return _redis_pool _redis_pool = redis.from_url( settings.REDIS_URL, encoding="utf-8", decode_responses=True, max_connections=20, socket_timeout=5.0, socket_connect_timeout=5.0, ) # 測試連線 try: await _redis_pool.ping() logger.info( "redis_pool_initialized", url=settings.REDIS_URL.split("@")[-1], # 隱藏密碼 ) except redis.ConnectionError as e: logger.error("redis_connection_failed", error=str(e)) raise return _redis_pool async def close_redis_pool() -> None: """ 關閉 Redis 連線池 統帥鐵律: 連線池在 Lifespan 關閉時回收 """ global _redis_pool if _redis_pool is not None: await _redis_pool.close() _redis_pool = None logger.info("redis_pool_closed") def get_redis() -> redis.Redis: """ 取得 Redis 連線 (API 用，短超時) Raises: RuntimeError: 若連線池未初始化 """ if _redis_pool is None: raise RuntimeError("Redis pool not initialized. Call init_redis_pool() first.") return _redis_pool # ============================================================================= # Worker 專屬連線池 (長超時) # ============================================================================= async def init_worker_redis_pool() -> redis.Redis: """ 初始化 Worker 專屬 Redis 連線池 統帥鐵律 2026-03-23: - Worker 使用 XREADGROUP 阻塞操作，需要長超時 - 絕對禁止與 API 共用短超時連線池 - socket_timeout=None 表示無限等待 """ global _worker_redis_pool if _worker_redis_pool is not None: return _worker_redis_pool _worker_redis_pool = redis.from_url( settings.REDIS_URL, encoding="utf-8", decode_responses=True, max_connections=5, # Worker 不需要太多連線 socket_timeout=None, # 無限等待 (XREADGROUP 阻塞操作) socket_connect_timeout=10.0, # 連線超時仍需設定 ) # 測試連線 try: await _worker_redis_pool.ping() logger.info( "worker_redis_pool_initialized", url=settings.REDIS_URL.split("@")[-1], socket_timeout="None (unlimited)", ) except redis.ConnectionError as e: logger.error("worker_redis_connection_failed", error=str(e)) raise return _worker_redis_pool async def close_worker_redis_pool() -> None: """ 關閉 Worker 專屬 Redis 連線池 """ global _worker_redis_pool if _worker_redis_pool is not None: await _worker_redis_pool.close() _worker_redis_pool = None logger.info("worker_redis_pool_closed") def get_worker_redis() -> redis.Redis: """ 取得 Worker 專屬 Redis 連線 (長超時，用於 XREADGROUP) Raises: RuntimeError: 若連線池未初始化 """ if _worker_redis_pool is None: raise RuntimeError("Worker Redis pool not initialized. Call init_worker_redis_pool() first.") return _worker_redis_pool # ============================================================================= # Distributed Lock (分散式鎖) # ============================================================================= class RedisLock: """ Redis 分散式鎖 防禦場景: - 防止 Web + Telegram 同時簽核導致 Race Condition - 防止 K8s Executor 被觸發兩次 使用方式: async with RedisLock("approval:123:lock", timeout=10): # Critical section await execute_approval() """ def __init__( self, key: str, timeout: int = 30, blocking_timeout: float = 5.0, ): """ Args: key: 鎖的 Redis Key timeout: 鎖的自動過期時間 (秒) blocking_timeout: 等待取得鎖的最大時間 (秒) """ self.key = f"lock:{key}" self.timeout = timeout self.blocking_timeout = blocking_timeout self._lock_value: str | None = None async def acquire(self) -> bool: """ 嘗試取得鎖 Returns: bool: 是否成功取得鎖 """ import uuid redis_client = get_redis() self._lock_value = str(uuid.uuid4()) # 使用 SET NX EX 實現原子操作 acquired = await redis_client.set( self.key, self._lock_value, nx=True, # Only set if not exists ex=self.timeout, # Expire in timeout seconds ) if acquired: logger.debug("redis_lock_acquired", key=self.key) return True # 如果沒有立即取得，則等待 start_time = asyncio.get_event_loop().time() while asyncio.get_event_loop().time() - start_time < self.blocking_timeout: await asyncio.sleep(0.1) acquired = await redis_client.set( self.key, self._lock_value, nx=True, ex=self.timeout, ) if acquired: logger.debug("redis_lock_acquired_after_wait", key=self.key) return True logger.warning("redis_lock_timeout", key=self.key) return False async def release(self) -> bool: """ 釋放鎖 使用 Lua Script 確保只釋放自己持有的鎖 (防止誤刪) Returns: bool: 是否成功釋放 """ if self._lock_value is None: return False redis_client = get_redis() # Lua script: 只有當值匹配時才刪除 (原子操作) lua_script = """ if redis.call("get", KEYS[1]) == ARGV[1] then return redis.call("del", KEYS[1]) else return 0 end """ result = await redis_client.eval(lua_script, 1, self.key, self._lock_value) if result: logger.debug("redis_lock_released", key=self.key) return True else: logger.warning("redis_lock_release_failed", key=self.key) return False async def __aenter__(self) -> "RedisLock": acquired = await self.acquire() if not acquired: raise RuntimeError(f"Failed to acquire lock: {self.key}") return self async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: await self.release() # ============================================================================= # Context Manager # ============================================================================= @asynccontextmanager async def redis_context() -> AsyncGenerator[redis.Redis, None]: """ Redis 連線 Context Manager 用於需要獨立連線的場景 """ client = get_redis() try: yield client finally: pass # 使用連線池，不需要關閉