# /// script
# requires-python = ">=3.10"
# dependencies = ["confluent-kafka"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Choose the partition count N by measurement.

Measure what one consumer of a simulated workload handles per second,
compute N >= target / per-consumer throughput, multiply by 2-3 and round up
to a highly composite number. Then compare topics with N and 2N partitions:
the time for the consumer group to recover after one consumer is killed, and
the files per hour a sink writes (one file per partition per batch).
Needs the Kafka service from the compose file (port 59092).
"""
import argparse, json, math, multiprocessing as mp, os, signal, time, uuid
from confluent_kafka import Consumer, Producer
from confluent_kafka.admin import AdminClient, NewTopic

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--bootstrap", default="localhost:59092", help="Kafka bootstrap server.")
p.add_argument("--target", type=float, default=20_000, help="Target throughput, records per second.")
p.add_argument("--work-us", type=float, default=200, help="Simulated processing time per record, microseconds.")
p.add_argument("--records", type=int, default=50_000, help="Records for the throughput measurement.")
p.add_argument("--consumers", type=int, default=3, help="Consumers in the group for the rebalance test.")
p.add_argument("--sink-batch-min", type=float, default=5, help="Minutes per sink batch (one file per partition per batch).")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
admin = AdminClient({"bootstrap.servers": args.bootstrap})
HCN = [1, 2, 4, 6, 12, 24, 36, 48, 60, 120, 180, 240, 360, 720]


def topic(parts):
    name = f"n-{uuid.uuid4().hex[:8]}"; admin.create_topics([NewTopic(name, parts, 1)])[name].result(); time.sleep(1); return name

def fill(name, n):
    pr = Producer({"bootstrap.servers": args.bootstrap, "linger.ms": 20})
    for i in range(n): pr.produce(name, key=str(i), value=b"x" * 200); pr.poll(0)
    pr.flush()

def member(name, group, q):
    c = Consumer({"bootstrap.servers": args.bootstrap, "group.id": group, "auto.offset.reset": "earliest",
                  "session.timeout.ms": 6000, "heartbeat.interval.ms": 1000})
    c.subscribe([name], on_assign=lambda c_, parts: q.put((os.getpid(), time.time(), len(parts))))
    while True: c.poll(0.2)


if __name__ == "__main__":
    # 1. per-consumer throughput, with the simulated work
    t1 = topic(1); fill(t1, args.records)
    c = Consumer({"bootstrap.servers": args.bootstrap, "group.id": uuid.uuid4().hex, "auto.offset.reset": "earliest"})
    c.subscribe([t1]); n, start = 0, None
    while n < args.records:
        m = c.poll(1.0)
        if m is None or m.error(): continue
        start = start or time.perf_counter()
        end = time.perf_counter() + args.work_us / 1e6
        while time.perf_counter() < end: pass                 # the work, as a busy wait
        n += 1
    per = n / (time.perf_counter() - start); c.close()
    need = math.ceil(args.target / per)
    N = next(h for h in HCN if h >= 2.5 * need)
    print(f"one consumer: {per:,.0f} records/s; target {args.target:,.0f}/s needs N >= {need}; x2.5 and rounded up to a highly composite number: N = {N}")

    # 2. rebalance time after killing one consumer, N and 2N partitions
    rows = []
    for parts in (N, 2 * N):
        name = topic(parts); group = uuid.uuid4().hex; q = mp.Queue()
        procs = [mp.Process(target=member, args=(name, group, q), daemon=True) for _ in range(args.consumers)]
        for x in procs: x.start()
        time.sleep(12)
        while not q.empty(): q.get()
        killed_at = time.time(); os.kill(procs[0].pid, signal.SIGKILL)
        survivors = {x.pid for x in procs[1:]}; got = {}
        while set(got) != survivors and time.time() - killed_at < 60:
            try: pid, t, k = q.get(timeout=1)
            except Exception: continue
            if pid in survivors: got[pid] = t
        recovery = max(got.values()) - killed_at if got else float("nan")
        files = parts * 60 / args.sink_batch_min
        rows.append({"partitions": parts, "recovery_s": recovery, "files_per_hour": files,
                     "max_consumers": parts, "mb_per_file": args.target * 200 * args.sink_batch_min * 60 / parts / 1e6})
        print(f"{parts:>4} partitions: group back to work {recovery:4.1f} s after a consumer was killed; "
              f"sink writes {files:,.0f} files/hour of {rows[-1]['mb_per_file']:.1f} MB each; up to {parts} consumers")
        for x in procs[1:]: x.kill()

    if args.json:
        with open(args.json, "w") as f:
            json.dump({"tables": {"partitions": rows, "summary": [{"per_consumer": per, "need": need, "N": N}]}}, f)