# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Ack before or after processing: loss or duplicates.

A consumer reads numbers 0..9999, "processes" each by appending it to a
file, and stores its offset. Version A saves the offset first; version B
processes first. The process is killed with SIGKILL at random moments and
restarted until it finishes; then the output is checked for missing and
duplicated numbers.
"""
import argparse, collections, json, os, random, signal, subprocess, sys, tempfile, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--records", type=int, default=10_000, help="Records to process.")
p.add_argument("--max-kill-delay", type=float, default=0.3, help="Largest random delay before each kill, in seconds.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--worker", help=argparse.SUPPRESS)
p.add_argument("--dir", help=argparse.SUPPRESS)
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

if args.worker:                                            # the consumer process
    os.chdir(args.dir)
    def load():
        try: return int(open("offset.txt").read())
        except Exception: return 0
    def save(o):
        open("offset.txt.tmp", "w").write(str(o)); os.replace("offset.txt.tmp", "offset.txt")
    out = open("out.txt", "a")
    for n in range(load(), args.records):
        if args.worker == "ack_first":
            save(n + 1); out.write(f"{n}\n"); out.flush()
        else:
            out.write(f"{n}\n"); out.flush(); save(n + 1)
        time.sleep(random.random() / 2000)
    sys.exit(0)

random.seed(args.seed)
rows = []
for mode, label in (("ack_first", "A: save the offset, then process"), ("process_first", "B: process, then save the offset")):
    d = tempfile.mkdtemp(prefix=f"ack-{mode}-")
    kills = 0
    while True:
        proc = subprocess.Popen([sys.executable, __file__, "--worker", mode, "--dir", d, "--records", str(args.records)])
        try:
            proc.wait(timeout=random.uniform(0.05, args.max_kill_delay)); break     # finished
        except subprocess.TimeoutExpired:
            os.kill(proc.pid, signal.SIGKILL); proc.wait(); kills += 1
    nums = [int(x) for x in open(os.path.join(d, "out.txt")).read().split()]
    counts = collections.Counter(nums)
    missing = args.records - len(counts); dups = sum(c - 1 for c in counts.values() if c > 1)
    rows.append({"version": label, "kills": kills, "missing": missing, "duplicated": dups})
    print(f"{label:<36} killed {kills:>3} times: {missing:>3} numbers missing, {dups:>3} duplicated")

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