server-core/app/infrastructure/cache.py

# -*- 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