# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Gray failure: a dead node against a slow one.

A query fans out to 200 search nodes and waits for all. Case A: one node is
dead and skipped after a health check. Case B: one node is 100 times slower.
Case C: peer-relative ejection (p99 above 3 x the peers' median); then all
nodes 5 times slower; then an eviction cap when a deep check fails everywhere.
"""
import argparse, json
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--nodes", type=int, default=200, help="Nodes.")
p.add_argument("--queries", type=int, default=10_000, help="Queries.")
p.add_argument("--cap", type=float, default=0.20, help="Largest share of nodes that may be evicted.")
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()
rng = np.random.default_rng(args.seed); N, Q = args.nodes, args.queries

def latencies(slow_factor=1.0, slow_node=None, scale=1.0):
    lat = rng.normal(10, 1, (Q, N)).clip(5) * scale
    if slow_node is not None: lat[:, slow_node] *= slow_factor
    return lat

def eject(lat):
    p99 = np.percentile(lat, 99, axis=0); med = np.median(p99)
    return np.where(p99 > 3 * med)[0]

rows = []
lat = latencies(); lat = np.delete(lat, 0, axis=1)                                   # A: dead node skipped
q = lat.max(axis=1); rows.append({"case": "A: one node dead, skipped", "median_ms": float(np.median(q)), "slow_pct": float(np.mean(q > 50) * 100)})
lat = latencies(100, 0); q = lat.max(axis=1)                                         # B: one node 100x slower
rows.append({"case": "B: one node 100x slower", "median_ms": float(np.median(q)), "slow_pct": float(np.mean(q > 50) * 100)})
ej = eject(lat); q = np.delete(lat, ej, axis=1).max(axis=1)                          # C: peer-relative ejection
rows.append({"case": f"C: peer ejection ({len(ej)} node ejected)", "median_ms": float(np.median(q)), "slow_pct": float(np.mean(q > 50) * 100)})
lat = latencies(scale=5); ej2 = eject(lat); q = np.delete(lat, ej2, axis=1).max(axis=1)
rows.append({"case": f"all nodes 5x slower: peer ejection fires on {len(ej2)}", "median_ms": float(np.median(q)), "slow_pct": float(np.mean(q > 50) * 100)})
failing = N                                                                         # a deep check (the database) fails everywhere
evicted = min(failing, int(args.cap * N)) if failing <= args.cap * N else 0
rows.append({"case": f"deep check fails on all {N}: evicted with a {args.cap:.0%} cap", "median_ms": float("nan"), "slow_pct": float("nan"), "evicted": evicted})
for r in rows:
    if "evicted" in r: print(f"{r['case']:<52} evicted {r['evicted']}")
    else: print(f"{r['case']:<52} median query {r['median_ms']:7.1f} ms   queries over 50 ms {r['slow_pct']:5.1f}%")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"cases": [r for r in rows if "evicted" not in r], "summary": [{"evicted": rows[-1]["evicted"], "all_slow_ejected": len(ej2)}]}}, f)