# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Queueing on the spill device: R = S / (1 − ρ).

Measures the device's sequential write throughput with one unthrottled
writer, then runs 1 to 5 rate-limited writers (each like fio --rw=write
--bs=1m --rate=...) at a fixed share of that throughput and records the mean
completion latency of a 1 MB write, measured from when it was due. Writes
bypass the page cache (F_NOCACHE on macOS, O_DIRECT is not used; on Linux
each write is followed by fdatasync).
"""
import argparse, fcntl, json, os, statistics, sys, tempfile, threading, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--share", type=float, default=0.2, help="Each job's rate as a share of the measured throughput.")
p.add_argument("--jobs", default="1,2,3,4,5", help="Numbers of jobs to run.")
p.add_argument("--seconds", type=float, default=15, help="Seconds per run.")
p.add_argument("--size-mb", type=int, default=1024, help="File size per job (the job wraps around).")
p.add_argument("--dir", default=tempfile.gettempdir(), help="Directory for the job files.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

BLOCK = os.urandom(1 << 20)
def opener(i):
    fd = os.open(os.path.join(args.dir, f"spill-job-{i}.bin"), os.O_WRONLY | os.O_CREAT)
    if sys.platform == "darwin": fcntl.fcntl(fd, fcntl.F_NOCACHE, 1)
    return fd
def write(fd, off):
    os.pwrite(fd, BLOCK, off)
    if sys.platform != "darwin": os.fdatasync(fd)

def job(i, rate_mb, end, lat):
    fd = opener(i); n = 0; t0 = time.perf_counter()
    while True:
        due = t0 + n / rate_mb if rate_mb else time.perf_counter()
        if due >= end: break
        delay = due - time.perf_counter()
        if delay > 0: time.sleep(delay)
        write(fd, (n % args.size_mb) << 20); lat.append(time.perf_counter() - due); n += 1
    os.close(fd)

def run(k, rate_mb):
    lats = [[] for _ in range(k)]; end = time.perf_counter() + args.seconds
    th = [threading.Thread(target=job, args=(i, rate_mb, end, lats[i])) for i in range(k)]
    for t in th: t.start()
    for t in th: t.join()
    allv = [x for l in lats for x in l]
    return allv, sum(len(l) for l in lats) / args.seconds

_, peak = run(1, 0)                                                    # unthrottled: the device's throughput
service = 1 / peak                                                     # seconds per 1 MB at full speed
rate = peak * args.share; rows = []
for k in map(int, args.jobs.split(",")):
    lat, got = run(k, rate)
    rho = k * args.share
    rows.append({"jobs": k, "rho": rho, "mb_per_s": got, "mean_ms": statistics.mean(lat) * 1000,
                 "p99_ms": sorted(lat)[int(0.99 * len(lat))] * 1000, "model_ms": service / (1 - rho) * 1000 if rho < 1 else None})
for r in rows: r["vs_one_job"] = r["mean_ms"] / rows[0]["mean_ms"]
for i in range(5): 
    try: os.remove(os.path.join(args.dir, f"spill-job-{i}.bin"))
    except FileNotFoundError: pass
summary = {"device_mb_s": peak, "service_ms": service * 1000, "rate_mb_s": rate}
print(f"device: {peak:.0f} MB/s sequential writes (1 MB takes {service * 1000:.2f} ms); each job writes {rate:.0f} MB/s")
for r in rows:
    m = f"{r['model_ms']:.2f} ms" if r["model_ms"] else "unbounded"
    print(f"{r['jobs']} job(s), ρ = {r['rho']:.1f}: mean latency {r['mean_ms']:8.2f} ms ({r['vs_one_job']:6.1f}x one job), p99 {r['p99_ms']:8.1f} ms, achieved {r['mb_per_s']:.0f} MB/s   S/(1-ρ) = {m}")

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