# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Head-of-line blocking: a log against a queue.

12 partitions of 100 records/s each and one poison record in partition 3.
Consumer lag under three policies with one consumer per partition, then with
one consumer serving all 12 that blocks, then with a shared queue.
"""
import argparse, json

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--partitions", type=int, default=12, help="Partitions.")
p.add_argument("--rate", type=int, default=100, help="Records per second per partition.")
p.add_argument("--seconds", type=int, default=3600, help="Simulated seconds.")
p.add_argument("--poison-at", type=int, default=60, help="Second at which the poison record arrives.")
p.add_argument("--blocked", type=int, default=3, help="Partition with the poison record.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
P, R, B = args.partitions, args.rate, args.blocked


def run(policy):
    lag = [0] * P; series = []; skipped = 0
    shared_capacity = R * P                               # one consumer for all partitions: same total capacity
    for sec in range(args.seconds):
        stuck = sec >= args.poison_at
        for q in range(P): lag[q] += R
        if policy == "block and retry (consumer per partition)":
            for q in range(P):
                if not (q == B and stuck): lag[q] = 0
        elif policy == "skip the poison record":
            if stuck and sec == args.poison_at: skipped += 1
            lag = [0] * P
        elif policy == "pause that partition only":
            for q in range(P):
                if not (q == B and stuck): lag[q] = 0     # the paused partition waits; the rest flow
        elif policy == "one consumer for all, blocks":
            if not stuck: lag = [0] * P                     # stuck on the poison record: nothing moves
        elif policy == "shared queue":
            lag = [0] * P                                   # any worker takes the next message; the poison is set aside
        if sec % 60 == 59:
            series.append({"minute": (sec + 1) // 60, "policy": policy, "blocked_lag": lag[B],
                           "max_other_lag": max(lag[q] for q in range(P) if q != B)})
    return series


rows = []
for pol in ("block and retry (consumer per partition)", "pause that partition only", "skip the poison record",
            "one consumer for all, blocks", "shared queue"):
    s = run(pol); rows += s
    print(f"{pol:<42} after {args.seconds // 60} min: blocked partition lag {s[-1]['blocked_lag']:>9,}   "
          f"other partitions max {s[-1]['max_other_lag']:>9,}")

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