# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""kill -9 against power loss: what survives?

1,000 numbered lines written with write() only and with write() + fsync().
The process is killed with SIGKILL at a random moment, again and again; after
each kill the file is checked against what the process had reported written.
(The power-loss half of the lab needs a virtual machine you can switch off;
this script runs the kill half.)
"""
import argparse, json, os, random, signal, subprocess, sys, tempfile, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--lines", type=int, default=1000, help="Lines to write.")
p.add_argument("--kills", type=int, default=50, help="Kills per version.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--worker", help=argparse.SUPPRESS)
p.add_argument("--out", help=argparse.SUPPRESS)
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

if args.worker:
    fd = os.open(args.out, os.O_WRONLY | os.O_CREAT | os.O_TRUNC)
    for i in range(args.lines):
        os.write(fd, f"{i}\n".encode())
        if args.worker == "fsync": os.fsync(fd)
        print(i, flush=True)                              # "line i is written": what the process believes
        time.sleep(0.001)
    sys.exit(0)

random.seed(args.seed)
rows = []
for mode in ("write only", "write + fsync"):
    lost_total = 0
    for _ in range(args.kills):
        out = os.path.join(tempfile.mkdtemp(), "lines.txt")
        proc = subprocess.Popen([sys.executable, __file__, "--worker", "fsync" if "fsync" in mode else "plain", "--out", out,
                                 "--lines", str(args.lines)], stdout=subprocess.PIPE, text=True)
        time.sleep(random.uniform(0.1, 1.0)); os.kill(proc.pid, signal.SIGKILL); proc.wait()
        acked = [int(x) for x in proc.stdout.read().split()]
        on_disk = set(int(x) for x in open(out).read().split())
        lost_total += sum(1 for i in acked if i not in on_disk)
    rows.append({"version": mode, "kills": args.kills, "acknowledged_lines_lost": lost_total})
    print(f"{mode:<14} {args.kills} kills: {lost_total} acknowledged lines missing after the kill")

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