# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Contended compare-and-swap: n²T/2.

n writers each need one successful conditional write on a shared pointer; a
round takes T = 50 ms and exactly one of the writers trying in a round wins.
Total time and wasted attempts with immediate retry and with jittered
exponential backoff, against n(n − 1)/2 and n²T/2.
"""
import argparse, json, random

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--ns", default="2,5,10,20,40", help="Writer counts.")
p.add_argument("--t", type=float, default=0.05, help="Round time T (s).")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

def drain(n, backoff=None, seed=1):
    rng = random.Random(seed); waiting = list(range(n)); wasted = rounds = 0
    next_try = {w: 0 for w in waiting}; fails = {w: 0 for w in waiting}
    while waiting:
        trying = [w for w in waiting if next_try[w] <= rounds]
        if trying:
            winner = rng.choice(trying); waiting.remove(winner)
            for w in trying:
                if w != winner:
                    wasted += 1; fails[w] += 1
                    next_try[w] = rounds + 1 + (rng.randint(0, 2 ** min(fails[w], 6)) if backoff == "jitter" else 0)
        rounds += 1
    return rounds * args.t, wasted
rows = []
for n in map(int, args.ns.split(",")):
    t0, w0 = drain(n); t1, w1 = drain(n, "jitter")
    rows.append({"n": n, "no_backoff_s": t0, "no_backoff_wasted": w0, "jitter_s": t1, "jitter_wasted": w1, "pairs": n * (n - 1) / 2, "formula_s": n * n * args.t / 2})
for r in rows: print(f"n = {r['n']:>2}: no backoff {r['no_backoff_s']:5.2f} s, {r['no_backoff_wasted']:>4} wasted   jitter {r['jitter_s']:5.2f} s, {r['jitter_wasted']:>4} wasted   n(n-1)/2 = {r['pairs']:.0f}, n²T/2 = {r['formula_s']:.2f} s")

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