# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Kill a node and measure the recovery.

Three instances of a small service (their own processes) sit behind a
health-checked proxy that sends traffic to the first healthy instance (the
primary) and fails over to the next. Hypothesis, written first: "after the
primary is killed, the error rate returns below 1% within 10 seconds". Abort
condition: errors above 50% for 60 s. A steady open-loop client measures from
the intended send time while the primary is killed with SIGKILL, three times.
Proxy settings mirror common load-balancer defaults (check every 2 s, down
after 3 failures).
"""
import argparse, asyncio, json, os, signal, subprocess, sys, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rate", type=float, default=200, help="Client requests per second.")
p.add_argument("--timeout", type=float, default=2.0, help="Client timeout per request (s).")
p.add_argument("--check-interval", type=float, default=2.0, help="Proxy health-check interval (s).")
p.add_argument("--fall", type=int, default=3, help="Failed checks before the proxy marks an instance down.")
p.add_argument("--hypothesis", type=float, default=10, help="Hypothesised recovery time (s).")
p.add_argument("--repeats", type=int, default=3, help="Kills to run.")
p.add_argument("--port", type=int, default=58200, help="First local port to use.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

BACKEND = r"""
import asyncio, sys
async def h(r, w):
    while await r.readline(): w.write(b"OK\n"); await w.drain()
async def main():
    s = await asyncio.start_server(h, "127.0.0.1", int(sys.argv[1])); print("ready", flush=True)
    async with s: await s.serve_forever()
asyncio.run(main())
"""
PROXY = r"""
import asyncio, sys
port, backends, inter, fall = int(sys.argv[1]), [int(x) for x in sys.argv[2].split(",")], float(sys.argv[3]), int(sys.argv[4])
up = {b: True for b in backends}; fails = {b: 0 for b in backends}
async def ping(b):
    try:
        r, w = await asyncio.wait_for(asyncio.open_connection("127.0.0.1", b), 1)
        w.write(b"PING\n"); await asyncio.wait_for(r.readline(), 1); w.close(); return True
    except Exception: return False
async def checks():
    while True:
        for b in backends:
            if await ping(b): fails[b] = 0; up[b] = True
            else:
                fails[b] += 1
                if fails[b] >= fall: up[b] = False
        await asyncio.sleep(inter)
async def h(r, w):
    while await r.readline():
        target = next((b for b in backends if up[b]), None)
        try:
            br, bw = await asyncio.open_connection("127.0.0.1", target)
            bw.write(b"GET\n"); reply = await br.readline(); bw.close()
            w.write(reply or b"ERR\n")
        except Exception: w.write(b"ERR\n")
        await w.drain()
async def main():
    asyncio.ensure_future(checks())
    s = await asyncio.start_server(h, "127.0.0.1", port); print("ready", flush=True)
    async with s: await s.serve_forever()
asyncio.run(main())
"""

def start(code, *a):
    pr = subprocess.Popen([sys.executable, "-c", code, *map(str, a)], stdout=subprocess.PIPE, text=True)
    pr.stdout.readline(); return pr

async def client(proxy_port, kill_at, total, primary):
    t0 = time.monotonic(); results = []
    async def one(intended):
        ok = False
        try:
            r, w = await asyncio.wait_for(asyncio.open_connection("127.0.0.1", proxy_port), args.timeout)
            w.write(b"GET\n"); ok = (await asyncio.wait_for(r.readline(), max(0.01, intended + args.timeout - time.monotonic()))) == b"OK\n"
            w.close()
        except Exception: ok = False
        results.append((intended - t0, ok, time.monotonic() - intended))
    tasks, killed, i = [], False, 0
    while (intended := t0 + i / args.rate) < t0 + total:
        await asyncio.sleep(max(0, intended - time.monotonic()))
        if not killed and intended - t0 >= kill_at: os.kill(primary.pid, signal.SIGKILL); killed = True
        tasks.append(asyncio.ensure_future(one(intended))); i += 1
    await asyncio.gather(*tasks); return results

def analyse(results, kill_at):
    """Recovery = first second after the kill from which every 1 s window stays below 1% errors."""
    buckets = {}
    for t, ok, _ in results: buckets.setdefault(int(t * 10), []).append(ok)
    def err(start):                                                    # error share in the 1 s window starting at tenth `start`
        xs = [ok for k in range(start, start + 10) for ok in buckets.get(k, [])]
        return 1 - sum(xs) / len(xs) if xs else 0
    last = max(buckets); k0 = int(kill_at * 10); bad = [k for k in range(k0, last - 9) if err(k) >= 0.01]
    recovered = (max(bad) + 10) / 10 - kill_at if bad else 0.0
    peak = max((err(k) for k in range(k0, last - 9)), default=0)
    return recovered, peak, [{"t": k / 10 - kill_at, "error_pct": err(k) * 100} for k in range(0, last - 9, 5)]

rows, timeline = [], []
for rep in range(args.repeats):
    base = args.port + rep * 10; ports = [base + 1, base + 2, base + 3]
    backends = [start(BACKEND, pt) for pt in ports]
    proxy = start(PROXY, base, ",".join(map(str, ports)), args.check_interval, args.fall)
    try:
        res = asyncio.run(client(base, 5, 5 + 30, backends[0]))
    finally:
        for pr in backends + [proxy]: pr.kill(); pr.wait()
    rec, peak, tl = analyse(res, 5)
    aborted = sum(1 for t, ok, _ in res if t > 5 and not ok) / max(1, sum(1 for t, *_ in res if t > 5)) > 0.5
    rows.append({"kill": rep + 1, "recovery_s": rec, "peak_error_pct": peak * 100, "failed_requests": sum(1 for _, ok, _ in res if not ok),
                 "hypothesis": "survived" if rec <= args.hypothesis else "refuted", "aborted": "yes" if aborted else "no"})
    timeline += [{"kill": f"kill {rep + 1}", **x} for x in tl]
    print(f"kill {rep + 1}: errors back below 1% after {rec:.1f} s (peak {peak:.0%}, {rows[-1]['failed_requests']} failed requests)")
rs = [r["recovery_s"] for r in rows]
summary = {"min_s": min(rs), "max_s": max(rs), "hypothesis_s": args.hypothesis, "detect_floor_s": args.check_interval * (args.fall - 1),
           "detect_ceiling_s": args.check_interval * args.fall}
print(f"hypothesis 'below 1% within {args.hypothesis:g} s': recovery ranged {min(rs):.1f}-{max(rs):.1f} s "
      f"(health checks alone need {summary['detect_floor_s']:g}-{summary['detect_ceiling_s']:g} s to mark the primary down)")

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