"""
Ch4 配套代码 1 / 3 —— SCAN vs KEYS 性能对比

  - 灌入 5 万个 Key
  - 分别用 KEYS 和 SCAN 遍历，对比单次阻塞时间
  - SCAN 总耗时可能更长，但单次极短，不会阻塞主线程
"""

import time
import redis

r = redis.Redis(host="127.0.0.1", port=6379, decode_responses=True)
PREFIX = "demo:scan:"


def section(t): print("\n" + "=" * 60 + f"\n{t}\n" + "=" * 60)


def setup(n: int) -> None:
    print(f"[setup] 灌入 {n} 个 Key ...")
    pipe = r.pipeline(transaction=False)
    for i in range(n):
        pipe.set(f"{PREFIX}{i}", "v")
        if i % 1000 == 999: pipe.execute()
    pipe.execute()


def cleanup() -> None:
    cursor = 0
    while True:
        cursor, keys = r.scan(cursor, match=f"{PREFIX}*", count=1000)
        if keys: r.delete(*keys)
        if cursor == 0: break


def use_keys() -> None:
    section("方案 A: KEYS（单次阻塞）⚠️ 生产禁用")
    start = time.perf_counter()
    keys = r.keys(f"{PREFIX}*")
    cost = (time.perf_counter() - start) * 1000
    print(f"  返回 {len(keys)} 个 Key")
    print(f"  ⚠️ 单次调用耗时 {cost:.2f} ms（这段时间主线程阻塞，所有客户端等待）")


def use_scan() -> None:
    section("方案 B: SCAN（多次小调用）✅ 安全")
    cursor = 0
    rounds = 0
    max_call_ms = 0
    total_keys = 0
    start = time.perf_counter()
    while True:
        s = time.perf_counter()
        cursor, keys = r.scan(cursor, match=f"{PREFIX}*", count=500)
        single = (time.perf_counter() - s) * 1000
        max_call_ms = max(max_call_ms, single)
        total_keys += len(keys)
        rounds += 1
        if cursor == 0: break
    total = (time.perf_counter() - start) * 1000
    print(f"  调用 {rounds} 次，累计返回 {total_keys} 个 Key（含可能的重复）")
    print(f"  总耗时 {total:.2f} ms（分摊到多次，主线程不阻塞）")
    print(f"  ✓ 单次最大耗时 {max_call_ms:.2f} ms（安全）")


if __name__ == "__main__":
    try:
        cleanup()
        setup(50000)
        use_keys()
        use_scan()
    finally:
        cleanup()