mreschke/speedtest-hd.py

speedtest-hd.py · 25 KiB · Python Raw

#!/usr/bin/env python3
"""speedtest-hd — a CrystalDiskMark-style storage benchmark for Linux, on fio.

This is the Python successor to ``speedtest-hd.sh``. It measures storage the way
CrystalDiskMark does, and adds a dedicated **SLOG / sync-write latency profile**
for diagnosing ZFS ZIL performance (NFS / iSCSI / VM sync workloads).

Three profiles
--------------
* **cdm**  — the four CrystalDiskMark default tests (``SEQ1M Q8T1``, ``SEQ1M
  Q1T1``, ``RND4K Q32T16``, ``RND4K Q1T1``), each measured for read *and*
  write, reported in MB/s and IOPS.
* **slog** — synchronous 4K random writes swept across thread counts, reporting
  IOPS, MB/s and p50/p99 *commit latency*. This is the load a ZFS SLOG sees.
* **dd**   — a dependency-free fallback when ``fio`` isn't installed.

Why fio JSON?
-------------
Every measurement asks fio for ``--output-format=json`` and parses it with the
standard library. That's the whole reason this is Python and not bash: robust,
unit-safe parsing of bandwidth / IOPS / latency percentiles with no fragile
text scraping.

See README.md for the full case study (diagnosing a "slow" Optane SLOG that
turned out to be CPU power management, not the disk).
"""

from __future__ import annotations

import argparse
import glob
import json
import os
import shutil
import statistics
import subprocess
import sys
import time
from dataclasses import dataclass, field
from typing import Callable, Optional, Sequence

# --------------------------------------------------------------------------- #
# Color / styling (standard library only -- raw ANSI escape codes)
# --------------------------------------------------------------------------- #

# SGR escape sequences. We stick to the basic 16-color set so output respects
# the user's terminal theme instead of hard-coding RGB.
RESET = "\033[0m"
BOLD = "\033[1m"
DIM = "\033[2m"
ITALIC = "\033[3m"
RED = "\033[31m"
GREEN = "\033[32m"
YELLOW = "\033[33m"
BLUE = "\033[34m"
MAGENTA = "\033[35m"
CYAN = "\033[36m"
WHITE = "\033[37m"
BRIGHT_CYAN = "\033[96m"

def supports_color(stream) -> bool:
    """Decide per-stream whether to emit ANSI codes.

Honors the de-facto ``NO_COLOR`` / ``FORCE_COLOR`` conventions and otherwise
    colors only when the stream is an interactive terminal.
    """
    if os.environ.get("NO_COLOR"):
        return False
    if os.environ.get("FORCE_COLOR"):
        return True
    if os.environ.get("TERM") == "dumb":
        return False
    return bool(getattr(stream, "isatty", None)) and stream.isatty()

class Painter:
    """Wraps text in SGR codes, or returns it untouched when color is off."""

def __init__(self, enabled: bool):
        self.enabled = enabled

def paint(self, text: str, *codes: str) -> str:
        if not self.enabled or not codes:
            return text
        return "".join(codes) + text + RESET

# stdout carries results (banner + tables); stderr carries progress + verbose
# fio dumps. Each gets its own enable flag so piping one but not the other works.
OUT = Painter(supports_color(sys.stdout))
ERR = Painter(supports_color(sys.stderr))

# --------------------------------------------------------------------------- #
# Constants
# --------------------------------------------------------------------------- #

# Preference order for the IO engine. io_uring is the modern, lowest-overhead
# Linux engine; libaio is older (and only truly async with O_DIRECT); posixaio
# and sync are the portable fallbacks (e.g. some NFS mounts).
ENGINE_CANDIDATES = ("io_uring", "libaio", "posixaio", "sync")

# Binary unit multipliers (fio treats "1g" as 1 GiB, so we match that).
_SIZE_UNITS = {"k": 1024, "m": 1024**2, "g": 1024**3, "t": 1024**4}

@dataclass(frozen=True)
class CdmTest:
    """One CrystalDiskMark test. ``seq`` picks sequential vs random patterns."""

label: str
    bs: str
    iodepth: int
    numjobs: int
    seq: bool

@property
    def read_rw(self) -> str:
        return "read" if self.seq else "randread"

@property
    def write_rw(self) -> str:
        return "write" if self.seq else "randwrite"

# The four CrystalDiskMark default tests, as data, in CrystalDiskMark's own
# display order. Q = queue depth (iodepth), T = threads (numjobs).
CDM_TESTS: tuple[CdmTest, ...] = (
    CdmTest("SEQ1M  Q8T1",   bs="1m", iodepth=8,  numjobs=1,  seq=True),
    CdmTest("SEQ1M  Q1T1",   bs="1m", iodepth=1,  numjobs=1,  seq=True),
    CdmTest("RND4K  Q32T16", bs="4k", iodepth=32, numjobs=16, seq=False),
    CdmTest("RND4K  Q1T1",   bs="4k", iodepth=1,  numjobs=1,  seq=False),
)

# Thread counts for the SLOG sweep. T1 is the headline single-stream latency;
# the higher counts show how the SLOG scales as concurrent sync writers pile on.
SLOG_THREADS: tuple[int, ...] = (1, 4, 8, 16)

# --------------------------------------------------------------------------- #
# Configuration
# --------------------------------------------------------------------------- #

@dataclass
class Config:
    """Resolved run settings. ``engine`` and ``direct`` may start unset and get
    filled in by :func:`detect_io_settings`."""

path: str
    mode: Optional[str]            # "cdm" | "slog" | "dd" | None (auto)
    engine: Optional[str]          # None => auto-detect
    direct: Optional[bool]         # None => auto-detect
    runtime: int
    size: str
    verbose: bool
    assume_yes: bool
    benchfile: str = field(default="")

@property
    def direct_flag(self) -> int:
        """fio's --direct takes 0/1; default to 0 until detection runs."""
        return 1 if self.direct else 0

# --------------------------------------------------------------------------- #
# Small helpers
# --------------------------------------------------------------------------- #

def log(message: str) -> None:
    """Progress/diagnostic line -> stderr, so stdout stays pure results."""
    print(message, file=sys.stderr, flush=True)

def step(message: str) -> None:
    """A single colored progress line ('measuring ...') -> stderr."""
    log(f"  {ERR.paint('▶', BOLD, CYAN)} {ERR.paint(message, DIM)}")

def vsection(title: str) -> None:
    """A clear, ruled header for one --verbose fio section -> stderr."""
    width = 78
    head = f"─── {title} "
    head += "─" * max(3, width - len(head))
    log("")
    log(ERR.paint(head, BOLD, CYAN))

def parse_size_bytes(size: str) -> int:
    """Turn an fio-style size ("4k", "1g", "512m", "2048") into bytes."""
    s = size.strip().lower()
    if s and s[-1] in _SIZE_UNITS:
        return int(float(s[:-1]) * _SIZE_UNITS[s[-1]])
    return int(s)

def render_table(
    headers: Sequence[str],
    rows: Sequence[Sequence[str]],
    aligns: Optional[Sequence[str]] = None,
    col_styles: Optional[Sequence[Optional[str]]] = None,
) -> str:
    """Render a bordered ASCII table whose columns auto-size to their content.

``aligns`` is a per-column "<" (left) or ">" (right); by default the first
    column is left-aligned (the label) and the rest are right-aligned (numbers).
    ``col_styles`` is an optional per-column ANSI style applied to body cells
    (header is always bold, borders dim). Widths are computed on the *plain*
    text and color is applied after padding, so escape codes never skew layout.
    """
    ncols = len(headers)
    if aligns is None:
        aligns = ["<"] + [">"] * (ncols - 1)
    if col_styles is None:
        col_styles = [None] * ncols

# Widest cell (header or any row) sets each column's width.
    widths = [
        max([len(str(headers[c]))] + [len(str(r[c])) for r in rows])
        for c in range(ncols)
    ]

bar = OUT.paint("|", DIM)
    sep = OUT.paint(" | ", DIM)

def rule() -> str:
        return OUT.paint("+" + "+".join("-" * (w + 2) for w in widths) + "+", DIM)

def line(cells: Sequence[str], *, header: bool = False) -> str:
        parts = []
        for c in range(ncols):
            padded = f"{str(cells[c]):{aligns[c]}{widths[c]}}"
            if header:
                parts.append(OUT.paint(padded, BOLD))
            elif col_styles[c]:
                parts.append(OUT.paint(padded, col_styles[c]))
            else:
                parts.append(padded)
        return f"{bar} " + sep.join(parts) + f" {bar}"

out = [rule(), line(headers, header=True), rule()]
    out += [line(r) for r in rows]
    out.append(rule())
    return "\n".join(out)

# --------------------------------------------------------------------------- #
# fio: results, invocation, parsing
# --------------------------------------------------------------------------- #

@dataclass(frozen=True)
class FioResult:
    """Parsed metrics for one direction (read or write) of an fio run.

Latencies are completion latency (clat) in microseconds; for synchronous
    writes that is effectively the durable-commit latency.
    """

bw_mbps: float        # decimal MB/s, matching CrystalDiskMark's SI figure
    iops: float
    p50_us: float
    p99_us: float
    mean_us: float

@classmethod
    def zero(cls) -> "FioResult":
        return cls(0.0, 0.0, 0.0, 0.0, 0.0)

def _direction_of(rw: str) -> str:
    """fio reports 'read' and 'write' sub-objects; pick the active one."""
    return "read" if "read" in rw else "write"

def _aggregate(jobs: list[dict], direction: str) -> FioResult:
    """Combine per-job fio JSON stats into a single :class:`FioResult`.

We deliberately run fio *without* --group_reporting and aggregate ourselves,
    which sidesteps fio's group-merge quirks: sum throughput, average the median
    latency, and take the worst-case tail (p99) across jobs.
    """

def section(job: dict) -> dict:
        return job.get(direction, {})

bw_mbps = sum(section(j).get("bw_bytes", 0) for j in jobs) / 1e6
    iops = sum(section(j).get("iops", 0.0) for j in jobs)

def percentiles(key: str) -> list[float]:
        vals = []
        for j in jobs:
            pct = section(j).get("clat_ns", {}).get("percentile", {})
            v = pct.get(key)
            if v is not None:
                vals.append(v)
        return vals

p50_ns = percentiles("50.000000")
    p99_ns = percentiles("99.000000")
    mean_ns = [section(j).get("clat_ns", {}).get("mean", 0.0) for j in jobs]

return FioResult(
        bw_mbps=bw_mbps,
        iops=iops,
        p50_us=(statistics.mean(p50_ns) / 1000.0) if p50_ns else 0.0,   # avg of medians
        p99_us=(max(p99_ns) / 1000.0) if p99_ns else 0.0,               # worst tail
        mean_us=(statistics.mean(mean_ns) / 1000.0) if mean_ns else 0.0,
    )

def run_fio(
    cfg: Config,
    *,
    rw: str,
    bs: str,
    iodepth: int,
    numjobs: int,
    engine: Optional[str] = None,
    direct: Optional[bool] = None,
    sync: bool = False,
    end_fsync: bool = False,
) -> FioResult:
    """Run a single fio job against the shared bench file and parse its JSON.

``engine``/``direct`` default to the detected config values; the SLOG path
    overrides them (psync + buffered + O_SYNC). ``end_fsync`` flushes the device
    cache at the end of a write run so cached writes can't inflate the result.
    """
    engine = engine if engine is not None else cfg.engine
    direct = direct if direct is not None else cfg.direct

cmd = [
        "sudo", "fio",
        "--name=speedtest",
        f"--filename={cfg.benchfile}",
        f"--ioengine={engine}",
        f"--direct={1 if direct else 0}",
        f"--rw={rw}",
        f"--bs={bs}",
        f"--size={cfg.size}",
        f"--numjobs={numjobs}",
        f"--iodepth={iodepth}",
        "--time_based",
        f"--runtime={cfg.runtime}",
        "--output-format=json",
    ]
    if sync:
        cmd.append("--sync=1")        # O_SYNC: every write is a durable commit
    if end_fsync:
        cmd.append("--end_fsync=1")   # flush device cache once at the end

proc = subprocess.run(cmd, capture_output=True, text=True)

if cfg.verbose:
        sync_tag = "  sync=1" if sync else ""
        vsection(f"fio · {rw}  bs={bs}  iodepth={iodepth}  numjobs={numjobs}{sync_tag}")
        log("  " + ERR.paint("$ " + " ".join(cmd[1:]), DIM, ITALIC))
        log(ERR.paint(proc.stdout.strip() or "(no stdout)", DIM))
        if proc.stderr.strip():
            log(ERR.paint(proc.stderr.strip(), YELLOW))

try:
        data = json.loads(proc.stdout)
        return _aggregate(data["jobs"], _direction_of(rw))
    except (json.JSONDecodeError, KeyError, ValueError):
        # fio failed or produced no parseable JSON; report zeros rather than
        # crashing the whole run (verbose mode above shows what went wrong).
        return FioResult.zero()

def fio_probe(path: str, engine: str, direct: bool) -> bool:
    """Throwaway 1s fio job to see if an (engine, O_DIRECT) combo works here.

This is how we stay accurate *and* portable across ext4/xfs/btrfs/ZFS/NFS
    without the caller needing to know what each filesystem supports.
    """
    cmd = [
        "sudo", "fio", "--name=probe",
        f"--directory={path}",
        f"--ioengine={engine}",
        "--rw=write", "--bs=4k", "--size=1m",
        f"--direct={1 if direct else 0}",
        "--time_based", "--runtime=1",
    ]
    proc = subprocess.run(cmd, capture_output=True, text=True)
    for leftover in glob.glob(os.path.join(path, "probe*")):
        try:
            os.remove(leftover)
        except OSError:
            pass
    return proc.returncode == 0

def detect_io_settings(cfg: Config) -> None:
    """Fill in ``cfg.engine`` and ``cfg.direct`` if the user didn't force them.

O_DIRECT bypasses the OS page cache so we measure the device, not RAM. Not
    every filesystem supports it (older OpenZFS, some NFS), so we probe and fall
    back to buffered rather than erroring out.
    """
    if cfg.engine is None:
        for engine in ENGINE_CANDIDATES:
            if fio_probe(cfg.path, engine, direct=False):
                cfg.engine = engine
                break
        else:
            cfg.engine = "sync"

if cfg.direct is None:
        cfg.direct = fio_probe(cfg.path, cfg.engine, direct=True)

# --------------------------------------------------------------------------- #
# Profiles
# --------------------------------------------------------------------------- #

def _meta_line(text: str) -> str:
    """Color the 'Label :' prefix of a banner metadata line, leave value as-is.

The value may already contain its own ANSI codes (e.g. a red O_DIRECT
    warning); since these lines aren't width-aligned that's harmless.
    """
    key, sep, val = text.partition(":")
    if not sep:
        return "  " + text
    return "  " + OUT.paint(key + ":", BOLD, BRIGHT_CYAN) + val

def banner(title: str, cfg: Config, extra: Sequence[str] = ()) -> None:
    full = f"speedtest-hd  ·  {title}"
    inner = len(full) + 4

print()
    print(OUT.paint("╭" + "─" * inner + "╮", BOLD, CYAN))
    print(
        OUT.paint("│", BOLD, CYAN)
        + "  " + OUT.paint(full, BOLD, WHITE) + "  "
        + OUT.paint("│", BOLD, CYAN)
    )
    print(OUT.paint("╰" + "─" * inner + "╯", BOLD, CYAN))
    print(_meta_line(f"Target  : {OUT.paint(cfg.path, CYAN)}"))
    for line_ in extra:
        print(_meta_line(line_))
    print()

def cdm_profile(cfg: Config) -> None:
    """The CrystalDiskMark-style profile: 4 tests x (read, write), MB/s + IOPS."""
    detect_io_settings(cfg)

if cfg.direct:
        direct_label = OUT.paint("enabled (device)", GREEN)
    else:
        direct_label = OUT.paint(
            "DISABLED (buffered -- may reflect RAM cache!)", BOLD, RED
        )

banner(
        "CrystalDiskMark-style storage benchmark",
        cfg,
        extra=[
            f"Engine  : {cfg.engine}    O_DIRECT: {direct_label}",
            f"Profile : size={cfg.size.upper()}  runtime={cfg.runtime}s/run  (8 runs)",
        ],
    )

rows: list[list[str]] = []
    for test in CDM_TESTS:
        step(f"measuring {test.label.strip()} ...")
        r = run_fio(cfg, rw=test.read_rw, bs=test.bs,
                    iodepth=test.iodepth, numjobs=test.numjobs)
        w = run_fio(cfg, rw=test.write_rw, bs=test.bs,
                    iodepth=test.iodepth, numjobs=test.numjobs, end_fsync=True)
        rows.append([
            test.label,
            f"{r.bw_mbps:.2f}", f"{w.bw_mbps:.2f}",
            f"{r.iops:.0f}", f"{w.iops:.0f}",
        ])

_cleanup(cfg)

print()
    print(render_table(
        ["Test", "Read (MB/s)", "Write (MB/s)", "Read (IOPS)", "Write (IOPS)"],
        rows,
        col_styles=[CYAN, GREEN, YELLOW, GREEN, YELLOW],
    ))
    print()

def slog_profile(cfg: Config) -> None:
    """SLOG / sync-write latency profile.

Forces synchronous 4K random writes (O_SYNC) via the portable psync engine,
    so every write is a ZIL commit -- exercising a ZFS SLOG exactly the way a
    sync=always dataset (NFS/iSCSI/VM) does, regardless of the dataset's own
    sync property. Engine/direct detection is only for the banner here.
    """
    detect_io_settings(cfg)

banner(
        "SLOG / sync-write latency profile",
        cfg,
        extra=[
            "Method  : fio randwrite bs=4k --sync=1 (O_SYNC), psync engine",
            f"Profile : runtime={cfg.runtime}s/run  size={cfg.size.upper()}",
            "Note    : every write is a synchronous ZIL commit -- the load your",
            "          SLOG actually sees. Watch it live in another shell with:",
            "            zpool iostat -vl <pool> 1",
        ],
    )

# Measure everything first (progress -> stderr), then draw the whole table,
    # so the "measuring..." lines can't interleave into the middle of it.
    rows: list[list[str]] = []
    for threads in SLOG_THREADS:
        step(f"measuring 4K sync randwrite T{threads} ...")
        res = run_fio(
            cfg, rw="randwrite", bs="4k", iodepth=1, numjobs=threads,
            engine="psync", direct=False, sync=True,
        )
        rows.append([
            f"4K sync T{threads}",
            f"{res.iops:.0f}", f"{res.bw_mbps:.2f}",
            f"{res.p50_us:.1f}", f"{res.p99_us:.1f}",
        ])

_cleanup(cfg)

print()
    print(render_table(
        ["Test", "IOPS", "MB/s", "p50 lat(us)", "p99 lat(us)"],
        rows,
        col_styles=[CYAN, GREEN, GREEN, YELLOW, MAGENTA],
    ))
    print()
    print(OUT.paint("  Healthy Optane SLOG (eg P1600X) single-stream (T1) target:", BOLD))
    print(OUT.paint("    ~15-25k IOPS, p50 latency ~40-65us.", GREEN)
          + OUT.paint(" Much higher latency usually means", DIM))
    print(OUT.paint("    CPU C-states / PCIe ASPM / BIOS power profile (eg Dell DAPC) throttling.", DIM))
    print()

def dd_profile(cfg: Config) -> None:
    """Dependency-free fallback when fio isn't installed.

Far cruder than fio: a single sequential stream, and the cached-read figure
    will reflect RAM. Use it only as a rough sanity check.
    """
    count_mib = max(1, parse_size_bytes(cfg.size) // (1024**2))
    nbytes = count_mib * 1024**2

banner(
        "basic dd benchmark (fio not installed)",
        cfg,
        extra=[f"Profile : {count_mib} MiB sequential stream"],
    )

def timed_dd(args: list[str]) -> float:
        start = time.monotonic()
        subprocess.run(args, capture_output=True, text=True)
        elapsed = time.monotonic() - start
        return (nbytes / elapsed / 1e6) if elapsed > 0 else 0.0  # decimal MB/s

step("measuring sequential write ...")
    write_mbps = timed_dd([
        "dd", "if=/dev/zero", f"of={cfg.benchfile}",
        "bs=1M", f"count={count_mib}", "conv=fdatasync,notrunc",
    ])

step("dropping caches for uncached read ...")
    subprocess.run(["sudo", "sh", "-c", "echo 3 > /proc/sys/vm/drop_caches"],
                   capture_output=True, text=True)

step("measuring uncached read ...")
    uncached_mbps = timed_dd([
        "dd", f"if={cfg.benchfile}", "of=/dev/null", "bs=1M", f"count={count_mib}",
    ])

step("measuring cached read ...")
    cached_mbps = timed_dd([
        "dd", f"if={cfg.benchfile}", "of=/dev/null", "bs=1M", f"count={count_mib}",
    ])

_cleanup(cfg)

print()
    print(render_table(
        ["Test", "MB/s"],
        [
            ["Sequential write", f"{write_mbps:.2f}"],
            ["Uncached read", f"{uncached_mbps:.2f}"],
            ["Cached read (RAM)", f"{cached_mbps:.2f}"],
        ],
        col_styles=[CYAN, GREEN],
    ))
    print()

def _cleanup(cfg: Config) -> None:
    """Remove the shared bench file."""
    try:
        os.remove(cfg.benchfile)
    except OSError:
        pass

# --------------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------------- #

def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="speedtest-hd.py",
        description="CrystalDiskMark-style storage benchmark built on fio.",
        epilog=(
            "Examples:\n"
            "  speedtest-hd.py .\n"
            "  speedtest-hd.py /mnt/nvmepool --runtime=10 --size=4g\n"
            "  speedtest-hd.py /mnt/nfsshare --buffered\n"
            "  speedtest-hd.py /mnt/nvme-ultra-r10/vm-root --slog --runtime=30\n"
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    p.add_argument("path", help="directory/mount to benchmark ('.' for cwd)")

mode = p.add_mutually_exclusive_group()
    mode.add_argument("--fio", action="store_const", const="cdm", dest="mode",
                      help="force the fio CrystalDiskMark-style profile")
    mode.add_argument("--dd", action="store_const", const="dd", dest="mode",
                      help="force the basic dd fallback test")
    mode.add_argument("--slog", action="store_const", const="slog", dest="mode",
                      help="SLOG / sync-write latency profile (ZFS ZIL)")

direct = p.add_mutually_exclusive_group()
    direct.add_argument("--direct", action="store_const", const=True, dest="direct",
                        help="force O_DIRECT (bypass page cache)")
    direct.add_argument("--buffered", action="store_const", const=False, dest="direct",
                        help="force buffered IO (e.g. if O_DIRECT unsupported)")

p.add_argument("--engine", choices=ENGINE_CANDIDATES,
                   help="force a specific IO engine (default: auto-detect)")
    p.add_argument("--runtime", type=int, default=5, metavar="SEC",
                   help="seconds per run (default: 5, like CrystalDiskMark)")
    p.add_argument("--size", default="1g", metavar="SIZE",
                   help="test file size (default: 1g)")
    p.add_argument("--verbose", action="store_true",
                   help="also print the full fio output for every run")
    p.add_argument("-y", "--yes", action="store_true", dest="assume_yes",
                   help="skip the confirmation prompt")

p.set_defaults(mode=None, direct=None)
    return p

def confirm(cfg: Config) -> None:
    """Guard prompt -- we're about to write a multi-GB file to the target."""
    if cfg.assume_yes:
        return
    print(OUT.paint("NOTICE:", BOLD, YELLOW)
          + f" {cfg.size.upper()} free space on "
          + OUT.paint(f"'{cfg.path}'", CYAN)
          + " is required to perform the benchmark.")
    answer = input(f"Are you ready to start a storage benchmark against "
                   f"'{cfg.path}' ? ")
    if not answer.strip().lower().startswith("y"):
        print(OUT.paint("Ok, cancelled!", YELLOW))
        sys.exit(0)
    print(OUT.paint("Great! Starting benchmark now!", BOLD, GREEN))

def main(argv: Optional[Sequence[str]] = None) -> int:
    args = build_parser().parse_args(argv)

path = os.getcwd() if args.path == "." else args.path
    if not os.path.exists(path):
        print(f"Path {path} does not exist", file=sys.stderr)
        return 1

cfg = Config(
        path=path,
        mode=args.mode,
        engine=args.engine,
        direct=args.direct,
        runtime=args.runtime,
        size=args.size,
        verbose=args.verbose,
        assume_yes=args.assume_yes,
        benchfile=os.path.join(path, "speedtest-hd.bench"),
    )

have_fio = shutil.which("fio") is not None

# Resolve the mode: explicit flag wins; otherwise fio if available, else dd.
    mode = cfg.mode
    if mode in ("cdm", "slog") and not have_fio:
        print(ERR.paint("ERROR:", BOLD, RED)
              + " --fio/--slog require fio (apt install fio / pacman -S fio).",
              file=sys.stderr)
        return 1
    if mode is None:
        if have_fio:
            mode = "cdm"
        else:
            print(OUT.paint("\nfio is not installed -- falling back to basic dd test.",
                            YELLOW))
            print(OUT.paint("Install fio for the full CrystalDiskMark-style benchmark.",
                            DIM))
            mode = "dd"

confirm(cfg)

# Dictionary dispatch: map each mode name to its handler function (the values
    # are the functions themselves -- no parentheses), look up the one for `mode`,
    # then call it with (cfg). Equivalent to an if/elif chain over `mode`. `mode`
    # is always one of these three keys by now, so the lookup can't KeyError.
    {"cdm": cdm_profile, "slog": slog_profile, "dd": dd_profile}[mode](cfg)
    return 0

if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        print(ERR.paint("\nInterrupted.", YELLOW), file=sys.stderr)
        sys.exit(130)

1	#!/usr/bin/env python3
2	"""speedtest-hd — a CrystalDiskMark-style storage benchmark for Linux, on fio.
3
4	This is the Python successor to ``speedtest-hd.sh``. It measures storage the way
5	CrystalDiskMark does, and adds a dedicated SLOG / sync-write latency profile
6	for diagnosing ZFS ZIL performance (NFS / iSCSI / VM sync workloads).
7
8	Three profiles
9	--------------
10	* cdm — the four CrystalDiskMark default tests (``SEQ1M Q8T1``, ``SEQ1M
11	Q1T1``, ``RND4K Q32T16``, ``RND4K Q1T1``), each measured for read and
12	write, reported in MB/s and IOPS.
13	* slog — synchronous 4K random writes swept across thread counts, reporting
14	IOPS, MB/s and p50/p99 commit latency. This is the load a ZFS SLOG sees.
15	* dd — a dependency-free fallback when ``fio`` isn't installed.
16
17	Why fio JSON?
18	-------------
19	Every measurement asks fio for ``--output-format=json`` and parses it with the
20	standard library. That's the whole reason this is Python and not bash: robust,
21	unit-safe parsing of bandwidth / IOPS / latency percentiles with no fragile
22	text scraping.
23
24	See README.md for the full case study (diagnosing a "slow" Optane SLOG that
25	turned out to be CPU power management, not the disk).
26	"""
27
28	from __future__ import annotations
29
30	import argparse
31	import glob
32	import json
33	import os
34	import shutil
35	import statistics
36	import subprocess
37	import sys
38	import time
39	from dataclasses import dataclass, field
40	from typing import Callable, Optional, Sequence
41
42	# --------------------------------------------------------------------------- #
43	# Color / styling (standard library only -- raw ANSI escape codes)
44	# --------------------------------------------------------------------------- #
45
46	# SGR escape sequences. We stick to the basic 16-color set so output respects
47	# the user's terminal theme instead of hard-coding RGB.
48	RESET = "\033[0m"
49	BOLD = "\033[1m"
50	DIM = "\033[2m"
51	ITALIC = "\033[3m"
52	RED = "\033[31m"
53	GREEN = "\033[32m"
54	YELLOW = "\033[33m"
55	BLUE = "\033[34m"
56	MAGENTA = "\033[35m"
57	CYAN = "\033[36m"
58	WHITE = "\033[37m"
59	BRIGHT_CYAN = "\033[96m"
60
61
62	def supports_color(stream) -> bool:
63	"""Decide per-stream whether to emit ANSI codes.
64
65	Honors the de-facto ``NO_COLOR`` / ``FORCE_COLOR`` conventions and otherwise
66	colors only when the stream is an interactive terminal.
67	"""
68	if os.environ.get("NO_COLOR"):
69	return False
70	if os.environ.get("FORCE_COLOR"):
71	return True
72	if os.environ.get("TERM") == "dumb":
73	return False
74	return bool(getattr(stream, "isatty", None)) and stream.isatty()
75
76
77	class Painter:
78	"""Wraps text in SGR codes, or returns it untouched when color is off."""
79
80	def __init__(self, enabled: bool):
81	self.enabled = enabled
82
83	def paint(self, text: str, *codes: str) -> str:
84	if not self.enabled or not codes:
85	return text
86	return "".join(codes) + text + RESET
87
88
89	# stdout carries results (banner + tables); stderr carries progress + verbose
90	# fio dumps. Each gets its own enable flag so piping one but not the other works.
91	OUT = Painter(supports_color(sys.stdout))
92	ERR = Painter(supports_color(sys.stderr))
93
94
95	# --------------------------------------------------------------------------- #
96	# Constants
97	# --------------------------------------------------------------------------- #
98
99	# Preference order for the IO engine. io_uring is the modern, lowest-overhead
100	# Linux engine; libaio is older (and only truly async with O_DIRECT); posixaio
101	# and sync are the portable fallbacks (e.g. some NFS mounts).
102	ENGINE_CANDIDATES = ("io_uring", "libaio", "posixaio", "sync")
103
104	# Binary unit multipliers (fio treats "1g" as 1 GiB, so we match that).
105	_SIZE_UNITS = {"k": 1024, "m": 10242, "g": 10243, "t": 1024**4}
106
107
108	@dataclass(frozen=True)
109	class CdmTest:
110	"""One CrystalDiskMark test. ``seq`` picks sequential vs random patterns."""
111
112	label: str
113	bs: str
114	iodepth: int
115	numjobs: int
116	seq: bool
117
118	@property
119	def read_rw(self) -> str:
120	return "read" if self.seq else "randread"
121
122	@property
123	def write_rw(self) -> str:
124	return "write" if self.seq else "randwrite"
125
126
127	# The four CrystalDiskMark default tests, as data, in CrystalDiskMark's own
128	# display order. Q = queue depth (iodepth), T = threads (numjobs).
129	CDM_TESTS: tuple[CdmTest, ...] = (
130	CdmTest("SEQ1M Q8T1", bs="1m", iodepth=8, numjobs=1, seq=True),
131	CdmTest("SEQ1M Q1T1", bs="1m", iodepth=1, numjobs=1, seq=True),
132	CdmTest("RND4K Q32T16", bs="4k", iodepth=32, numjobs=16, seq=False),
133	CdmTest("RND4K Q1T1", bs="4k", iodepth=1, numjobs=1, seq=False),
134	)
135
136	# Thread counts for the SLOG sweep. T1 is the headline single-stream latency;
137	# the higher counts show how the SLOG scales as concurrent sync writers pile on.
138	SLOG_THREADS: tuple[int, ...] = (1, 4, 8, 16)
139
140
141	# --------------------------------------------------------------------------- #
142	# Configuration
143	# --------------------------------------------------------------------------- #
144
145
146	@dataclass
147	class Config:
148	"""Resolved run settings. ``engine`` and ``direct`` may start unset and get
149	filled in by :func:`detect_io_settings`."""
150
151	path: str
152	mode: Optional[str] # "cdm" \| "slog" \| "dd" \| None (auto)
153	engine: Optional[str] # None => auto-detect
154	direct: Optional[bool] # None => auto-detect
155	runtime: int
156	size: str
157	verbose: bool
158	assume_yes: bool
159	benchfile: str = field(default="")
160
161	@property
162	def direct_flag(self) -> int:
163	"""fio's --direct takes 0/1; default to 0 until detection runs."""
164	return 1 if self.direct else 0
165
166
167	# --------------------------------------------------------------------------- #
168	# Small helpers
169	# --------------------------------------------------------------------------- #
170
171
172	def log(message: str) -> None:
173	"""Progress/diagnostic line -> stderr, so stdout stays pure results."""
174	print(message, file=sys.stderr, flush=True)
175
176
177	def step(message: str) -> None:
178	"""A single colored progress line ('measuring ...') -> stderr."""
179	log(f" {ERR.paint('▶', BOLD, CYAN)} {ERR.paint(message, DIM)}")
180
181
182	def vsection(title: str) -> None:
183	"""A clear, ruled header for one --verbose fio section -> stderr."""
184	width = 78
185	head = f"─── {title} "
186	head += "─" * max(3, width - len(head))
187	log("")
188	log(ERR.paint(head, BOLD, CYAN))
189
190
191	def parse_size_bytes(size: str) -> int:
192	"""Turn an fio-style size ("4k", "1g", "512m", "2048") into bytes."""
193	s = size.strip().lower()
194	if s and s[-1] in _SIZE_UNITS:
195	return int(float(s[:-1]) * _SIZE_UNITS[s[-1]])
196	return int(s)
197
198
199	def render_table(
200	headers: Sequence[str],
201	rows: Sequence[Sequence[str]],
202	aligns: Optional[Sequence[str]] = None,
203	col_styles: Optional[Sequence[Optional[str]]] = None,
204	) -> str:
205	"""Render a bordered ASCII table whose columns auto-size to their content.
206
207	``aligns`` is a per-column "<" (left) or ">" (right); by default the first
208	column is left-aligned (the label) and the rest are right-aligned (numbers).
209	``col_styles`` is an optional per-column ANSI style applied to body cells
210	(header is always bold, borders dim). Widths are computed on the plain
211	text and color is applied after padding, so escape codes never skew layout.
212	"""
213	ncols = len(headers)
214	if aligns is None:
215	aligns = ["<"] + [">"] * (ncols - 1)
216	if col_styles is None:
217	col_styles = [None] * ncols
218
219	# Widest cell (header or any row) sets each column's width.
220	widths = [
221	max([len(str(headers[c]))] + [len(str(r[c])) for r in rows])
222	for c in range(ncols)
223	]
224
225	bar = OUT.paint("\|", DIM)
226	sep = OUT.paint(" \| ", DIM)
227
228	def rule() -> str:
229	return OUT.paint("+" + "+".join("-" * (w + 2) for w in widths) + "+", DIM)
230
231	def line(cells: Sequence[str], *, header: bool = False) -> str:
232	parts = []
233	for c in range(ncols):
234	padded = f"{str(cells[c]):{aligns[c]}{widths[c]}}"
235	if header:
236	parts.append(OUT.paint(padded, BOLD))
237	elif col_styles[c]:
238	parts.append(OUT.paint(padded, col_styles[c]))
239	else:
240	parts.append(padded)
241	return f"{bar} " + sep.join(parts) + f" {bar}"
242
243	out = [rule(), line(headers, header=True), rule()]
244	out += [line(r) for r in rows]
245	out.append(rule())
246	return "\n".join(out)
247
248
249	# --------------------------------------------------------------------------- #
250	# fio: results, invocation, parsing
251	# --------------------------------------------------------------------------- #
252
253
254	@dataclass(frozen=True)
255	class FioResult:
256	"""Parsed metrics for one direction (read or write) of an fio run.
257
258	Latencies are completion latency (clat) in microseconds; for synchronous
259	writes that is effectively the durable-commit latency.
260	"""
261
262	bw_mbps: float # decimal MB/s, matching CrystalDiskMark's SI figure
263	iops: float
264	p50_us: float
265	p99_us: float
266	mean_us: float
267
268	@classmethod
269	def zero(cls) -> "FioResult":
270	return cls(0.0, 0.0, 0.0, 0.0, 0.0)
271
272
273	def _direction_of(rw: str) -> str:
274	"""fio reports 'read' and 'write' sub-objects; pick the active one."""
275	return "read" if "read" in rw else "write"
276
277
278	def _aggregate(jobs: list[dict], direction: str) -> FioResult:
279	"""Combine per-job fio JSON stats into a single :class:`FioResult`.
280
281	We deliberately run fio without --group_reporting and aggregate ourselves,
282	which sidesteps fio's group-merge quirks: sum throughput, average the median
283	latency, and take the worst-case tail (p99) across jobs.
284	"""
285
286	def section(job: dict) -> dict:
287	return job.get(direction, {})
288
289	bw_mbps = sum(section(j).get("bw_bytes", 0) for j in jobs) / 1e6
290	iops = sum(section(j).get("iops", 0.0) for j in jobs)
291
292	def percentiles(key: str) -> list[float]:
293	vals = []
294	for j in jobs:
295	pct = section(j).get("clat_ns", {}).get("percentile", {})
296	v = pct.get(key)
297	if v is not None:
298	vals.append(v)
299	return vals
300
301	p50_ns = percentiles("50.000000")
302	p99_ns = percentiles("99.000000")
303	mean_ns = [section(j).get("clat_ns", {}).get("mean", 0.0) for j in jobs]
304
305	return FioResult(
306	bw_mbps=bw_mbps,
307	iops=iops,
308	p50_us=(statistics.mean(p50_ns) / 1000.0) if p50_ns else 0.0, # avg of medians
309	p99_us=(max(p99_ns) / 1000.0) if p99_ns else 0.0, # worst tail
310	mean_us=(statistics.mean(mean_ns) / 1000.0) if mean_ns else 0.0,
311	)
312
313
314	def run_fio(
315	cfg: Config,
316	*,
317	rw: str,
318	bs: str,
319	iodepth: int,
320	numjobs: int,
321	engine: Optional[str] = None,
322	direct: Optional[bool] = None,
323	sync: bool = False,
324	end_fsync: bool = False,
325	) -> FioResult:
326	"""Run a single fio job against the shared bench file and parse its JSON.
327
328	``engine``/``direct`` default to the detected config values; the SLOG path
329	overrides them (psync + buffered + O_SYNC). ``end_fsync`` flushes the device
330	cache at the end of a write run so cached writes can't inflate the result.
331	"""
332	engine = engine if engine is not None else cfg.engine
333	direct = direct if direct is not None else cfg.direct
334
335	cmd = [
336	"sudo", "fio",
337	"--name=speedtest",
338	f"--filename={cfg.benchfile}",
339	f"--ioengine={engine}",
340	f"--direct={1 if direct else 0}",
341	f"--rw={rw}",
342	f"--bs={bs}",
343	f"--size={cfg.size}",
344	f"--numjobs={numjobs}",
345	f"--iodepth={iodepth}",
346	"--time_based",
347	f"--runtime={cfg.runtime}",
348	"--output-format=json",
349	]
350	if sync:
351	cmd.append("--sync=1") # O_SYNC: every write is a durable commit
352	if end_fsync:
353	cmd.append("--end_fsync=1") # flush device cache once at the end
354
355	proc = subprocess.run(cmd, capture_output=True, text=True)
356
357	if cfg.verbose:
358	sync_tag = " sync=1" if sync else ""
359	vsection(f"fio · {rw} bs={bs} iodepth={iodepth} numjobs={numjobs}{sync_tag}")
360	log(" " + ERR.paint("$ " + " ".join(cmd[1:]), DIM, ITALIC))
361	log(ERR.paint(proc.stdout.strip() or "(no stdout)", DIM))
362	if proc.stderr.strip():
363	log(ERR.paint(proc.stderr.strip(), YELLOW))
364
365	try:
366	data = json.loads(proc.stdout)
367	return _aggregate(data["jobs"], _direction_of(rw))
368	except (json.JSONDecodeError, KeyError, ValueError):
369	# fio failed or produced no parseable JSON; report zeros rather than
370	# crashing the whole run (verbose mode above shows what went wrong).
371	return FioResult.zero()
372
373
374	def fio_probe(path: str, engine: str, direct: bool) -> bool:
375	"""Throwaway 1s fio job to see if an (engine, O_DIRECT) combo works here.
376
377	This is how we stay accurate and portable across ext4/xfs/btrfs/ZFS/NFS
378	without the caller needing to know what each filesystem supports.
379	"""
380	cmd = [
381	"sudo", "fio", "--name=probe",
382	f"--directory={path}",
383	f"--ioengine={engine}",
384	"--rw=write", "--bs=4k", "--size=1m",
385	f"--direct={1 if direct else 0}",
386	"--time_based", "--runtime=1",
387	]
388	proc = subprocess.run(cmd, capture_output=True, text=True)
389	for leftover in glob.glob(os.path.join(path, "probe*")):
390	try:
391	os.remove(leftover)
392	except OSError:
393	pass
394	return proc.returncode == 0
395
396
397	def detect_io_settings(cfg: Config) -> None:
398	"""Fill in ``cfg.engine`` and ``cfg.direct`` if the user didn't force them.
399
400	O_DIRECT bypasses the OS page cache so we measure the device, not RAM. Not
401	every filesystem supports it (older OpenZFS, some NFS), so we probe and fall
402	back to buffered rather than erroring out.
403	"""
404	if cfg.engine is None:
405	for engine in ENGINE_CANDIDATES:
406	if fio_probe(cfg.path, engine, direct=False):
407	cfg.engine = engine
408	break
409	else:
410	cfg.engine = "sync"
411
412	if cfg.direct is None:
413	cfg.direct = fio_probe(cfg.path, cfg.engine, direct=True)
414
415
416	# --------------------------------------------------------------------------- #
417	# Profiles
418	# --------------------------------------------------------------------------- #
419
420
421	def _meta_line(text: str) -> str:
422	"""Color the 'Label :' prefix of a banner metadata line, leave value as-is.
423
424	The value may already contain its own ANSI codes (e.g. a red O_DIRECT
425	warning); since these lines aren't width-aligned that's harmless.
426	"""
427	key, sep, val = text.partition(":")
428	if not sep:
429	return " " + text
430	return " " + OUT.paint(key + ":", BOLD, BRIGHT_CYAN) + val
431
432
433	def banner(title: str, cfg: Config, extra: Sequence[str] = ()) -> None:
434	full = f"speedtest-hd · {title}"
435	inner = len(full) + 4
436
437	print()
438	print(OUT.paint("╭" + "─" * inner + "╮", BOLD, CYAN))
439	print(
440	OUT.paint("│", BOLD, CYAN)
441	+ " " + OUT.paint(full, BOLD, WHITE) + " "
442	+ OUT.paint("│", BOLD, CYAN)
443	)
444	print(OUT.paint("╰" + "─" * inner + "╯", BOLD, CYAN))
445	print(_meta_line(f"Target : {OUT.paint(cfg.path, CYAN)}"))
446	for line_ in extra:
447	print(_meta_line(line_))
448	print()
449
450
451	def cdm_profile(cfg: Config) -> None:
452	"""The CrystalDiskMark-style profile: 4 tests x (read, write), MB/s + IOPS."""
453	detect_io_settings(cfg)
454
455	if cfg.direct:
456	direct_label = OUT.paint("enabled (device)", GREEN)
457	else:
458	direct_label = OUT.paint(
459	"DISABLED (buffered -- may reflect RAM cache!)", BOLD, RED
460	)
461
462	banner(
463	"CrystalDiskMark-style storage benchmark",
464	cfg,
465	extra=[
466	f"Engine : {cfg.engine} O_DIRECT: {direct_label}",
467	f"Profile : size={cfg.size.upper()} runtime={cfg.runtime}s/run (8 runs)",
468	],
469	)
470
471	rows: list[list[str]] = []
472	for test in CDM_TESTS:
473	step(f"measuring {test.label.strip()} ...")
474	r = run_fio(cfg, rw=test.read_rw, bs=test.bs,
475	iodepth=test.iodepth, numjobs=test.numjobs)
476	w = run_fio(cfg, rw=test.write_rw, bs=test.bs,
477	iodepth=test.iodepth, numjobs=test.numjobs, end_fsync=True)
478	rows.append([
479	test.label,
480	f"{r.bw_mbps:.2f}", f"{w.bw_mbps:.2f}",
481	f"{r.iops:.0f}", f"{w.iops:.0f}",
482	])
483
484	_cleanup(cfg)
485
486	print()
487	print(render_table(
488	["Test", "Read (MB/s)", "Write (MB/s)", "Read (IOPS)", "Write (IOPS)"],
489	rows,
490	col_styles=[CYAN, GREEN, YELLOW, GREEN, YELLOW],
491	))
492	print()
493
494
495	def slog_profile(cfg: Config) -> None:
496	"""SLOG / sync-write latency profile.
497
498	Forces synchronous 4K random writes (O_SYNC) via the portable psync engine,
499	so every write is a ZIL commit -- exercising a ZFS SLOG exactly the way a
500	sync=always dataset (NFS/iSCSI/VM) does, regardless of the dataset's own
501	sync property. Engine/direct detection is only for the banner here.
502	"""
503	detect_io_settings(cfg)
504
505	banner(
506	"SLOG / sync-write latency profile",
507	cfg,
508	extra=[
509	"Method : fio randwrite bs=4k --sync=1 (O_SYNC), psync engine",
510	f"Profile : runtime={cfg.runtime}s/run size={cfg.size.upper()}",
511	"Note : every write is a synchronous ZIL commit -- the load your",
512	" SLOG actually sees. Watch it live in another shell with:",
513	" zpool iostat -vl <pool> 1",
514	],
515	)
516
517	# Measure everything first (progress -> stderr), then draw the whole table,
518	# so the "measuring..." lines can't interleave into the middle of it.
519	rows: list[list[str]] = []
520	for threads in SLOG_THREADS:
521	step(f"measuring 4K sync randwrite T{threads} ...")
522	res = run_fio(
523	cfg, rw="randwrite", bs="4k", iodepth=1, numjobs=threads,
524	engine="psync", direct=False, sync=True,
525	)
526	rows.append([
527	f"4K sync T{threads}",
528	f"{res.iops:.0f}", f"{res.bw_mbps:.2f}",
529	f"{res.p50_us:.1f}", f"{res.p99_us:.1f}",
530	])
531
532	_cleanup(cfg)
533
534	print()
535	print(render_table(
536	["Test", "IOPS", "MB/s", "p50 lat(us)", "p99 lat(us)"],
537	rows,
538	col_styles=[CYAN, GREEN, GREEN, YELLOW, MAGENTA],
539	))
540	print()
541	print(OUT.paint(" Healthy Optane SLOG (eg P1600X) single-stream (T1) target:", BOLD))
542	print(OUT.paint(" ~15-25k IOPS, p50 latency ~40-65us.", GREEN)
543	+ OUT.paint(" Much higher latency usually means", DIM))
544	print(OUT.paint(" CPU C-states / PCIe ASPM / BIOS power profile (eg Dell DAPC) throttling.", DIM))
545	print()
546
547
548	def dd_profile(cfg: Config) -> None:
549	"""Dependency-free fallback when fio isn't installed.
550
551	Far cruder than fio: a single sequential stream, and the cached-read figure
552	will reflect RAM. Use it only as a rough sanity check.
553	"""
554	count_mib = max(1, parse_size_bytes(cfg.size) // (1024**2))
555	nbytes = count_mib * 1024**2
556
557	banner(
558	"basic dd benchmark (fio not installed)",
559	cfg,
560	extra=[f"Profile : {count_mib} MiB sequential stream"],
561	)
562
563	def timed_dd(args: list[str]) -> float:
564	start = time.monotonic()
565	subprocess.run(args, capture_output=True, text=True)
566	elapsed = time.monotonic() - start
567	return (nbytes / elapsed / 1e6) if elapsed > 0 else 0.0 # decimal MB/s
568
569	step("measuring sequential write ...")
570	write_mbps = timed_dd([
571	"dd", "if=/dev/zero", f"of={cfg.benchfile}",
572	"bs=1M", f"count={count_mib}", "conv=fdatasync,notrunc",
573	])
574
575	step("dropping caches for uncached read ...")
576	subprocess.run(["sudo", "sh", "-c", "echo 3 > /proc/sys/vm/drop_caches"],
577	capture_output=True, text=True)
578
579	step("measuring uncached read ...")
580	uncached_mbps = timed_dd([
581	"dd", f"if={cfg.benchfile}", "of=/dev/null", "bs=1M", f"count={count_mib}",
582	])
583
584	step("measuring cached read ...")
585	cached_mbps = timed_dd([
586	"dd", f"if={cfg.benchfile}", "of=/dev/null", "bs=1M", f"count={count_mib}",
587	])
588
589	_cleanup(cfg)
590
591	print()
592	print(render_table(
593	["Test", "MB/s"],
594	[
595	["Sequential write", f"{write_mbps:.2f}"],
596	["Uncached read", f"{uncached_mbps:.2f}"],
597	["Cached read (RAM)", f"{cached_mbps:.2f}"],
598	],
599	col_styles=[CYAN, GREEN],
600	))
601	print()
602
603
604	def _cleanup(cfg: Config) -> None:
605	"""Remove the shared bench file."""
606	try:
607	os.remove(cfg.benchfile)
608	except OSError:
609	pass
610
611
612	# --------------------------------------------------------------------------- #
613	# CLI
614	# --------------------------------------------------------------------------- #
615
616
617	def build_parser() -> argparse.ArgumentParser:
618	p = argparse.ArgumentParser(
619	prog="speedtest-hd.py",
620	description="CrystalDiskMark-style storage benchmark built on fio.",
621	epilog=(
622	"Examples:\n"
623	" speedtest-hd.py .\n"
624	" speedtest-hd.py /mnt/nvmepool --runtime=10 --size=4g\n"
625	" speedtest-hd.py /mnt/nfsshare --buffered\n"
626	" speedtest-hd.py /mnt/nvme-ultra-r10/vm-root --slog --runtime=30\n"
627	),
628	formatter_class=argparse.RawDescriptionHelpFormatter,
629	)
630	p.add_argument("path", help="directory/mount to benchmark ('.' for cwd)")
631
632	mode = p.add_mutually_exclusive_group()
633	mode.add_argument("--fio", action="store_const", const="cdm", dest="mode",
634	help="force the fio CrystalDiskMark-style profile")
635	mode.add_argument("--dd", action="store_const", const="dd", dest="mode",
636	help="force the basic dd fallback test")
637	mode.add_argument("--slog", action="store_const", const="slog", dest="mode",
638	help="SLOG / sync-write latency profile (ZFS ZIL)")
639
640	direct = p.add_mutually_exclusive_group()
641	direct.add_argument("--direct", action="store_const", const=True, dest="direct",
642	help="force O_DIRECT (bypass page cache)")
643	direct.add_argument("--buffered", action="store_const", const=False, dest="direct",
644	help="force buffered IO (e.g. if O_DIRECT unsupported)")
645
646	p.add_argument("--engine", choices=ENGINE_CANDIDATES,
647	help="force a specific IO engine (default: auto-detect)")
648	p.add_argument("--runtime", type=int, default=5, metavar="SEC",
649	help="seconds per run (default: 5, like CrystalDiskMark)")
650	p.add_argument("--size", default="1g", metavar="SIZE",
651	help="test file size (default: 1g)")
652	p.add_argument("--verbose", action="store_true",
653	help="also print the full fio output for every run")
654	p.add_argument("-y", "--yes", action="store_true", dest="assume_yes",
655	help="skip the confirmation prompt")
656
657	p.set_defaults(mode=None, direct=None)
658	return p
659
660
661	def confirm(cfg: Config) -> None:
662	"""Guard prompt -- we're about to write a multi-GB file to the target."""
663	if cfg.assume_yes:
664	return
665	print(OUT.paint("NOTICE:", BOLD, YELLOW)
666	+ f" {cfg.size.upper()} free space on "
667	+ OUT.paint(f"'{cfg.path}'", CYAN)
668	+ " is required to perform the benchmark.")
669	answer = input(f"Are you ready to start a storage benchmark against "
670	f"'{cfg.path}' ? ")
671	if not answer.strip().lower().startswith("y"):
672	print(OUT.paint("Ok, cancelled!", YELLOW))
673	sys.exit(0)
674	print(OUT.paint("Great! Starting benchmark now!", BOLD, GREEN))
675
676
677	def main(argv: Optional[Sequence[str]] = None) -> int:
678	args = build_parser().parse_args(argv)
679
680	path = os.getcwd() if args.path == "." else args.path
681	if not os.path.exists(path):
682	print(f"Path {path} does not exist", file=sys.stderr)
683	return 1
684
685	cfg = Config(
686	path=path,
687	mode=args.mode,
688	engine=args.engine,
689	direct=args.direct,
690	runtime=args.runtime,
691	size=args.size,
692	verbose=args.verbose,
693	assume_yes=args.assume_yes,
694	benchfile=os.path.join(path, "speedtest-hd.bench"),
695	)
696
697	have_fio = shutil.which("fio") is not None
698
699	# Resolve the mode: explicit flag wins; otherwise fio if available, else dd.
700	mode = cfg.mode
701	if mode in ("cdm", "slog") and not have_fio:
702	print(ERR.paint("ERROR:", BOLD, RED)
703	+ " --fio/--slog require fio (apt install fio / pacman -S fio).",
704	file=sys.stderr)
705	return 1
706	if mode is None:
707	if have_fio:
708	mode = "cdm"
709	else:
710	print(OUT.paint("\nfio is not installed -- falling back to basic dd test.",
711	YELLOW))
712	print(OUT.paint("Install fio for the full CrystalDiskMark-style benchmark.",
713	DIM))
714	mode = "dd"
715
716	confirm(cfg)
717
718	# Dictionary dispatch: map each mode name to its handler function (the values
719	# are the functions themselves -- no parentheses), look up the one for `mode`,
720	# then call it with (cfg). Equivalent to an if/elif chain over `mode`. `mode`
721	# is always one of these three keys by now, so the lookup can't KeyError.
722	{"cdm": cdm_profile, "slog": slog_profile, "dd": dd_profile}[mode](cfg)
723	return 0
724
725
726	if __name__ == "__main__":
727	try:
728	sys.exit(main())
729	except KeyboardInterrupt:
730	print(ERR.paint("\nInterrupted.", YELLOW), file=sys.stderr)
731	sys.exit(130)
732