# /// script
# requires-python = ">=3.10"
# dependencies = ["confluent-kafka"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Break per-key order in a log, one clause at a time.

10,000 events {order_id, seq} for 100 orders, keyed by order_id, into a
3-partition topic; a checker verifies that seq increases per order_id. Then:
(1) a thread pool of 8 workers inside the consumer; (2) a producer with
retries, 5 requests in flight and idempotence off, behind a proxy that drops
5% of produce responses; (3) the topic grown to 6 partitions mid-run. Then
the repairs for (1) and (2).
Needs the Kafka service from the compose file (ports 59092 and 59094).
"""
import argparse, json, random, socket, struct, threading, time, uuid
from concurrent.futures import ThreadPoolExecutor
from confluent_kafka import Consumer, Producer
from confluent_kafka.admin import AdminClient, NewPartitions, NewTopic

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--bootstrap", default="localhost:59092", help="Kafka bootstrap server.")
p.add_argument("--events", type=int, default=10_000, help="Events produced per experiment.")
p.add_argument("--orders", type=int, default=100, help="Distinct order ids (keys).")
p.add_argument("--drop", type=float, default=0.05, help="Share of produce responses the proxy drops in experiment 2.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
random.seed(args.seed)
admin = AdminClient({"bootstrap.servers": args.bootstrap})


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

def produce(name, conf=None, grow_at=None):
    pr = Producer({"bootstrap.servers": args.bootstrap, "linger.ms": 5, "enable.idempotence": True, "batch.num.messages": 20,
                   "topic.metadata.refresh.interval.ms": 1000, **(conf or {})})
    seq = {}
    for i in range(args.events):
        o = random.randrange(args.orders); seq[o] = seq.get(o, 0) + 1
        pr.produce(name, key=str(o), value=json.dumps({"order_id": o, "seq": seq[o]}))
        pr.poll(0)
        if grow_at and i == grow_at:
            pr.flush(); admin.create_partitions([NewPartitions(name, 6)])[name].result(); time.sleep(2)
            pr.list_topics(name, timeout=10)                 # the producer learns the new partition count
    pr.flush(60)

def consume(name, workers=1, by_key=False):
    c = Consumer({"bootstrap.servers": args.bootstrap, "group.id": uuid.uuid4().hex, "auto.offset.reset": "earliest",
                  "topic.metadata.refresh.interval.ms": 1000,
                  "max.partition.fetch.bytes": 2048, "fetch.max.bytes": 8192, "message.max.bytes": 4096})       # small fetches interleave the partitions, as a lagging consumer's do
    c.subscribe([name])
    applied, lock, got = [], threading.Lock(), 0
    def work(v):
        time.sleep(random.random() / 1000)               # a little uneven processing time
        with lock: applied.append(v)
    pools = [ThreadPoolExecutor(1) for _ in range(workers)] if by_key else None
    pool = ThreadPoolExecutor(workers) if workers > 1 and not by_key else None
    idle = 0
    while idle < 10:
        m = c.poll(0.5)
        if m is None or m.error(): idle += 1; continue
        idle = 0; got += 1
        v = json.loads(m.value())
        if by_key: pools[v["order_id"] % workers].submit(work, v)              # one worker per key: order kept
        elif pool: pool.submit(work, v)
        else: work(v)
    for x in (pools or []) + ([pool] if pool else []): x.shutdown(wait=True)
    c.close()
    last, bad = {}, 0
    for v in applied:
        if v["seq"] < last.get(v["order_id"], 0): bad += 1
        last[v["order_id"]] = max(last.get(v["order_id"], 0), v["seq"])
    return bad, got


class DroppingProxy:
    """Forwards Kafka traffic; drops a share of Produce responses (api key 0), matched by correlation id."""
    def __init__(self, listen, target, drop):
        self.listen, self.target, self.drop, self.dropped = listen, target, drop, 0
    def start(self):
        s = socket.socket(); s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1); s.bind(("127.0.0.1", self.listen)); s.listen(16)
        self.s = s; threading.Thread(target=self.loop, daemon=True).start()
    def loop(self):
        while True:
            try: c, _ = self.s.accept()
            except OSError: return
            u = socket.create_connection(("127.0.0.1", self.target))
            produce_ids = set()
            threading.Thread(target=self.pipe, args=(c, u, produce_ids, True), daemon=True).start()
            threading.Thread(target=self.pipe, args=(u, c, produce_ids, False), daemon=True).start()
    def pipe(self, a, b, produce_ids, requests):
        f = a.makefile("rb")
        try:
            while True:
                head = f.read(4)
                if len(head) < 4: break
                body = f.read(struct.unpack(">i", head)[0])
                if requests:
                    api_key, _, corr = struct.unpack(">hhi", body[:8])
                    if api_key == 0: produce_ids.add(corr)
                else:
                    corr = struct.unpack(">i", body[:4])[0]
                    if corr in produce_ids:
                        produce_ids.discard(corr)
                        if random.random() < self.drop: self.dropped += 1; continue      # the ack is lost
                b.sendall(head + body)
        except OSError: pass
        finally:
            for x in (a, b):
                try: x.close()
                except OSError: pass


rows = []
def record(name, bad, got, note=""):
    rows.append({"experiment": name, "violations": bad, "consumed": got, "note": note})
    print(f"{name:<52} {bad:>5} out-of-order events of {got:,} consumed {note}")

t = topic(3); produce(t); record("baseline: 1 consumer thread, idempotent producer", *consume(t))
record("(1) consumer with a pool of 8 workers", *consume(t, workers=8))
record("    repair: one worker per key (hash to a thread)", *consume(t, workers=8, by_key=True))

proxy = DroppingProxy(59093, 59094, args.drop); proxy.start()
lossy = {"bootstrap.servers": "localhost:59093", "enable.idempotence": False, "max.in.flight.requests.per.connection": 5,
         "retries": 10, "request.timeout.ms": 2000, "socket.timeout.ms": 2000, "message.timeout.ms": 120000,
         "retry.backoff.ms": 50, "linger.ms": 1, "batch.num.messages": 50}
t = topic(3); produce(t, lossy); bad, got = consume(t)
record("(2) producer: retries, 5 in flight, no idempotence", bad, got, f"({got - args.events:+,} vs produced; {proxy.dropped} acks dropped)")
proxy.dropped = 0
t = topic(3); produce(t, {**lossy, "enable.idempotence": True}); bad, got = consume(t)
record("    repair: enable.idempotence=true", bad, got, f"({got - args.events:+,} vs produced; {proxy.dropped} acks dropped)")

t = topic(3); produce(t, grow_at=args.events // 2); record("(3) topic grown from 3 to 6 partitions mid-run", *consume(t))

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