#!/usr/bin/env python3
"""Measure whether concurrent writes to one pipe arrive whole, with real processes.

The Shell Expansion page folds the pipe in as one more dup2, and its ledger
carries the pipe's other promise: the Seventh Edition manual says writes of
4096 bytes or less are atomic, POSIX says PIPE_BUF is at least 512, Linux's
pipe(7) says PIPE_BUF is 4096 and the capacity 65,536. Nothing in a manual
says what happens above those numbers on a given kernel, so this measures it.

Each round makes one pipe, forks WRITERS processes, releases them together,
and has every one of them make exactly ONE write(2) of SIZE bytes, all bytes
equal to that writer's own tag. The parent reads until every writer has
closed. A write arrived whole when its tag forms one contiguous run in what
was read; it was torn when its bytes arrive in more than one run.

Threads would not do: CPython's lock serialises them, and an earlier attempt
at this measurement saw nothing tear at any size for that reason. Separate
processes contend inside the kernel, which is the only place the promise
means anything.

    python3 scripts/shell-expansion-pipe-atomicity.py            print a table
    python3 scripts/shell-expansion-pipe-atomicity.py --json F   and write it to F
"""
import argparse
import fcntl
import json
import os
import platform
import re
import sys
from datetime import date

F_GETPIPE_SZ = getattr(fcntl, "F_GETPIPE_SZ", 1032)
SIZES = [512, 4096, 4097, 65536, 65537, 200000]
WRITERS = 4
ROUNDS = 200


def one_round(size: int, writers: int) -> tuple[int, int]:
    """Returns (writes torn in this round, bytes read)."""
    r, w = os.pipe()
    go_r, go_w = os.pipe()
    kids = []
    for k in range(writers):
        pid = os.fork()
        if pid == 0:
            os.close(r)
            os.close(go_w)
            os.read(go_r, 1)          # wait for the starting gun
            buf = bytes([65 + k]) * size
            n = os.write(w, buf)      # exactly one write(2)
            os._exit(0 if n == size else 3)
        kids.append(pid)
    os.close(w)
    os.close(go_r)
    os.write(go_w, b"x" * writers)    # release them together
    os.close(go_w)
    chunks = []
    while True:
        b = os.read(r, 1 << 20)
        if not b:
            break
        chunks.append(b)
    os.close(r)
    for pid in kids:
        _, status = os.waitpid(pid, 0)
        if os.waitstatus_to_exitcode(status) != 0:
            raise SystemExit(f"a writer made a short write at size {size}")
    data = b"".join(chunks)
    if len(data) != size * writers:
        raise SystemExit(f"read {len(data)} bytes, expected {size * writers}")
    # Maximal runs of one tag. A write that arrived whole is exactly one run.
    runs = {}
    for m in re.finditer(rb"(.)\1*", data, re.S):
        c = m.group(1)
        runs[c] = runs.get(c, 0) + 1
    torn = sum(1 for c in runs if runs[c] > 1)
    return torn, len(data)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--json", help="also write the result here")
    ap.add_argument("--rounds", type=int, default=ROUNDS)
    args = ap.parse_args()
    r, w = os.pipe()
    capacity = fcntl.fcntl(w, F_GETPIPE_SZ)
    pipe_buf = os.fpathconf(w, "PC_PIPE_BUF")
    os.close(r)
    os.close(w)
    rows = []
    for size in SIZES:
        torn_writes = 0
        torn_rounds = 0
        for _ in range(args.rounds):
            t, _n = one_round(size, WRITERS)
            torn_writes += t
            torn_rounds += 1 if t else 0
        rows.append({"size": size, "rounds": args.rounds, "writers": WRITERS,
                     "writes": args.rounds * WRITERS,
                     "torn_writes": torn_writes, "torn_rounds": torn_rounds})
        print(f"{size:>7} bytes  {WRITERS} writers x {args.rounds} rounds: "
              f"{torn_writes} of {args.rounds * WRITERS} writes torn, "
              f"in {torn_rounds} rounds", flush=True)
    out = {"measured": date.today().isoformat(), "kernel": platform.release(),
           "machine": platform.machine(), "cpus": os.cpu_count(),
           "pipe_buf": pipe_buf, "capacity": capacity, "rows": rows}
    print(f"kernel {out['kernel']}, PIPE_BUF {pipe_buf}, capacity {capacity}")
    if args.json:
        with open(args.json, "w") as f:
            json.dump(out, f, indent=1)
            f.write("\n")
    return 0


if __name__ == "__main__":
    sys.exit(main())