#!/usr/bin/env python3
"""
Stress reproducer: a cached generic plan stays stale after ATTACH PARTITION
to a sub-partitioned table commits while the plan is being built.

Requires only psycopg 3 (pip install "psycopg[binary]") and an UNMODIFIED
PostgreSQL server.  Connection parameters come from --dsn or the usual libpq
environment variables (PGHOST, PGPORT, PGDATABASE, PGUSER, ...).

Layout (created in its own schema, dropped at the end unless --keep):
    root   LIST(grp)  -> p_a (100 rows)
                      -> mid RANGE(k) -> mid_0 .. mid_<N-1>   (1 row each)
    float_sub (100 rows) is attached to / detached from mid by the driver.
Correct answers for  SELECT count(*) FROM root WHERE v >= 0  are therefore
100 + N (float_sub detached) or 200 + N (attached), nothing else.

Each round:
  1. every worker session loops   DISCARD PLANS; EXECUTE q(0)
  2. the driver sleeps a random fraction of one iteration, then runs
     ALTER TABLE mid ATTACH PARTITION float_sub ...  (autocommit)
  3. as soon as that commits, workers stop after their current iteration
  4. quiescent check: each worker runs, 3 times, each in its own
     transaction,  EXECUTE q(0)  and the same query unprepared
  5. a worker whose 3 EXECUTEs ALL differ from the unprepared query is a hit
  6. the driver DETACHes float_sub (workers idle) and the next round starts

Server settings: each worker locks every partition of mid in one transaction,
so the shared lock table must hold about sessions * partitions entries
(max_locks_per_transaction * (max_connections + max_prepared_transactions)).
The script checks this before starting.

Example:
    PGHOST=/tmp PGPORT=5432 python3 plancache_stress.py --sessions 8 --minutes 5
"""
import argparse
import json
import random
import sys
import threading
import time

import psycopg

QUERY = "SELECT count(*) FROM {s}.root WHERE v >= $1"
PLAIN = "SELECT count(*) FROM {s}.root WHERE v >= 0"


def connect(dsn, app):
    conn = psycopg.connect(dsn, autocommit=True, application_name=app)
    # psycopg's own statement cache sends DEALLOCATE ALL after DROP/ALTER/
    # ROLLBACK, which would also drop our SQL-level prepared statement.
    conn.prepare_threshold = None
    return conn


def one(conn, sql, params=None):
    return conn.execute(sql, params, prepare=False).fetchone()[0]


def setup(conn, s, nparts):
    conn.execute(f"DROP SCHEMA IF EXISTS {s} CASCADE")
    conn.execute(f"CREATE SCHEMA {s}")
    conn.execute(f"CREATE TABLE {s}.root (grp int, k int, v int) PARTITION BY LIST (grp)")
    conn.execute(f"CREATE TABLE {s}.p_a PARTITION OF {s}.root FOR VALUES IN (1)")
    conn.execute(f"CREATE TABLE {s}.mid PARTITION OF {s}.root FOR VALUES IN (2) PARTITION BY RANGE (k)")
    # one statement per partition (autocommit) keeps setup's lock usage small
    for i in range(nparts):
        conn.execute(f"CREATE TABLE {s}.mid_{i} PARTITION OF {s}.mid "
                     f"FOR VALUES FROM ({i * 10}) TO ({i * 10 + 10})")
    conn.execute(f"INSERT INTO {s}.p_a SELECT 1, i, i FROM generate_series(1, 100) i")
    conn.execute(f"INSERT INTO {s}.mid SELECT 2, i * 10, i FROM generate_series(0, {nparts - 1}) i")
    conn.execute(f"CREATE TABLE {s}.float_sub (grp int, k int, v int)")
    conn.execute(f"INSERT INTO {s}.float_sub SELECT 2, 100000 + i, i FROM generate_series(1, 100) i")
    # lets ATTACH skip the validation scan
    conn.execute(f"ALTER TABLE {s}.float_sub ADD CONSTRAINT float_sub_chk "
                 f"CHECK (grp = 2 AND k >= 100000 AND k < 200000)")
    conn.execute(f"ANALYZE {s}.root, {s}.float_sub")


def teardown(conn, s, nparts):
    # drop partitions one by one so a small lock table does not overflow
    for i in range(nparts):
        conn.execute(f"DROP TABLE IF EXISTS {s}.mid_{i}")
    conn.execute(f"DROP SCHEMA IF EXISTS {s} CASCADE")


class Worker:
    def __init__(self, dsn, idx, s, mode):
        self.conn = connect(dsn, f"plancache_stress_A{idx}")
        self.pid = one(self.conn, "SELECT pg_backend_pid()")
        self.conn.execute(f"SET plan_cache_mode = {mode}")
        self.conn.execute(f"PREPARE q(int) AS {QUERY.format(s=s)}")
        self.go = threading.Event()
        self.stop = threading.Event()
        self.idle = threading.Event()
        self.iterations = 0
        self.error = None
        threading.Thread(target=self._loop, daemon=True).start()

    def _loop(self):
        try:
            while True:
                self.idle.set()
                self.go.wait()
                self.go.clear()
                self.idle.clear()
                while not self.stop.is_set():
                    self.conn.execute("DISCARD PLANS")
                    one(self.conn, "EXECUTE q(0)")
                    self.iterations += 1
        except Exception as e:  # surfaced by the driver
            self.error = repr(e)
            self.idle.set()

    def start(self):
        self.stop.clear()
        self.idle.clear()
        self.go.set()

    def halt(self):
        self.stop.set()
        self.idle.wait()
        if self.error:
            raise RuntimeError(f"worker pid {self.pid}: {self.error}")


def check_lock_table(conn, sessions, nparts):
    mlpt = int(one(conn, "SHOW max_locks_per_transaction"))
    mc = int(one(conn, "SHOW max_connections"))
    mpt = int(one(conn, "SHOW max_prepared_transactions"))
    capacity = mlpt * (mc + mpt)
    need = (sessions + 1) * (nparts + 10)
    if sessions + 3 > mc:
        sys.exit(f"--sessions {sessions} needs max_connections >= {sessions + 3} (have {mc})")
    if need > capacity:
        sys.exit(f"lock table too small: ~{need} entries needed, capacity {capacity} "
                 f"(max_locks_per_transaction={mlpt}, max_connections={mc}); raise "
                 f"max_locks_per_transaction or lower --partitions/--sessions")


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--dsn", default="", help="libpq connection string (default: environment)")
    ap.add_argument("--schema", default="plancache_stress")
    ap.add_argument("--partitions", type=int, default=800)
    ap.add_argument("--sessions", type=int, default=1)
    ap.add_argument("--minutes", type=float, default=5)
    ap.add_argument("--max-rounds", type=int, default=0, help="0 = no limit")
    ap.add_argument("--mode", default="force_generic_plan", choices=["force_generic_plan", "auto"])
    ap.add_argument("--log", help="write one JSON line per round to this file")
    ap.add_argument("--keep", action="store_true", help="do not drop the schema at the end")
    a = ap.parse_args()
    s = a.schema

    ctl = connect(a.dsn, "plancache_stress_ctl")
    print(one(ctl, "SELECT version()"))
    check_lock_table(ctl, a.sessions, a.partitions)
    print(f"setup: schema {s}, {a.partitions} partitions under mid ...", flush=True)
    setup(ctl, s, a.partitions)
    base = 100 + a.partitions
    attach = (f"ALTER TABLE {s}.mid ATTACH PARTITION {s}.float_sub "
              f"FOR VALUES FROM (100000) TO (200000)")
    detach = f"ALTER TABLE {s}.mid DETACH PARTITION {s}.float_sub"
    plain = PLAIN.format(s=s)
    logf = open(a.log, "w") if a.log else None

    workers = [Worker(a.dsn, i, s, a.mode) for i in range(a.sessions)]
    w0 = workers[0].conn
    t0 = time.time()
    for _ in range(5):
        w0.execute("DISCARD PLANS")
        one(w0, "EXECUTE q(0)")
    iter_s = (time.time() - t0) / 5
    print(f"sessions={a.sessions} mode={a.mode} one DISCARD+EXECUTE ~ {iter_s * 1000:.1f} ms; "
          f"correct answers: {base} (detached) / {base + 100} (attached)", flush=True)

    start = time.time()
    rounds = hits = 0
    first_hit = None
    try:
        while time.time() - start < a.minutes * 60 and (a.max_rounds == 0 or rounds < a.max_rounds):
            rounds += 1
            for w in workers:
                w.start()
            time.sleep(iter_s * 1.5)                     # let every worker get going
            time.sleep(random.uniform(0, iter_s * 1.2))  # random landing point
            ctl.execute(attach)
            for w in workers:
                w.halt()
            hit_workers = []
            for w in workers:
                pairs = []
                for _ in range(3):
                    w.conn.execute("BEGIN")
                    e = one(w.conn, "EXECUTE q(0)")
                    w.conn.execute("COMMIT")
                    w.conn.execute("BEGIN")
                    p = one(w.conn, plain)
                    w.conn.execute("COMMIT")
                    pairs.append((e, p))
                if all(e != p for e, p in pairs):
                    plan = [r[0] for r in w.conn.execute("EXPLAIN (COSTS OFF) EXECUTE q(0)").fetchall()]
                    hit_workers.append({"pid": w.pid, "execute_vs_plain": pairs,
                                        "plan_scans_float_sub": any("float_sub" in l for l in plan)})
            if hit_workers:
                hits += 1
                if first_hit is None:
                    first_hit = time.time() - start
                    print(f"first hit: round {rounds}, {first_hit:.1f} s: {hit_workers[0]}", flush=True)
            if logf:
                logf.write(json.dumps({"round": rounds, "t": round(time.time() - start, 3),
                                       "hits": hit_workers}) + "\n")
            ctl.execute(detach)
            if rounds % 100 == 0:
                print(f"  {rounds} rounds, {hits} hit rounds, {time.time() - start:.0f} s", flush=True)
    except KeyboardInterrupt:
        print("interrupted")
    finally:
        for w in workers:
            w.stop.set()
        elapsed = time.time() - start
        print(json.dumps({"rounds": rounds, "hit_rounds": hits,
                          "hit_rate": round(hits / rounds, 3) if rounds else None,
                          "first_hit_s": round(first_hit, 1) if first_hit is not None else None,
                          "elapsed_s": round(elapsed, 1), "sessions": a.sessions,
                          "partitions": a.partitions, "mode": a.mode,
                          "executes": sum(w.iterations for w in workers)}))
        if not a.keep:
            for w in workers:
                w.conn.close()
            teardown(ctl, s, a.partitions)


if __name__ == "__main__":
    main()