"""
Ch7 配套代码 2 / 4 —— Pipeline / 单条 / 事务 三种方式压测

演示：
  1. 1 万次 SET，单条调用 —— 每次都走一次 RTT
  2. 1 万次 SET，Pipeline (batch=100) —— 100 倍少 RTT
  3. 1 万次 SET，MULTI/EXEC 整体事务 —— 也省 RTT 但服务端要排队
"""

import time
import redis

HOST, PORT = "127.0.0.1", 6379
N = 10_000          # 总命令数
BATCH = 100         # Pipeline / Tx 批大小


def make_client() -> redis.Redis:
    return redis.Redis(host=HOST, port=PORT, decode_responses=True)


def section(title: str) -> None:
    print("\n" + "=" * 62)
    print(title)
    print("=" * 62)


def warmup() -> None:
    r = make_client()
    r.flushdb()
    for _ in range(50): r.ping()


def bench_single() -> float:
    r = make_client()
    r.flushdb()
    t0 = time.perf_counter()
    for i in range(N):
        r.set(f"bench:s:{i}", "v")
    return time.perf_counter() - t0


def bench_pipeline(transaction: bool = False) -> float:
    r = make_client()
    r.flushdb()
    t0 = time.perf_counter()
    pipe = r.pipeline(transaction=transaction)
    for i in range(N):
        pipe.set(f"bench:p:{i}", "v")
        if (i + 1) % BATCH == 0:
            pipe.execute()
            pipe = r.pipeline(transaction=transaction)
    pipe.execute()
    return time.perf_counter() - t0


def bench_full_transaction() -> float:
    """整个 1 万条全塞一个 MULTI/EXEC 里"""
    r = make_client()
    r.flushdb()
    t0 = time.perf_counter()
    pipe = r.pipeline(transaction=True)
    for i in range(N):
        pipe.set(f"bench:t:{i}", "v")
    pipe.execute()
    return time.perf_counter() - t0


def fmt_row(name: str, sec: float) -> str:
    qps = N / sec if sec > 0 else float("inf")
    return f"  {name:<32} {sec*1000:8.1f} ms     {qps:>10,.0f} qps"


if __name__ == "__main__":
    try:
        warmup()
        section(f"Pipeline 性能压测  N={N}  batch={BATCH}")

        results = []
        results.append(("① 单条 SET（每次一次 RTT）",     bench_single()))
        results.append(("② Pipeline 非事务",              bench_pipeline(False)))
        results.append(("③ Pipeline + 事务（默认）",       bench_pipeline(True)))
        results.append(("④ 单个超大事务（1 万条全打包）",   bench_full_transaction()))

        print()
        print(f"  {'方式':<32} {'耗时':>10}     {'吞吐':>14}")
        print("  " + "-" * 60)
        baseline = results[0][1]
        for name, sec in results:
            print(fmt_row(name, sec))

        print("\n  📊 加速比（相对于「单条」）：")
        for name, sec in results:
            speed = baseline / sec if sec > 0 else float("inf")
            print(f"    {name:<32} ×{speed:>5.1f}")

        print("\n  💡 结论：")
        print("    - Pipeline 主要省 RTT，本机 50ms RTT 都能让 1 万次提速 50x")
        print("    - 大事务（④）也省 RTT 但服务端要单线程排队执行 + 占内存")
        print("    - 生产推荐 Pipeline + 中等批大小（100~1000）")

        make_client().flushdb()
    except redis.ConnectionError as e:
        print(f"❌ Redis 连接失败: {e}")