#!/bin/bash
# pg_xact_status after a crash under load, on a fresh PostgreSQL 17 container.
# Four sessions loop over begin/insert/pg_current_xact_id()/commit. A backend gets
# kill -9, recovery runs, and every in-flight transaction is compared with its row.
# Usage: PREFIX=nn02r ./xid-load.sh     (optional: ROUNDS=15 SESSIONS=4)
# Removes only its own containers and network, also on failure.
set -eu
PREFIX=${PREFIX:?set PREFIX}
ROUNDS=${ROUNDS:-15}; SESSIONS=${SESSIONS:-4}
NET="$PREFIX-xl-net"; PG="$PREFIX-xl-pg"; CLI="$PREFIX-xl-cli"
cleanup() { docker rm -f "$PG" "$CLI" >/dev/null 2>&1 || true; docker network rm "$NET" >/dev/null 2>&1 || true; }
trap cleanup EXIT
cleanup
docker network create "$NET" >/dev/null
docker run -d --name "$PG" --network "$NET" --hostname pg -e POSTGRES_PASSWORD=test postgres:17 >/dev/null
docker run -d --name "$CLI" --network "$NET" -e PGPASSWORD=test python:3.12-slim sleep 3600 >/dev/null
for _ in $(seq 120); do
  docker logs "$PG" 2>&1 | grep -q 'listening on IPv4' && docker exec "$PG" pg_isready -q -U postgres && break
  sleep .5
done
docker exec "$CLI" pip install -q --root-user-action=ignore psycopg2-binary==2.9.13 2>&1 | grep -v notice || true
docker exec "$PG" psql -U postgres -qc "create table t(id serial primary key, tag text)"

docker exec -i "$CLI" sh -c 'cat > /tmp/load.py' <<'EOF'
import psycopg2,os,threading,json,time
N=int(os.environ['N']); R=os.environ['R']; res=[]
def conn():
    while True:
        try:
            c=psycopg2.connect("host=pg user=postgres password=test dbname=postgres application_name=lw sslmode=disable"); c.autocommit=True; return c
        except psycopg2.OperationalError: time.sleep(.05)
def w(k):
    cur=conn().cursor(); i=0
    while True:
        i+=1; tag=f"{R}-{k}-{i}"; xid=None
        try:
            cur.execute('begin'); cur.execute('insert into t(tag) values(%s)',(tag,))
            cur.execute('select pg_current_xact_id()'); xid=cur.fetchone()[0]; cur.execute('commit')
        except psycopg2.Error:
            res.append((tag,xid)); return
ts=[threading.Thread(target=w,args=(k,)) for k in range(N)]
[t.start() for t in ts]; [t.join() for t in ts]
json.dump(res,open('/tmp/res.json','w'))
EOF
docker exec -i "$CLI" sh -c 'cat > /tmp/chk.py' <<'EOF'
import psycopg2,json,time
res=json.load(open('/tmp/res.json'))
while True:
    try:
        c=psycopg2.connect("host=pg user=postgres password=test dbname=postgres sslmode=disable"); c.autocommit=True; break
    except psycopg2.OperationalError: time.sleep(.05)
cur=c.cursor()
def st(x):
    try:
        cur.execute("select pg_xact_status(%s::xid8)",(str(x),)); return cur.fetchone()[0]
    except psycopg2.Error as e:
        return 'future' if 'future' in str(e) else 'error'
a={x:st(x) for _,x in res if x}
for _ in range(30):
    cur.execute("insert into t(tag) values('other')")
for tag,x in res:
    cur.execute("select count(*) from t where tag=%s",(tag,)); n=cur.fetchone()[0]
    print(json.dumps([x,a.get(x),st(x) if x else None,n]))
EOF

all=$(mktemp)
for r in $(seq "$ROUNDS"); do
  docker exec -e N="$SESSIONS" -e R="r$r" "$CLI" python /tmp/load.py & p=$!
  sleep 1.2; pid=''
  for _ in $(seq 50); do
    pid=$(docker exec "$PG" psql -U postgres -Atc "select pid from pg_stat_activity where application_name='lw' limit 1" 2>/dev/null || true)
    [ -n "$pid" ] && break; sleep .1
  done
  docker exec "$PG" kill -9 "$pid"
  wait "$p"
  docker exec "$CLI" python /tmp/chk.py >> "$all"
done
python3 -I - "$all" "$ROUNDS" <<'EOF'
import json,sys,collections
rows=[json.loads(l) for l in open(sys.argv[1])]
C=collections.Counter((a,b,'row' if n else 'no row') for x,a,b,n in rows)
print('crashes',sys.argv[2]); print('sessions in flight',len(rows))
print('committed with row',sum(1 for x,a,b,n in rows if x and b=='committed' and n))
print('committed without row',sum(1 for x,a,b,n in rows if x and b=='committed' and not n))
print('no xid yet',sum(1 for x,a,b,n in rows if not x))
print('other',sum(1 for x,a,b,n in rows if x and b!='committed'))
print('detail (right-after, after 30 commits, row):')
for k,v in sorted(C.items(),key=str): print(' ',k,v)
EOF
rm -f "$all"