redis_cache_health_check.py 12.4 KB
#!/usr/bin/env python3
"""
缓存状态巡检脚本

按「缓存类型」维度(embedding / translation / anchors)查看:
- 估算 key 数量
- TTL 分布(采样)
- 近期活跃 key(按 IDLETIME 近似)
- 若干样本 key 及 value 概览

使用示例:

  # 默认:检查已知三类缓存,使用 env_config 中的 Redis 配置
  python scripts/redis/redis_cache_health_check.py

  # 只看某一类缓存
  python scripts/redis/redis_cache_health_check.py --type embedding
  python scripts/redis/redis_cache_health_check.py --type translation anchors

  # 自定义前缀(pattern),不限定缓存类型
  python scripts/redis/redis_cache_health_check.py --pattern "mycache:*"

  # 调整采样规模
  python scripts/redis/redis_cache_health_check.py --sample-size 100 --max-scan 50000
"""

from __future__ import annotations

import argparse
import json
import pickle
import sys
from collections import defaultdict
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Tuple

import redis
import numpy as np

# 让脚本可以直接使用 config/env_config 与 services_config
PROJECT_ROOT = Path(__file__).parent.parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

from config.env_config import REDIS_CONFIG  # type: ignore
from config.services_config import get_translation_cache_config  # type: ignore


@dataclass
class CacheTypeConfig:
    name: str
    pattern: str
    description: str


def _load_known_cache_types() -> Dict[str, CacheTypeConfig]:
    """根据当前配置装配三种已知缓存类型及其前缀 pattern。"""
    cache_types: Dict[str, CacheTypeConfig] = {}

    # embedding 缓存:固定 embedding:* 前缀
    cache_types["embedding"] = CacheTypeConfig(
        name="embedding",
        pattern="embedding:*",
        description="文本向量缓存(embeddings/text_encoder.py)",
    )

    # translation 缓存:prefix 来自 services.translation.cache.key_prefix
    cache_cfg = get_translation_cache_config()
    trans_prefix = cache_cfg.get("key_prefix", "trans:v2")
    cache_types["translation"] = CacheTypeConfig(
        name="translation",
        pattern=f"{trans_prefix}:*",
        description="翻译结果缓存(query/qwen_mt_translate.Translator)",
    )

    # anchors 缓存:prefix 来自 REDIS_CONFIG['anchor_cache_prefix'](若存在),否则 product_anchors
    anchor_prefix = REDIS_CONFIG.get("anchor_cache_prefix", "product_anchors")
    cache_types["anchors"] = CacheTypeConfig(
        name="anchors",
        pattern=f"{anchor_prefix}:*",
        description="商品内容理解缓存(indexer/product_enrich.py,anchors/语义属性/tags)",
    )

    return cache_types


def get_redis_client(db: int = 0) -> redis.Redis:
    return redis.Redis(
        host=REDIS_CONFIG.get("host", "localhost"),
        port=REDIS_CONFIG.get("port", 6479),
        password=REDIS_CONFIG.get("password"),
        db=db,
        decode_responses=False,  # 原始 bytes,方便区分 pickle / str
        socket_timeout=10,
        socket_connect_timeout=10,
    )


def scan_keys(
    client: redis.Redis, pattern: str, max_scan: int, scan_count: int = 1000
) -> List[bytes]:
    """使用 SCAN 扫描匹配 pattern 的 key,最多扫描 max_scan 个结果。"""
    keys: List[bytes] = []
    cursor: int = 0
    scanned = 0
    while True:
        cursor, batch = client.scan(cursor=cursor, match=pattern, count=scan_count)
        for k in batch:
            keys.append(k)
            scanned += 1
            if scanned >= max_scan:
                return keys
        if cursor == 0:
            break
    return keys


def ttl_bucket(ttl: int) -> str:
    """将 TTL(秒)归类到简短区间标签中。"""
    if ttl < 0:
        # -1: 永不过期;-2: 不存在
        return "no-expire-or-expired"
    if ttl <= 3600:
        return "0-1h"
    if ttl <= 86400:
        return "1h-1d"
    if ttl <= 30 * 86400:
        return "1d-30d"
    return ">30d"


def format_seconds(sec: int) -> str:
    if sec < 0:
        return str(sec)
    if sec < 60:
        return f"{sec}s"
    if sec < 3600:
        return f"{sec // 60}m{sec % 60}s"
    if sec < 86400:
        h = sec // 3600
        m = (sec % 3600) // 60
        return f"{h}h{m}m"
    d = sec // 86400
    h = (sec % 86400) // 3600
    return f"{d}d{h}h"


def decode_value_preview(
    cache_type: str, key: bytes, raw_value: Optional[bytes]
) -> str:
    """根据缓存类型生成简短的 value 概览字符串。"""
    if raw_value is None:
        return "<nil>"

    # embedding: pickle 序列化的 numpy.ndarray
    if cache_type == "embedding":
        try:
            arr = pickle.loads(raw_value)
            if isinstance(arr, np.ndarray):
                return f"ndarray shape={arr.shape} dtype={arr.dtype}"
            return f"pickle object type={type(arr).__name__}"
        except Exception:
            return f"<binary {len(raw_value)} bytes>"

    # anchors: JSON dict
    if cache_type == "anchors":
        try:
            text = raw_value.decode("utf-8", errors="replace")
            obj = json.loads(text)
            if isinstance(obj, dict):
                brief = {
                    k: obj.get(k)
                    for k in ["id", "lang", "title_input", "title", "category_path", "anchor_text"]
                    if k in obj
                }
                return "json " + json.dumps(brief, ensure_ascii=False)[:200]
            # 其他情况简单截断
            return "json " + text[:200]
        except Exception:
            return raw_value.decode("utf-8", errors="replace")[:200]

    # translation: 纯字符串
    if cache_type == "translation":
        try:
            text = raw_value.decode("utf-8", errors="replace")
            return text[:200]
        except Exception:
            return f"<binary {len(raw_value)} bytes>"

    # 兜底:尝试解码为 UTF-8
    try:
        text = raw_value.decode("utf-8", errors="replace")
        return text[:200]
    except Exception:
        return f"<binary {len(raw_value)} bytes>"


def analyze_cache_type(
    client: redis.Redis,
    cache_type: str,
    cfg: CacheTypeConfig,
    sample_size: int,
    max_scan: int,
) -> None:
    """对单个缓存类型做统计与样本展示。"""
    print("=" * 80)
    print(f"Cache type: {cache_type}  pattern={cfg.pattern}")
    print(f"Description: {cfg.description}")
    print("=" * 80)

    keys = scan_keys(client, cfg.pattern, max_scan=max_scan)
    total_scanned = len(keys)
    print(f"Scanned up to {max_scan} keys, matched {total_scanned} keys.")
    if total_scanned == 0:
        print("No keys found for this cache type.\n")
        return

    # 采样
    if total_scanned <= sample_size:
        sampled_keys = keys
    else:
        # 简单取前 sample_size 个,scan 本身已是渐进遍历,足够近似随机
        sampled_keys = keys[:sample_size]

    ttl_hist: Dict[str, int] = defaultdict(int)
    recent_keys: List[Tuple[bytes, int, int]] = []  # (key, ttl, idletime)
    samples: List[Tuple[bytes, int, int, str]] = []  # (key, ttl, idletime, preview)

    for k in sampled_keys:
        try:
            ttl = client.ttl(k)
        except Exception:
            ttl = -3  # 表示 TTL 查询失败

        # TTL 分布
        ttl_hist[ttl_bucket(ttl)] += 1

        # 近期活跃判断(idletime 越小越“新”)
        idletime = -1
        try:
            # OBJECT IDLETIME 返回秒数(整数)
            idletime = client.object("idletime", k)  # type: ignore[arg-type]
        except Exception:
            pass

        # 记录近期活跃样本
        if idletime >= 0 and idletime <= 600:
            recent_keys.append((k, ttl, idletime))

        # 收集样本 value 预览(控制数量)
        if len(samples) < 5:
            raw_val = None
            try:
                raw_val = client.get(k)
            except Exception:
                pass
            preview = decode_value_preview(cache_type, k, raw_val)
            samples.append((k, ttl, idletime, preview))

    # TTL 分布输出
    print("\nTTL distribution (sampled):")
    total_sampled = len(sampled_keys)
    for bucket in ["no-expire-or-expired", "0-1h", "1h-1d", "1d-30d", ">30d"]:
        cnt = ttl_hist.get(bucket, 0)
        pct = (cnt / total_sampled * 100.0) if total_sampled else 0.0
        print(f"  {bucket:<18}: {cnt:>4} ({pct:>5.1f}%)")

    # 近期活跃 key
    recent_keys = sorted(recent_keys, key=lambda x: x[2])[:5]
    print("\nRecent active keys (idletime <= 600s, from sampled set):")
    if not recent_keys:
        print("  (none in sampled set)")
    else:
        for k, ttl, idle in recent_keys:
            try:
                k_str = k.decode("utf-8", errors="replace")
            except Exception:
                k_str = repr(k)
            if len(k_str) > 80:
                k_str = k_str[:77] + "..."
            print(
                f"  key={k_str}  ttl={ttl} ({format_seconds(ttl)})  "
                f"idletime={idle} ({format_seconds(idle)})"
            )

    # 样本 value 概览
    print("\nSample keys & value preview:")
    if not samples:
        print("  (no samples)")
    else:
        for k, ttl, idle, preview in samples:
            try:
                k_str = k.decode("utf-8", errors="replace")
            except Exception:
                k_str = repr(k)
            if len(k_str) > 80:
                k_str = k_str[:77] + "..."
            print(f"  key={k_str}")
            print(f"    ttl={ttl} ({format_seconds(ttl)})  idletime={idle} ({format_seconds(idle)})")
            print(f"    value: {preview}")

    print()  # 结尾空行


def main() -> None:
    parser = argparse.ArgumentParser(description="Redis 缓存状态巡检(按缓存类型)")
    parser.add_argument(
        "--type",
        dest="types",
        nargs="+",
        choices=["embedding", "translation", "anchors"],
        help="指定要检查的缓存类型(默认:三种全部)",
    )
    parser.add_argument(
        "--pattern",
        type=str,
        help="自定义 key pattern(如 'mycache:*')。设置后将忽略 --type,仅按 pattern 进行一次检查。",
    )
    parser.add_argument(
        "--db",
        type=int,
        default=0,
        help="Redis 数据库编号(默认 0)",
    )
    parser.add_argument(
        "--sample-size",
        type=int,
        default=50,
        help="每种缓存类型采样的 key 数量(默认 50)",
    )
    parser.add_argument(
        "--max-scan",
        type=int,
        default=20000,
        help="每种缓存类型最多 SCAN 的 key 数量上限(默认 20000)",
    )

    args = parser.parse_args()

    print("Redis 缓存状态巡检")
    print("=" * 80)
    print(f"Checked at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"Redis host={REDIS_CONFIG.get('host', 'localhost')} port={REDIS_CONFIG.get('port', 6479)} db={args.db}")
    print()

    try:
        client = get_redis_client(db=args.db)
        client.ping()
        print("✅ Redis 连接成功\n")
    except Exception as exc:
        print(f"❌ Redis 连接失败: {exc}")
        print(f"  Host: {REDIS_CONFIG.get('host', 'localhost')}")
        print(f"  Port: {REDIS_CONFIG.get('port', 6479)}")
        print(f"  Password: {'已配置' if REDIS_CONFIG.get('password') else '未配置'}")
        return

    # 如果指定了自定义 pattern,则只做一次“匿名类型”的巡检
    if args.pattern:
        anon_cfg = CacheTypeConfig(
            name=f"pattern:{args.pattern}",
            pattern=args.pattern,
            description="自定义 pattern 检查",
        )
        analyze_cache_type(
            client=client,
            cache_type="custom",
            cfg=anon_cfg,
            sample_size=args.sample_size,
            max_scan=args.max_scan,
        )
        print("巡检完成。")
        return

    # 否则根据已知缓存类型巡检
    known_types = _load_known_cache_types()
    types_to_check: Iterable[str]
    if args.types:
        types_to_check = args.types
    else:
        types_to_check = known_types.keys()

    for t in types_to_check:
        cfg = known_types.get(t)
        if not cfg:
            print(f"⚠️ 未知缓存类型: {t},跳过")
            continue
        analyze_cache_type(
            client=client,
            cache_type=t,
            cfg=cfg,
            sample_size=args.sample_size,
            max_scan=args.max_scan,
        )

    print("巡检完成。")


if __name__ == "__main__":
    main()