# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Pessimistic against optimistic: the crossover.

A PostgreSQL counter table of K rows; 16 workers each do read-modify-write on
one random row for 15 seconds, either with SELECT ... FOR UPDATE (pessimistic)
or with a version column and retry when the conditional UPDATE changes 0 rows
(optimistic). For K from 1,000,000 down to 1: attempts per second, successful
commits per second (goodput) and the abort probability.
Needs the PostgreSQL service from the compose file (port 55804).
"""
import argparse, json, random, threading, time
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:55804/postgres", help="Connection string.")
p.add_argument("--keys", default="1000000,10000,100,1", help="Values of K.")
p.add_argument("--workers", type=int, default=16, help="Workers.")
p.add_argument("--seconds", type=float, default=15, help="Seconds per run.")
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 run(K, mode):
    admin.execute("DROP TABLE IF EXISTS counters"); admin.execute("CREATE TABLE counters (id int PRIMARY KEY, n int, version int)")
    admin.execute("INSERT INTO counters SELECT g, 0, 0 FROM generate_series(1, %s) g", (K,)); admin.execute("VACUUM ANALYZE counters")
    stats = {"attempts": 0, "commits": 0}; lock = threading.Lock(); stop = time.time() + args.seconds
    def worker(seed):
        rng = random.Random(seed); c = psycopg.connect(args.dsn); a = ok = 0
        while time.time() < stop:
            k = rng.randrange(1, K + 1); a += 1
            if mode == "pessimistic":
                with c.transaction():
                    n = c.execute("SELECT n FROM counters WHERE id = %s FOR UPDATE", (k,)).fetchone()[0]
                    c.execute("UPDATE counters SET n = %s WHERE id = %s", (n + 1, k))
                ok += 1
            else:
                with c.transaction():
                    n, v = c.execute("SELECT n, version FROM counters WHERE id = %s", (k,)).fetchone()
                    cur = c.execute("UPDATE counters SET n = %s, version = version + 1 WHERE id = %s AND version = %s", (n + 1, k, v))
                    ok += cur.rowcount                                     # 0 rows: someone else won, retry
        with lock: stats["attempts"] += a; stats["commits"] += ok
    ts = [threading.Thread(target=worker, args=(i,)) for i in range(args.workers)]
    for t in ts: t.start()
    for t in ts: t.join()
    return {"K": K, "mode": mode, "attempts_per_s": stats["attempts"] / args.seconds, "goodput_per_s": stats["commits"] / args.seconds,
            "abort_pct": (1 - stats["commits"] / stats["attempts"]) * 100 if stats["attempts"] else None}
rows = [run(int(K), m) for K in args.keys.split(",") for m in ("pessimistic", "optimistic")]
for r in rows: print(f"K = {r['K']:>9,} {r['mode']:<12} attempts {r['attempts_per_s']:7.0f}/s  goodput {r['goodput_per_s']:7.0f}/s  aborts {r['abort_pct']:5.1f}%")

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