# /// 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: one queue, mixed jobs.

A pool of 4 workers and one shared queue: 4 big jobs (5 s each) followed by
200 small ones (1 ms). Then the big jobs on their own pool, and cut into
100 ms chunks. The wait of the small jobs, and the age of the oldest waiting
job, which is the metric that shows the problem.
"""
import argparse, json, queue, threading, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--big-s", type=float, default=5.0, help="Seconds per big job.")
p.add_argument("--small", type=int, default=200, help="Small jobs.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

def pool(n):
    q = queue.Queue(); waits = []; lock = threading.Lock()
    def worker():
        while True:
            t0, cost, small = q.get()
            if small:
                with lock: waits.append(time.time() - t0)
            time.sleep(cost); q.task_done()
    for _ in range(n): threading.Thread(target=worker, daemon=True).start()
    return q, waits

def oldest(q):
    with q.mutex: items = list(q.queue)
    return max((time.time() - t0 for t0, _, _ in items), default=0.0)

rows = []
for label in ("one shared queue", "big jobs on their own pool", "big jobs cut into 100 ms chunks"):
    q, waits = pool(4)
    if label == "one shared queue":
        for _ in range(4): q.put((time.time(), args.big_s, False))
    elif label == "big jobs on their own pool":
        bq, _ = pool(4)
        for _ in range(4): bq.put((time.time(), args.big_s, False))
    else:
        n = int(args.big_s / 0.1)
        for c in range(n):                                   # interleave chunks with the small jobs
            for _ in range(4): q.put((time.time(), 0.1, False))
            if c == 0:
                for _ in range(args.small): q.put((time.time(), 0.001, True))
    if label != "big jobs cut into 100 ms chunks":
        time.sleep(0.05)
        for _ in range(args.small): q.put((time.time(), 0.001, True))
    time.sleep(0.5); age = oldest(q)
    while len(waits) < args.small: time.sleep(0.05)
    rows.append({"setup": label, "max_small_wait_s": max(waits), "median_small_wait_s": sorted(waits)[len(waits) // 2], "oldest_age_at_0_5s": age})
    print(f"{label:<34} small jobs: max wait {max(waits):5.2f} s, median {rows[-1]['median_small_wait_s']:5.3f} s   oldest waiting job after 0.5 s: {age:4.2f} s")

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