#!/usr/bin/env python3
"""
sequential_vs_random.py
=======================

用 Python `os.write` + `os.lseek` 直接对裸文件做 4KB 块的「顺序写」vs「随机写」吞吐对比，
直观感受为什么 Kafka 的「顺序追加日志」是性能制胜法宝。

测量逻辑：
  - 顺序写：在一个空文件里从头到尾追加 N 个 4KB 块。
  - 随机写：在一个**预先分配好** N×4KB 大小的文件里，每次写一个随机偏移的 4KB 块。
  - 都用 `os.fsync` 强制刷盘，避免 PageCache 假象。

输出指标：
  - 总写入字节数
  - 耗时
  - 吞吐（MB/s）
  - 顺序 / 随机 倍数比

依赖：标准库（os / time / random / argparse），不需要任何第三方包。

用法：
  python3 sequential_vs_random.py
  python3 sequential_vs_random.py --blocks 50000 --block-size 4096
  python3 sequential_vs_random.py --no-fsync          # 看 PageCache 加持下的吞吐
  python3 sequential_vs_random.py --workdir /tmp      # 换数据目录（建议放 SSD/HDD/NVMe 各跑一次）
"""

from __future__ import annotations

import argparse
import os
import random
import sys
import time
import shutil
from pathlib import Path


def write_sequential(path: Path, blocks: int, block_size: int, fsync: bool) -> float:
    """
    在一个空文件中从头到尾追加 blocks 个 block_size 字节。
    返回耗时（秒）。
    """
    payload = b"\xAB" * block_size
    fd = os.open(str(path), os.O_CREAT | os.O_WRONLY | os.O_TRUNC, 0o644)
    t0 = time.perf_counter()
    try:
        for _ in range(blocks):
            os.write(fd, payload)
        if fsync:
            os.fsync(fd)
    finally:
        os.close(fd)
    return time.perf_counter() - t0


def write_random(path: Path, blocks: int, block_size: int, fsync: bool, seed: int) -> float:
    """
    在一个预分配 blocks*block_size 字节的文件里随机偏移写 blocks 个块。
    返回耗时（秒）。
    """
    total = blocks * block_size
    # 预分配
    with open(path, "wb") as f:
        f.truncate(total)

    rng = random.Random(seed)
    # 预生成所有偏移（不计入计时）
    offsets = [rng.randrange(0, blocks) * block_size for _ in range(blocks)]
    payload = b"\xCD" * block_size

    fd = os.open(str(path), os.O_WRONLY)
    t0 = time.perf_counter()
    try:
        for off in offsets:
            os.lseek(fd, off, os.SEEK_SET)
            os.write(fd, payload)
        if fsync:
            os.fsync(fd)
    finally:
        os.close(fd)
    return time.perf_counter() - t0


def human_mb(b: int) -> str:
    return f"{b / 1024 / 1024:.2f} MB"


def detect_fs_type(path: Path) -> str:
    """尽力探测挂载点 + 文件系统类型，方便读者解读。"""
    try:
        with open("/proc/self/mounts") as f:
            mounts = f.read().splitlines()
    except Exception:
        return "?"
    parent = path.resolve().parent
    while parent != Path("/") and parent.exists():
        for line in mounts:
            parts = line.split()
            if len(parts) >= 3 and parts[1] == str(parent):
                return f"{parts[0]} ({parts[2]})"
        parent = parent.parent
    return "?"


def main() -> None:
    parser = argparse.ArgumentParser(description="顺序写 vs 随机写吞吐对比")
    parser.add_argument("--workdir", type=str, default="./_seq_rand_bench",
                        help="测试数据目录（运行后会被删除）")
    parser.add_argument("--blocks", type=int, default=10000,
                        help="块数，默认 10000 个")
    parser.add_argument("--block-size", type=int, default=4096,
                        help="每块字节数，默认 4096")
    parser.add_argument("--no-fsync", action="store_true",
                        help="不调 fsync（看 PageCache 加持下的吞吐）")
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    workdir = Path(args.workdir)
    workdir.mkdir(parents=True, exist_ok=True)

    seq_file = workdir / "sequential.bin"
    rnd_file = workdir / "random.bin"
    total_bytes = args.blocks * args.block_size

    print("=" * 72)
    print(" 顺序写 vs 随机写 吞吐对比")
    print("-" * 72)
    print(f"  workdir   : {workdir.resolve()}")
    print(f"  fs        : {detect_fs_type(workdir)}")
    print(f"  blocks    : {args.blocks:,}")
    print(f"  block_size: {args.block_size} B")
    print(f"  total     : {human_mb(total_bytes)}")
    print(f"  fsync     : {'NO (依赖 PageCache)' if args.no_fsync else 'YES (强制刷盘)'}")
    print("=" * 72)

    print("\n[1/2] 顺序写…")
    seq_time = write_sequential(seq_file, args.blocks, args.block_size, fsync=not args.no_fsync)
    seq_mbps = total_bytes / seq_time / 1024 / 1024
    print(f"  耗时 {seq_time*1000:8.1f} ms   吞吐 {seq_mbps:7.1f} MB/s")

    print("\n[2/2] 随机写…")
    rnd_time = write_random(rnd_file, args.blocks, args.block_size,
                            fsync=not args.no_fsync, seed=args.seed)
    rnd_mbps = total_bytes / rnd_time / 1024 / 1024
    print(f"  耗时 {rnd_time*1000:8.1f} ms   吞吐 {rnd_mbps:7.1f} MB/s")

    ratio = seq_mbps / rnd_mbps if rnd_mbps > 0 else float("inf")

    print("\n" + "=" * 72)
    print(f"  顺序 / 随机 = {ratio:.2f}x")
    if ratio > 50:
        print("  ✅ 数量级差异（典型 HDD 行为）。Kafka 的顺序追加日志极大利用了这点。")
    elif ratio > 5:
        print("  🟡 倍数差异（典型 SSD 行为）。即使 SSD，顺序写也明显占优。")
    elif ratio > 1.5:
        print("  🟢 小幅领先（NVMe 或 PageCache 假象）。如果用了 --no-fsync，差异会被 OS 缓存掩盖。")
    else:
        print("  ⚠️  几乎无差异。可能 fsync 没生效，或 workdir 在 tmpfs 上。")
    print("=" * 72)

    # 清理
    try:
        shutil.rmtree(workdir)
        print(f"\n  已清理 {workdir}")
    except Exception as e:
        print(f"\n  ⚠️  清理 {workdir} 失败：{e}")

    print("\n实验提示：")
    print("  - 在 HDD 上跑预期看到 50~200x 顺序优势")
    print("  - 在 SATA SSD 上跑预期看到 5~15x")
    print("  - 在 NVMe SSD 上跑预期看到 3~6x")
    print("  - 加 --no-fsync 后，PageCache 会让随机写也变快，但生产环境必须 fsync")


if __name__ == "__main__":
    sys.exit(main())