# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Group commit and the square-root wait.

Part 1: append 100-byte records and flush after every 1, 10 and 100 records
(F_FULLFSYNC on macOS, fsync elsewhere): records per second. Part 2: records
arrive at rate λ; a timer-based batcher flushes every W seconds. The mean
latency per record (waiting for the batch, then the flush) for several W,
against W* = √(2 · fsync / λ).
"""
import argparse, json, math, os, random, sys, tempfile, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--records", type=int, default=2000, help="Records per batch-size run.")
p.add_argument("--rate", type=float, default=2000, help="Arrival rate λ (records/s) in part 2.")
p.add_argument("--waits-ms", default="0.5,1,2,4,8,16", help="Batch waits W to try (ms).")
p.add_argument("--seconds", type=float, default=4, help="Seconds per W in part 2.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

def flush(fd):
    if sys.platform == "darwin":
        import fcntl; fcntl.fcntl(fd, fcntl.F_FULLFSYNC)
    else: os.fsync(fd)
path = os.path.join(tempfile.mkdtemp(), "gc.log"); fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_APPEND)
part1 = []
for every in (1, 10, 100):
    t = time.perf_counter()
    for i in range(args.records):
        os.write(fd, b"x" * 100)
        if (i + 1) % every == 0: flush(fd)
    el = time.perf_counter() - t; part1.append({"flush_every": every, "records_per_s": args.records / el})
t = time.perf_counter()
for _ in range(50): os.write(fd, b"x" * 100); flush(fd)
fsync_s = (time.perf_counter() - t) / 50
part2 = []
for w_ms in map(float, args.waits_ms.split(",")):
    W = w_ms / 1000; rng = random.Random(1); lat = []; now = time.perf_counter(); end = now + args.seconds
    next_arrival = now; pending = []; next_flush = now + W
    while time.perf_counter() < end:
        t_now = time.perf_counter()
        while next_arrival <= t_now: pending.append(next_arrival); os.write(fd, b"x" * 100); next_arrival += rng.expovariate(args.rate)
        if t_now >= next_flush:
            if pending:
                flush(fd); done = time.perf_counter(); lat += [done - a for a in pending]; pending = []
            next_flush = max(next_flush + W, time.perf_counter())
    part2.append({"wait_ms": w_ms, "mean_latency_ms": sum(lat) / len(lat) * 1000, "records": len(lat)})
os.close(fd); os.remove(path)
best = min(part2, key=lambda r: r["mean_latency_ms"])
summary = {"fsync_ms": fsync_s * 1000, "w_star_ms": math.sqrt(2 * fsync_s / args.rate) * 1000, "best_w_ms": best["wait_ms"]}
for r in part1: print(f"flush every {r['flush_every']:>3} records: {r['records_per_s']:10,.0f} records/s")
print(f"one flush = {summary['fsync_ms']:.2f} ms; W* = √(2·fsync/λ) = {summary['w_star_ms']:.2f} ms at λ = {args.rate:g}/s")
for r in part2: print(f"  W = {r['wait_ms']:>4g} ms: mean latency {r['mean_latency_ms']:6.2f} ms over {r['records']:,} records")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"summary": summary, "tables": {"batches": part1, "waits": part2}}, f)