# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Write skew under snapshot isolation, and the cost of SERIALIZABLE.

on_call(agent, on_shift) with two agents on shift. Two sessions each check
that two are on shift, then each takes a different agent off, and both
commit: under REPEATABLE READ nobody is left on shift; under SERIALIZABLE one
commit fails. Then 50 clients run the same check-and-update against a pool of
agents under SERIALIZABLE with a retry loop, and the retry rate is measured.
Needs the PostgreSQL service from the compose file (port 55803).
"""
import argparse, json, random, threading, time
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:55803/postgres", help="Connection string.")
p.add_argument("--clients", type=int, default=50, help="Concurrent clients in the retry test.")
p.add_argument("--agents", type=int, default=100, help="Agents in the retry test.")
p.add_argument("--per-client", type=int, default=40, help="Transactions per client.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

admin = psycopg.connect(args.dsn, autocommit=True)
def reset(n=2):
    admin.execute("DROP TABLE IF EXISTS on_call"); admin.execute("CREATE TABLE on_call (agent text PRIMARY KEY, on_shift boolean)")
    admin.execute("INSERT INTO on_call SELECT 'a' || g, true FROM generate_series(1, %s) g", (n,))
rows = []
for level in ("REPEATABLE READ", "SERIALIZABLE"):
    reset(); s1, s2 = psycopg.connect(args.dsn), psycopg.connect(args.dsn); outcome = []
    for s in (s1, s2): s.execute(f"BEGIN ISOLATION LEVEL {level}")
    seen = [s.execute("SELECT count(*) FROM on_call WHERE on_shift").fetchone()[0] for s in (s1, s2)]
    s1.execute("UPDATE on_call SET on_shift = false WHERE agent = 'a1'"); s2.execute("UPDATE on_call SET on_shift = false WHERE agent = 'a2'")
    for s in (s1, s2):
        try: s.execute("COMMIT"); outcome.append("committed")
        except psycopg.errors.SerializationFailure: outcome.append("serialization failure"); s.rollback()
    left = admin.execute("SELECT count(*) FROM on_call WHERE on_shift").fetchone()[0]
    rows.append({"isolation": level, "both_saw": seen[0], "session_1": outcome[0], "session_2": outcome[1], "on_shift_after": left})
reset(args.agents); stats = {"commits": 0, "retries": 0}; lock = threading.Lock()
def client(seed):
    rng = random.Random(seed); c = psycopg.connect(args.dsn)
    for _ in range(args.per_client):
        while True:
            try:
                c.execute("BEGIN ISOLATION LEVEL SERIALIZABLE")
                n = c.execute("SELECT count(*) FROM on_call WHERE on_shift").fetchone()[0]
                a = f"a{rng.randrange(1, args.agents + 1)}"
                c.execute("UPDATE on_call SET on_shift = NOT on_shift WHERE agent = %s AND (on_shift = false OR %s > 1)", (a, n))
                c.execute("COMMIT")
                with lock: stats["commits"] += 1
                break
            except psycopg.errors.SerializationFailure:
                c.rollback()
                with lock: stats["retries"] += 1
t = time.perf_counter(); ts = [threading.Thread(target=client, args=(i,)) for i in range(args.clients)]
for th in ts: th.start()
for th in ts: th.join()
el = time.perf_counter() - t
summary = {"clients": args.clients, "commits": stats["commits"], "retries": stats["retries"], "retry_rate_pct": stats["retries"] / (stats["commits"] + stats["retries"]) * 100,
           "commits_per_s": stats["commits"] / el}
for r in rows: print(f"{r['isolation']:<16} both saw {r['both_saw']} on shift; session 1 {r['session_1']}, session 2 {r['session_2']}; on shift after: {r['on_shift_after']}")
print(f"SERIALIZABLE with {args.clients} clients: {stats['commits']:,} commits, {stats['retries']:,} retries ({summary['retry_rate_pct']:.1f}% of attempts), {summary['commits_per_s']:.0f} commits/s")

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