# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""A circuit breaker that protects the caller.

A service with a pool of 20 threads calls a dependency at 50 calls/s; the
dependency hangs for 3 s on every call between seconds 10 and 40. Without a
breaker (3 s timeout), and with a rate breaker (1 s timeout, opens at 50%
errors over the last 20 calls, 10 s cooldown, fallback answer while open).
Then, with no outage and 5% ordinary errors, how often a rate breaker and a
count breaker trip wrongly at low and at high traffic. A discrete-event simulation.
"""
import argparse, heapq, json, random

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--threads", type=int, default=20, help="Caller thread pool.")
p.add_argument("--rate", type=float, default=50, help="Incoming calls per second.")
p.add_argument("--seconds", type=float, default=60, help="Length of the run.")
p.add_argument("--outage", type=lambda s: [float(v) for v in s.split(",")], default=[10, 40], help="Outage start,end in seconds.")
p.add_argument("--cooldown", type=float, default=10, help="Seconds the breaker stays open.")
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(breaker, rate, timeout, seconds, outage, error_rate=0.0):
    """breaker: none | rate (>= 50% of the last 20 calls failed) | count (>= 5 failures in the last 10 s)."""
    free, waiting, events, per_sec = args.threads, [], [], {}
    open_until, recent, fail_times, trips = -1.0, [], [], 0
    t = 0.0
    while t < seconds:
        t += random.expovariate(rate); heapq.heappush(events, (t, 0, "arrive", t))
    def sec(x): return per_sec.setdefault(int(x), {"second": int(x), "lat": [], "queue": 0, "fallback": 0})
    def finish(now, born, fallback=False):
        s = sec(now); s["lat"].append(now - born); s["fallback"] += fallback
    def start(now, born):
        nonlocal free
        free -= 1
        hang = outage[0] <= now < outage[1]
        fail = hang or random.random() < error_rate
        dur = timeout if hang else (timeout if fail else random.uniform(0.005, 0.02))
        heapq.heappush(events, (now + dur, 1, "fail" if fail else "ok", born))
    def trip(now):
        nonlocal open_until, trips, recent, fail_times
        open_until, trips, recent, fail_times = now + args.cooldown, trips + 1, [], []
        while waiting: finish(now, waiting.pop(0), True)              # queued calls get the fallback at once
    while events:
        t, _, kind, born = heapq.heappop(events)
        if t > seconds + 30: break
        if kind == "arrive":
            if breaker != "none" and t < open_until: finish(t, born, True); continue
            if free > 0: start(t, born)
            else: waiting.append(born)
        else:
            free += 1
            finish(t, born)
            ok = kind == "ok"
            if breaker == "rate":
                recent = (recent + [ok])[-20:]
                if len(recent) == 20 and recent.count(False) >= 10 and t >= open_until: trip(t)
            elif breaker == "count":
                if not ok: fail_times = [x for x in fail_times + [t] if t - x < 10]
                if len(fail_times) >= 5 and t >= open_until: trip(t)
            if waiting and free > 0: start(t, waiting.pop(0))
        sec(t)["queue"] = max(sec(t)["queue"], len(waiting))
    rows = []
    for s in sorted(per_sec.values(), key=lambda x: x["second"]):
        lat = sorted(s.pop("lat")); s["p99_s"] = lat[max(0, int(len(lat) * 0.99) - 1)] if lat else 0.0; rows.append(s)
    return rows, trips


rows = []
summary = []
for breaker, label, timeout in (("none", "no breaker, 3 s timeout", 3.0), ("rate", "rate breaker, 1 s timeout", 1.0)):
    secs, trips = run(breaker, args.rate, timeout, args.seconds, args.outage)
    for s in secs: s["setup"] = label
    rows += secs
    during = [s for s in secs if args.outage[0] + 2 <= s["second"] < args.outage[1]]
    recovered = next((s["second"] for s in secs if s["second"] >= args.outage[1] and s["p99_s"] < 0.1 and s["queue"] == 0), None)
    summary.append({"setup": label, "max_queue": max(s["queue"] for s in secs), "worst_p99_s": max(s["p99_s"] for s in during),
                    "fallbacks": sum(s["fallback"] for s in during), "recovered_at_s": recovered if recovered is not None else -1})
    x = summary[-1]
    print(f"{label:<26} callers waiting at most {x['max_queue']:>5}   worst p99 during the outage {x['worst_p99_s']:6.2f} s   "
          f"fallbacks {x['fallbacks']:>5}   back to p99 < 0.1 s at second {x['recovered_at_s']}")

wrong = []
for traffic in (1.0, args.rate):
    for breaker in ("rate", "count"):
        _, trips = run(breaker, traffic, 1.0, 600, [1e9, 1e9], error_rate=0.05)
        wrong.append({"breaker": f"{breaker} breaker", "calls_per_s": traffic, "wrong_trips_per_10_min": trips})
        print(f"no outage, 5% ordinary errors, {traffic:>4g} calls/s: {breaker} breaker tripped {trips} times in 10 minutes")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"seconds": rows, "summary": summary, "wrong_trips": wrong}}, f)