# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Anomaly rate equals window times rate.

A user writes a row at a steady rate lambda on the primary and keeps reading
it from a replica that lags by W. The share of stale reads should be about
lambda x W, as long as that is well below 1. Then the user's reads go to the
primary for 2 s after each write.
Needs the primary and replica from the compose file (ports 55472, 55473).
"""
import argparse, json, random, threading, time
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--primary", default="postgresql://postgres:lab@localhost:55472/postgres", help="Primary connection string.")
p.add_argument("--replica", default="postgresql://postgres:lab@localhost:55473/postgres", help="Replica connection string.")
p.add_argument("--rate", type=float, default=0.2, help="lambda: writes per second (Poisson).")
p.add_argument("--lags-ms", type=lambda s: [int(v) for v in s.split(",")], default=[20, 200, 2000],
               help="Replica lags W to try (recovery_min_apply_delay), comma-separated.")
p.add_argument("--seconds", type=float, default=60, help="Seconds of traffic per lag.")
p.add_argument("--read-every", type=float, default=0.01, help="Seconds between the user's reads.")
p.add_argument("--sticky", type=float, default=2.0, help="Seconds the user's reads go to the primary after a write, in the fix.")
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)

pri = psycopg.connect(args.primary, autocommit=True)
rep = psycopg.connect(args.replica, autocommit=True)
pri.execute("DROP TABLE IF EXISTS cart"); pri.execute("CREATE TABLE cart (id int PRIMARY KEY, n int)")
pri.execute("INSERT INTO cart VALUES (1, 0)")


def phase(lag_ms, sticky):
    rep.execute(f"ALTER SYSTEM SET recovery_min_apply_delay = '{lag_ms}ms'"); rep.execute("SELECT pg_reload_conf()")
    time.sleep(lag_ms / 1000 + 1)
    w = psycopg.connect(args.primary, autocommit=True)
    state = {"n": pri.execute("SELECT n FROM cart").fetchone()[0], "last_write": 0.0}
    stop = threading.Event()
    def writer():                                        # Poisson writes at rate lambda
        while not stop.wait(random.expovariate(args.rate)):
            state["n"] = w.execute("UPDATE cart SET n = n + 1 WHERE id = 1 RETURNING n").fetchone()[0]
            state["last_write"] = time.time()
    th = threading.Thread(target=writer); th.start()
    reads = stale = 0; t0 = time.time()
    while time.time() - t0 < args.seconds:
        expected = state["n"]
        use_primary = sticky and time.time() - state["last_write"] < args.sticky
        got = (pri if use_primary else rep).execute("SELECT n FROM cart").fetchone()[0]
        reads += 1; stale += got < expected
        time.sleep(args.read_every)
    stop.set(); th.join(); w.close()
    return {"lag_ms": lag_ms, "lambda_W": args.rate * lag_ms / 1000, "reads": reads, "stale": stale,
            "stale_share": stale / reads, "stale_per_s": stale / args.seconds, "sticky": "yes" if sticky else "no"}


rows = []
for lag in args.lags_ms:
    r = phase(lag, False); rows.append(r)
    print(f"W = {lag:>5} ms  lambda*W = {r['lambda_W']:.3f}   stale reads {r['stale']:>5} of {r['reads']} "
          f"= {r['stale_share']:.3f} of reads ({r['stale_per_s']:.2f}/s)")
fix = phase(args.lags_ms[-1], True)
print(f"W = {args.lags_ms[-1]} ms, reads to the primary for {args.sticky:g} s after each write: "
      f"{fix['stale']} stale of {fix['reads']}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"lags": rows, "summary": [{"sticky_stale": fix["stale"], "sticky_reads": fix["reads"],
                   "rate": args.rate, "max_lag_ms": args.lags_ms[-1]}]}}, f)