# -*- coding: utf-8 -*- """基础设施层 · Redis 异步封装(缓存 / 限流 / 分布式锁)。 提供: - 基础 get/set/delete/incr - JSON 序列化 get_json/set_json - 缓存装饰器 @cached(自动缓存函数返回值) - 按前缀批量删除 delete_prefix - 分布式锁 lock/unlock - 全局单例 cache(应用启动时 connect,关闭时 close) """ from __future__ import annotations import asyncio import functools import json import logging from typing import Any, Callable, Optional from redis.asyncio import Redis from .. import config logger = logging.getLogger(__name__) class Cache: """redis.asyncio 封装:连接池 / get / set / incr / 分布式锁(SET NX EX)。""" def __init__(self, url: str | None = None) -> None: self._redis: Optional[Redis] = None self._url = url or config.REDIS_URL self._enabled = config.REDIS_ENABLED async def connect(self) -> None: if not self._enabled: logger.info("Redis 未配置,缓存降级为内存(无缓存)") return if self._redis is None: try: self._redis = Redis.from_url( self._url, decode_responses=True, max_connections=50, socket_timeout=5.0, socket_connect_timeout=5.0, ) # 测试连接 await self._redis.ping() logger.info(f"Redis 连接成功: {self._url}") except Exception as e: logger.warning(f"Redis 连接失败,缓存降级: {e}") self._redis = None async def close(self) -> None: if self._redis is not None: await self._redis.aclose() self._redis = None @property def available(self) -> bool: return self._redis is not None # ── 基础字符串操作 ────────────────────────────────────────────── async def get(self, key: str) -> str | None: if self._redis is None: return None try: return await self._redis.get(key) except Exception as e: logger.warning(f"Redis get 失败: {key}, {e}") return None async def set(self, key: str, value: str, ttl: int | None = None) -> None: if self._redis is None: return try: await self._redis.set(key, value, ex=ttl) except Exception as e: logger.warning(f"Redis set 失败: {key}, {e}") async def delete(self, key: str) -> None: if self._redis is None: return try: await self._redis.delete(key) except Exception as e: logger.warning(f"Redis delete 失败: {key}, {e}") async def incr(self, key: str) -> int: if self._redis is None: return 0 try: return int(await self._redis.incr(key)) except Exception as e: logger.warning(f"Redis incr 失败: {key}, {e}") return 0 # ── JSON 序列化操作 ────────────────────────────────────────────── async def get_json(self, key: str) -> Any | None: """获取并反序列化 JSON 值。""" raw = await self.get(key) if raw is None: return None try: return json.loads(raw) except Exception as e: logger.warning(f"Redis JSON 反序列化失败: {key}, {e}") return None async def set_json(self, key: str, value: Any, ttl: int | None = None) -> None: """序列化并存储 JSON 值。""" try: raw = json.dumps(value, ensure_ascii=False, default=str) except Exception as e: logger.warning(f"Redis JSON 序列化失败: {key}, {e}") return await self.set(key, raw, ttl=ttl) # ── 批量操作 ───────────────────────────────────────────────────── async def delete_prefix(self, prefix: str) -> int: """按前缀批量删除 key,返回删除数量。""" if self._redis is None: return 0 try: count = 0 async for key in self._redis.scan_iter(match=f"{prefix}*", count=100): await self._redis.delete(key) count += 1 return count except Exception as e: logger.warning(f"Redis delete_prefix 失败: {prefix}, {e}") return 0 # ── 分布式锁 ───────────────────────────────────────────────────── async def lock(self, key: str, ttl: int = 30) -> bool: """分布式锁:SET key token NX EX ttl(成功拿到锁返回 True)。""" if self._redis is None: return True # 无 Redis 时降级为总是通过(单机部署) token = f"{asyncio.get_event_loop().time()}" try: ok = await self._redis.set(key, token, nx=True, ex=ttl) return bool(ok) except Exception as e: logger.warning(f"Redis lock 失败: {key}, {e}") return True async def unlock(self, key: str) -> None: if self._redis is None: return try: await self._redis.delete(key) except Exception as e: logger.warning(f"Redis unlock 失败: {key}, {e}") # ── 全局单例 ────────────────────────────────────────────────────────── cache = Cache() # ── 缓存装饰器 ──────────────────────────────────────────────────────── def cached(key_prefix: str, ttl: int = 60, *, include_args: bool = True) -> Callable: """异步函数缓存装饰器。 用法: @cached("hall:tasks", ttl=10) async def get_tasks(mode="", category=""): ... key 格式:{key_prefix}:{args_hash}(include_args=True 时) 写操作后需手动调用 cache.delete_prefix(key_prefix) 失效缓存。 """ def decorator(func: Callable) -> Callable: @functools.wraps(func) async def wrapper(*args, **kwargs): if not cache.available: return await func(*args, **kwargs) # 构建缓存 key if include_args: import hashlib arg_str = f"{args}:{sorted(kwargs.items())}" arg_hash = hashlib.md5(arg_str.encode()).hexdigest()[:12] cache_key = f"{key_prefix}:{arg_hash}" else: cache_key = key_prefix # 尝试读缓存 cached_val = await cache.get_json(cache_key) if cached_val is not None: return cached_val # 缓存未命中,执行函数 result = await func(*args, **kwargs) # 写缓存(异步,不阻塞返回) asyncio.create_task(cache.set_json(cache_key, result, ttl=ttl)) return result return wrapper return decorator