# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Group commit: measure fsync, and share it.

Part A: the cost of one durable 200-byte write (fsync; F_FULLFSYNC on macOS,
which is what reaches the device there). Part B: 20 records, then one flush.
Part C: pgbench in the PostgreSQL container, 1 and 32 clients, with
synchronous_commit on and off.
Part C needs the PostgreSQL service from the compose file (port 55805) and
the docker CLI.
"""
import argparse, json, os, re, subprocess, sys, tempfile, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--writes", type=int, default=2000, help="Durable writes in part A.")
p.add_argument("--batch", type=int, default=20, help="Records per flush in part B.")
p.add_argument("--container", default="group-commit-measure-fsync-and-the-postgres-1", help="PostgreSQL container for part C.")
p.add_argument("--seconds", type=int, default=30, help="pgbench duration.")
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(), "probe.log"); fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_APPEND)
t = time.perf_counter()
for _ in range(args.writes): os.write(fd, b"x" * 200); flush(fd)
per_a = (time.perf_counter() - t) / args.writes * 1e6
t = time.perf_counter(); n = 0
while n < args.writes:
    for _ in range(args.batch): os.write(fd, b"x" * 200); n += 1
    flush(fd)
per_b = (time.perf_counter() - t) / args.writes * 1e6
os.close(fd); os.remove(path)
parts = [{"part": "A: flush every record", "us_per_record": per_a}, {"part": f"B: flush every {args.batch} records", "us_per_record": per_b}]
def sh(*a): return subprocess.run(["docker", "exec", args.container, *a], capture_output=True, text=True)
pg = []
if sh("true").returncode == 0:
    sh("pgbench", "-U", "postgres", "-i", "-s", "10", "postgres")
    for sync in ("on", "off"):
        for clients in (1, 32):
            out = subprocess.run(["docker", "exec", "-e", f"PGOPTIONS=-c synchronous_commit={sync}", args.container, "pgbench", "-U", "postgres",
                                  "-c", str(clients), "-j", str(min(clients, 4)), "-T", str(args.seconds), "-N", "postgres"], capture_output=True, text=True)
            tps = float(re.search(r"tps = ([\d.]+)", out.stdout).group(1))
            lat = float(re.search(r"latency average = ([\d.]+) ms", out.stdout).group(1))
            pg.append({"synchronous_commit": sync, "clients": clients, "tps": tps, "latency_ms": lat})
summary = {"batch_gain": per_a / per_b}
for r in parts: print(f"{r['part']:<28} {r['us_per_record']:9.1f} µs per record")
for r in pg: print(f"pgbench synchronous_commit={r['synchronous_commit']:<3} clients={r['clients']:>2}: {r['tps']:9.0f} tps, latency {r['latency_ms']:.2f} ms")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"summary": summary, "tables": {"flushes": parts, "pgbench": pg}}, f)