# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Fencing tokens versus a paused leader.

Leader A holds epoch 1 and pauses (a GC pause) past its lease; leader B gets
epoch 2 and writes; A wakes up and writes. The seat store accepts any write,
or rejects writes with an epoch lower than the highest it has seen. Then 1,000
runs with random pauses.
"""
import argparse, json, random, threading, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--pause", type=float, default=3.0, help="Seconds leader A is paused.")
p.add_argument("--b-at", type=float, default=1.2, help="Seconds after start that leader B writes.")
p.add_argument("--runs", type=int, default=1000, help="Runs with random pauses (scaled 100x faster).")
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 once(fenced, pause_a, b_at):
    store, state, lock, out = {}, {"highest": 0}, threading.Lock(), []
    def write(owner, epoch):
        with lock:
            if fenced and epoch < state["highest"]: return "REJECTED"
            state["highest"] = max(state["highest"], epoch); store["A12"] = (owner, epoch); return "OK"
    def a(): time.sleep(pause_a); out.append(("A", write("A", 1)))
    def b(): time.sleep(b_at); out.append(("B", write("B", 2)))
    ts = [threading.Thread(target=a), threading.Thread(target=b)]
    [t.start() for t in ts]; [t.join() for t in ts]
    return out, store["A12"][0]


rows = []
for pause in (args.pause, 0.5):
    for fenced in (False, True):
        out, owner = once(fenced, pause, args.b_at)
        rows.append({"pause_s": pause, "fenced": "yes" if fenced else "no", "writes": ", ".join(f"{w} {r}" for w, r in out), "final_owner": owner})
        print(f"A paused {pause:g} s, fenced = {fenced!s:<5}  {rows[-1]['writes']:<24} final owner {owner}")

wrong = {False: 0, True: 0}
for _ in range(args.runs):
    pause, b_at = random.uniform(0, 4) / 100, random.uniform(0, 4) / 100      # the same timings, 100x faster
    for fenced in (False, True):
        _, owner = once(fenced, pause, b_at)
        wrong[fenced] += owner == "A" and pause > b_at                        # A wrote after B: a stale leader won
print(f"{args.runs} random runs: stale leader's write kept in {wrong[False]} unfenced, {wrong[True]} fenced")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"cases": rows, "summary": [{"runs": args.runs, "wrong_unfenced": wrong[False], "wrong_fenced": wrong[True]}]}}, f)