# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Bulkheads: one slow dependency, one shared pool.

A service with 10 threads calls dependency B and dependency C (always fast).
B hangs from second 5. With one shared pool, B's hung calls take every
thread and C's calls fail too; with two pools of 5, C stays healthy.
A discrete-event simulation.
"""
import argparse, heapq, json, random

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--threads", type=int, default=10, help="Threads in total.")
p.add_argument("--rate", type=float, default=40, help="Calls per second to each dependency.")
p.add_argument("--hang-at", type=float, default=5, help="Second at which B starts hanging.")
p.add_argument("--seconds", type=float, default=20, help="Length of the run.")
p.add_argument("--queue-timeout", type=float, default=0.5, help="A call waiting longer than this for a thread fails.")
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)


def run(split):
    pools = {"B": args.threads // 2, "C": args.threads // 2} if split else {"all": args.threads}
    pool_of = (lambda d: d) if split else (lambda d: "all")
    free = dict(pools); waiting = {k: [] for k in pools}
    events, res = [], {"B": [0, 0], "C": [0, 0]}                  # [ok, failed] after the hang
    busy_samples = []
    for dep in ("B", "C"):
        t = 0.0
        while t < args.seconds:
            t += random.expovariate(args.rate); heapq.heappush(events, (t, "arrive", dep, t))
    def start(t, dep, born):
        k = pool_of(dep); free[k] -= 1
        hang = dep == "B" and t >= args.hang_at
        heapq.heappush(events, (t + (1e9 if hang else random.uniform(0.01, 0.05)), "done", dep, born))
    while events:
        t, kind, dep, born = heapq.heappop(events)
        if t > args.seconds: break
        k = pool_of(dep)
        if kind == "arrive":
            if free[k] > 0: start(t, dep, born)
            else: waiting[k].append((dep, born))
        else:
            if born >= args.hang_at: res[dep][0] += 1
            free[k] += 1
            while waiting[k]:
                d2, b2 = waiting[k].pop(0)
                if t - b2 > args.queue_timeout:
                    if b2 >= args.hang_at: res[d2][1] += 1
                    continue
                start(t, d2, b2); break
        if t < args.hang_at: busy_samples.append(sum(pools[x] - free[x] for x in pools) / args.threads)
    for k in waiting:                                              # whatever is still queued at the end failed
        for d2, b2 in waiting[k]:
            if b2 >= args.hang_at: res[d2][1] += 1
    c_ok, c_fail = res["C"]
    return {"design": "two pools of 5" if split else "one shared pool of 10",
            "c_success_pct": c_ok / max(1, c_ok + c_fail) * 100, "idle_before_pct": (1 - sum(busy_samples) / len(busy_samples)) * 100}


rows = [run(False), run(True)]
for r in rows:
    print(f"{r['design']:<24} C calls succeeding after B hangs: {r['c_success_pct']:5.1f}%   idle threads before: {r['idle_before_pct']:4.1f}%")

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