# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Queue depth and the device limit.

A 2 GB file read in random 4 KB blocks, bypassing the page cache (F_NOCACHE
on macOS, posix_fadvise elsewhere), by 1, 4, 16, 64 and 128 threads each with
one request in flight: IOPS against requests in flight. Then the same random
reads through mmap with one thread.
"""
import argparse, json, mmap, os, random, sys, tempfile, threading, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--file-mb", type=int, default=2048, help="Test file (MB).")
p.add_argument("--depths", default="1,4,16,64,128", help="Requests in flight.")
p.add_argument("--seconds", type=float, default=5, help="Seconds per depth.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

path = os.path.join(tempfile.mkdtemp(), "qd.bin")
with open(path, "wb") as f:
    for _ in range(args.file_mb): f.write(os.urandom(1 << 20))
blocks = args.file_mb * 256
def open_nocache():
    fd = os.open(path, os.O_RDONLY)
    if sys.platform == "darwin":
        import fcntl; fcntl.fcntl(fd, fcntl.F_NOCACHE, 1)
    return fd
def run(depth):
    counts = [0] * depth; stop = time.perf_counter() + args.seconds
    def worker(i):
        fd = open_nocache(); rng = random.Random(i); n = 0
        while time.perf_counter() < stop:
            os.pread(fd, 4096, rng.randrange(blocks) * 4096); n += 1
        counts[i] = n; os.close(fd)
    ts = [threading.Thread(target=worker, args=(i,)) for i in range(depth)]
    for t in ts: t.start()
    for t in ts: t.join()
    iops = sum(counts) / args.seconds
    return {"in_flight": depth, "iops": iops, "mb_s": iops * 4096 / 2**20, "latency_us": depth / iops * 1e6}
if hasattr(os, "posix_fadvise"):
    fd = os.open(path, os.O_RDONLY); os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED); os.close(fd)
rows = [run(int(d)) for d in args.depths.split(",")]
with open(path, "rb") as f:
    m = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ); rng = random.Random(9); n = 0; stop = time.perf_counter() + args.seconds
    while time.perf_counter() < stop: m[rng.randrange(blocks) * 4096]; n += 1
    m.close()
os.remove(path)
summary = {"mmap_iops": n / args.seconds, "depth1_iops": rows[0]["iops"], "depth1_latency_us": rows[0]["latency_us"], "peak_iops": max(r["iops"] for r in rows)}
for r in rows: print(f"{r['in_flight']:>4} in flight: {r['iops']:>9,.0f} IOPS  ({r['mb_s']:7.1f} MB/s, {r['latency_us']:7.1f} µs per read)")
print(f"mmap, one thread: {summary['mmap_iops']:,.0f} reads/s (pages may come from the cache after the first touch)")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"summary": summary, "tables": {"depths": rows}}, f)