# /// script
# requires-python = ">=3.10"
# dependencies = ["kafka-python"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Auto-commit against commit-after-process.

10,000 numbered records in a one-partition topic. A consumer "processes"
each by appending its number to a file after a 5 ms sleep. Once with
auto-commit every second, once with a manual commit after every 100 records.
Each run is killed with SIGKILL at a random moment, restarted, and finished;
then the file is checked for missing and duplicated numbers.
Needs the Kafka service from the compose file (port 59092).
"""
import argparse, collections, json, os, random, signal, subprocess, sys, tempfile, time, uuid

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--bootstrap", default="localhost:59092", help="Kafka bootstrap server.")
p.add_argument("--records", type=int, default=10_000, help="Records.")
p.add_argument("--kill-between", type=lambda s: [float(v) for v in s.split(",")], default=[10, 30],
               help="Kill the consumer after a random number of seconds in this range.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--worker", help=argparse.SUPPRESS)
p.add_argument("--topic", 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:                                         # the consumer process
    from kafka import KafkaConsumer
    auto = args.worker == "auto"
    c = KafkaConsumer(args.topic, bootstrap_servers=args.bootstrap, group_id="lab-" + args.topic, enable_auto_commit=auto,
                      auto_commit_interval_ms=1000, auto_offset_reset="earliest", max_poll_records=100,
                      session_timeout_ms=6000, heartbeat_interval_ms=1000,
                      consumer_timeout_ms=30000)          # waits out the dead member's session before giving up
    out = open(args.out, "a")
    for msg in c:
        time.sleep(0.005)
        out.write(msg.value.decode() + "\n"); out.flush()
        if not auto and msg.offset % 100 == 99: c.commit()
    if not auto: c.commit()
    sys.exit(0)

from kafka import KafkaProducer
from kafka.admin import KafkaAdminClient, NewTopic
random.seed(args.seed)
admin = KafkaAdminClient(bootstrap_servers=args.bootstrap)
rows = []
for mode, label in (("auto", "auto-commit every 1 s"), ("manual", "commit after every 100 processed")):
    topic = f"numbers-{uuid.uuid4().hex[:8]}"
    admin.create_topics([NewTopic(topic, 1, 1)]); time.sleep(1)
    pr = KafkaProducer(bootstrap_servers=args.bootstrap)
    for i in range(args.records): pr.send(topic, str(i).encode())
    pr.flush()
    out = os.path.join(tempfile.mkdtemp(), "processed.txt")
    cmd = [sys.executable, __file__, "--worker", mode, "--topic", topic, "--out", out, "--bootstrap", args.bootstrap]
    proc = subprocess.Popen(cmd); time.sleep(random.uniform(*args.kill_between))
    os.kill(proc.pid, signal.SIGKILL); proc.wait()
    subprocess.run(cmd, check=True)                     # restart and finish
    nums = [int(x) for x in open(out).read().split()]
    cnt = collections.Counter(nums)
    missing = args.records - len(cnt); dups = sum(v - 1 for v in cnt.values() if v > 1)
    rows.append({"commit": label, "missing": missing, "duplicated": dups})
    print(f"{label:<34} killed once: {missing:>4} numbers missing, {dups:>4} processed twice")

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