# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Reproduce a deadlock, then remove it.

Two PostgreSQL sessions move 30 between anna and ben in opposite orders and
deadlock; one is aborted after deadlock_timeout. Then both lock the two rows
in a fixed order first (SELECT ... ORDER BY owner FOR UPDATE) and no deadlock
occurs. Repeated 20 times each, and with deadlock_timeout = 100ms, measuring
how long the victim waited.
Needs the PostgreSQL service from the compose file (port 55801).
"""
import argparse, json, threading, time
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:55801/postgres", help="Connection string.")
p.add_argument("--trials", type=int, default=20, help="Trials per variant.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

setup = psycopg.connect(args.dsn, autocommit=True)
setup.execute("DROP TABLE IF EXISTS wallets"); setup.execute("CREATE TABLE wallets (owner text PRIMARY KEY, balance int)")
setup.execute("INSERT INTO wallets VALUES ('anna', 1000), ('ben', 1000)")
def session(first, second, ordered, timeout, out, barrier):
    c = psycopg.connect(args.dsn); c.execute(f"SET deadlock_timeout = '{timeout}'")
    t = time.perf_counter()
    try:
        with c.transaction():
            if ordered: c.execute("SELECT * FROM wallets WHERE owner IN ('anna', 'ben') ORDER BY owner FOR UPDATE")
            c.execute("UPDATE wallets SET balance = balance - 30 WHERE owner = %s", (first,))
            if barrier: barrier.wait()                                      # both hold their first row before taking the second
            time.sleep(0.05)
            c.execute("UPDATE wallets SET balance = balance + 30 WHERE owner = %s", (second,))
        out.append(("ok", time.perf_counter() - t))
    except psycopg.errors.DeadlockDetected:
        out.append(("deadlock", time.perf_counter() - t))
    finally: c.close()
rows = []
for ordered, timeout in ((False, "1s"), (False, "100ms"), (True, "1s")):
    deadlocks, waits = 0, []
    for _ in range(args.trials):
        out = []; barrier = None if ordered else threading.Barrier(2, timeout=10)
        ts = [threading.Thread(target=session, args=("anna", "ben", ordered, timeout, out, barrier)),
              threading.Thread(target=session, args=("ben", "anna", ordered, timeout, out, barrier))]
        for t in ts: t.start()
        for t in ts: t.join()
        victims = [w for r, w in out if r == "deadlock"]; deadlocks += bool(victims); waits += victims
    total = setup.execute("SELECT sum(balance) FROM wallets").fetchone()[0]
    rows.append({"locking": "fixed order (FOR UPDATE first)" if ordered else "opposite orders", "deadlock_timeout": timeout,
                 "trials": args.trials, "deadlocks": deadlocks, "victim_wait_ms": (sum(waits) / len(waits) * 1000) if waits else None, "total_balance": total})
for r in rows: print(f"{r['locking']:<32} deadlock_timeout {r['deadlock_timeout']:<6} deadlocks {r['deadlocks']:>2}/{r['trials']}" +
                      (f"   victim waited {r['victim_wait_ms']:.0f} ms" if r["victim_wait_ms"] else "") + f"   total balance {r['total_balance']}")

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