6f049877fc
- Python: ruff --fix 修復 280 個 lint 錯誤 - lewooogo-core: src/ 目錄未追蹤，導致 CI eslint 失敗 Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
303 lines
8.2 KiB
Python
303 lines
8.2 KiB
Python
"""
|
|
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 # 使用連線池，不需要關閉
|