#!/bin/bash
# Starts everything for the "server closed the connection unexpectedly"
# measurements: PostgreSQL 13, 15, 17 and 18 on a private Docker network, a
# cgroup-limited 17 for the OOM test, and one container per client stack.
# Nothing is published on the host. The password "test" is a test value.
#
#   PREFIX=pgclosed ./setup.sh          # start, install, build (a few minutes)
#   PREFIX=pgclosed ./setup.sh clean    # remove these containers and the network
#
# The probes are written to ./work/ next to this script, which is mounted into
# the client containers as /w.
set -eu
PREFIX=${PREFIX:-pgclosed}
NET="$PREFIX-net"
HERE=$(cd "$(dirname "$0")" && pwd)
W="$HERE/work"
SERVERS="pg13 pg15 pg17 pg18 pgnet pgxid pgoom"
CLIENTS="cli node java go old"

if [ "${1:-}" = clean ]; then
  [ -d "$W" ] && docker run --rm -v "$W:/w" python:3.12-slim rm -rf /w/shared /w/out >/dev/null 2>&1 || true
  for c in $SERVERS $CLIENTS bou; do docker rm -f "$PREFIX-$c" >/dev/null 2>&1 || true; done
  docker network rm "$NET" >/dev/null 2>&1 || true
  rm -rf "$W"
  echo "removed"; exit 0
fi

mkdir -p "$W/shared" "$W/out"; chmod 777 "$W/shared" "$W/out"
docker network inspect "$NET" >/dev/null 2>&1 || docker network create "$NET" >/dev/null

server() { # short image [docker args --] [server args]
  s=$1; image=$2; shift 2; dargs=()
  while [ $# -gt 0 ] && [ "$1" != -- ]; do dargs+=("$1"); shift; done
  [ "${1:-}" = -- ] && shift
  docker rm -f "$PREFIX-$s" >/dev/null 2>&1 || true
  docker run -d --name "$PREFIX-$s" --network "$NET" --hostname "$s" "${dargs[@]}" \
    -e POSTGRES_PASSWORD=test "$image" -c "log_line_prefix=%m [%p] " "$@" >/dev/null
  echo "$PREFIX-$s $image"
}
client() { # short image [docker args]
  s=$1; image=$2; shift 2
  docker rm -f "$PREFIX-$s" >/dev/null 2>&1 || true
  docker run -d --name "$PREFIX-$s" --network "$NET" --hostname "$s" "$@" -v "$W:/w" \
    -e PGPASSWORD=test "$image" sleep 86400 >/dev/null
  echo "$PREFIX-$s $image"
}

server pg13 postgres:13
server pg15 postgres:15
server pg17 postgres:17
server pg18 postgres:18
server pgnet postgres:17 --cap-add NET_ADMIN
server pgxid postgres:17
server pgoom postgres:17 --memory=300m --memory-swap=300m -- \
  -c work_mem=2GB -c shared_buffers=32MB -c hash_mem_multiplier=4
client cli  python:3.12-slim --cap-add NET_ADMIN
client node node:22-slim
client java maven:3-eclipse-temurin-21
client go   golang:1.25
client old  python:3.9-slim

for s in pg13 pg15 pg17 pg18 pgnet pgxid pgoom; do
  c="$PREFIX-$s"; n=0
  until docker logs "$c" 2>&1 | grep -q 'listening on IPv4' && docker exec "$c" pg_isready -q -U postgres; do
    n=$((n + 1)); [ $n -gt 60 ] && { echo "$c did not come up" >&2; exit 1; }
    sleep 1
  done
done

docker exec "$PREFIX-pgnet" sh -c 'apt-get update -qq >/dev/null 2>&1; apt-get install -y -qq iptables >/dev/null 2>&1'
docker exec "$PREFIX-cli" sh -c 'apt-get update -qq >/dev/null 2>&1; apt-get install -y -qq postgresql-client iptables util-linux >/dev/null 2>&1
  pip install -q --root-user-action=ignore psycopg2-binary==2.9.13 "psycopg[binary]==3.3.6" sqlalchemy==2.1.6 2>&1 | grep -v notice || true'
docker exec "$PREFIX-node" sh -c 'mkdir -p /app && cd /app && npm init -y >/dev/null && npm install --silent pg@8.23.1 >/dev/null'
docker exec "$PREFIX-java" sh -c 'mkdir -p /pgjdbc && curl -fsSL -o /pgjdbc/pg.jar https://repo1.maven.org/maven2/org/postgresql/postgresql/42.7.13/postgresql-42.7.13.jar'
docker exec "$PREFIX-old" sh -c 'for v in 2.8.6 2.9.3 2.9.9; do python -m venv /v$v && /v$v/bin/pip install -q psycopg2-binary==$v 2>&1 | grep -v -E "notice|WARNING" || true; done'

# ---- probes ---------------------------------------------------------------
cat > "$W/p.py" <<'EOF'
import os,sys,time
lib=sys.argv[1]; mode=os.environ.get('MODE','idle'); go=os.environ['GOF']
base=f"host={os.environ['H']} port={os.environ.get('P','5432')} user=postgres password=test dbname=postgres application_name=p-{lib} sslmode=disable {os.environ.get('CI','')}"
if lib=='psycopg2':
    import psycopg2 as D; state=lambda e:e.pgcode
    c=D.connect(base); c.autocommit=True
else:
    import psycopg as D; state=lambda e:e.sqlstate
    c=D.connect(base,autocommit=True)
def q(tag,s):
    t0=time.time()
    try:
        cur=c.cursor(); cur.execute(s); print(f"[{lib}] {tag} OK +{time.time()-t0:.2f}s",flush=True)
    except Exception as e:
        print(f"[{lib}] {tag} +{time.time()-t0:.2f}s {type(e).__module__}.{type(e).__name__} sqlstate={state(e)} msg={str(e)!r}",flush=True)
print(f"[{lib}] READY",flush=True)
if mode=='txn': q('begin','begin')
if mode=='busy': q('q0-running','select pg_sleep(45)')
else:
    while not os.path.exists(go): time.sleep(.05)
for i in (1,2,3): q(f'q{i}',f'select {i}')
EOF
cat > "$W/psql.sh" <<'EOF'
#!/bin/bash
{ [ "$MODE" = txn ] && echo "begin;"
  [ "$MODE" = busy ] && echo "select pg_sleep(45);"
  if [ "$MODE" != busy ]; then while [ ! -e "$GOF" ]; do sleep .05; done; fi
  echo "select 1;"; sleep .5; echo "select 2;"; } | \
 PGAPPNAME=p-psql PGSSLMODE=disable psql -X -h "$H" -p ${P:-5432} -U postgres -d postgres -At 2>&1 | sed 's/^/[psql] /'
echo "[psql] exit=${PIPESTATUS[1]}"
EOF
chmod +x "$W/psql.sh"
cat > "$W/p.js" <<'EOF'
const {Client}=require('/app/node_modules/pg');
const fs=require('fs');const mode=process.env.MODE||'idle';
const log=(...a)=>console.log('[node-pg]',...a);
const sh=e=>`${e.constructor.name} code=${e.code} msg=${JSON.stringify(e.message)}`;
(async()=>{
 const c=new Client({host:process.env.H,port:+(process.env.P||5432),user:'postgres',password:'test',database:'postgres',application_name:'p-node-pg'});
 if(!process.env.NOLISTENER){c.on('error',e=>log(`EVENT 'error' ${sh(e)}`));}
 await c.connect(); log('READY');
 const q=async(tag,s)=>{const t=Date.now();try{await c.query(s);log(`${tag} OK +${((Date.now()-t)/1000).toFixed(2)}s`)}catch(e){log(`${tag} +${((Date.now()-t)/1000).toFixed(2)}s ${sh(e)}`)}};
 if(mode=='txn')await q('begin','begin');
 if(mode=='busy')await q('q0-running','select pg_sleep(45)');
 else while(!fs.existsSync(process.env.GOF))await new Promise(r=>setTimeout(r,50));
 for(const i of [1,2,3])await q('q'+i,`select ${i}`);
 process.exit(0);
})().catch(e=>{log('TOP',sh(e));process.exit(1)});
EOF
cat > "$W/P.java" <<'EOF'
import java.sql.*;import java.nio.file.*;import java.util.*;
public class P{static Connection c;
 static void q(String t,String s){long t0=System.currentTimeMillis();try(Statement st=c.createStatement()){st.execute(s);System.out.printf("[pgjdbc] %s OK +%.2fs%n",t,(System.currentTimeMillis()-t0)/1000.0);}catch(SQLException e){
  String cz=e.getCause()==null?"":" cause="+e.getCause();
  System.out.printf("[pgjdbc] %s +%.2fs %s sqlstate=%s msg=%s%s%n",t,(System.currentTimeMillis()-t0)/1000.0,e.getClass().getName(),e.getSQLState(),e.getMessage().replace("\n","\\n"),cz.replace("\n","\\n"));}
  System.out.flush();}
 public static void main(String[] a)throws Exception{String mode=System.getenv().getOrDefault("MODE","idle");
  Properties p=new Properties();p.setProperty("user","postgres");p.setProperty("password","test");p.setProperty("ApplicationName","p-pgjdbc");
  c=DriverManager.getConnection("jdbc:postgresql://"+System.getenv("H")+":"+System.getenv().getOrDefault("P","5432")+"/postgres?sslmode=disable",p);
  System.out.println("[pgjdbc] READY");System.out.flush();
  if(mode.equals("txn"))q("begin","begin");
  if(mode.equals("busy"))q("q0-running","select pg_sleep(45)");
  else while(!Files.exists(Path.of(System.getenv("GOF"))))Thread.sleep(50);
  for(int i=1;i<=3;i++)q("q"+i,"select "+i);System.exit(0);}}
EOF
mkdir -p "$W/gop"
cat > "$W/gop/main.go" <<'EOF'
package main
import("context";"errors";"fmt";"os";"time";"github.com/jackc/pgx/v5";"github.com/jackc/pgx/v5/pgconn")
func main(){ctx:=context.Background();h:=os.Getenv("H");pt:=os.Getenv("P");if pt==""{pt="5432"}
 c,err:=pgx.Connect(ctx,fmt.Sprintf("postgres://postgres:test@%s:%s/postgres?application_name=p-pgx&sslmode=disable",h,pt))
 if err!=nil{fmt.Println("[pgx] conn",err);return}
 fmt.Println("[pgx] READY")
 q:=func(t,s string){t0:=time.Now();_,e:=c.Exec(ctx,s);d:=time.Since(t0).Seconds()
  if e==nil{fmt.Printf("[pgx] %s OK +%.2fs\n",t,d);return}
  var pe *pgconn.PgError;st:="-";if errors.As(e,&pe){st=pe.Code}
  fmt.Printf("[pgx] %s +%.2fs %T sqlstate=%s safeToRetry=%v msg=%q\n",t,d,e,st,pgconn.SafeToRetry(e),e.Error())}
 m:=os.Getenv("MODE")
 if m=="txn"{q("begin","begin")}
 if m=="busy"{q("q0-running","select pg_sleep(45)")}else{for{if _,e:=os.Stat(os.Getenv("GOF"));e==nil{break};time.Sleep(50*time.Millisecond)}}
 for i:=1;i<=3;i++{q(fmt.Sprintf("q%d",i),fmt.Sprintf("select %d",i))}}
EOF
docker exec "$PREFIX-go" sh -c 'cd /w/gop && (go mod init p >/dev/null 2>&1 || true) && go get github.com/jackc/pgx/v5@v5.11.0 >/dev/null 2>&1 && CGO_ENABLED=0 go build -o /w/pgxbin . '
cat > "$W/n.py" <<'EOF'
# One psycopg2 connection into the black hole. WAIT = seconds between the drop and the query.
import psycopg2,os,sys,time
mode=sys.argv[1]
c=psycopg2.connect(f"host={os.environ['H']} user=postgres password=test dbname=postgres application_name=n-py sslmode=disable {os.environ.get('CI','')}"); c.autocommit=True
cur=c.cursor(); s='select pg_sleep(300)' if mode=='busy' else 'select 1'
open('/w/shared/n-ready','w').close()
if mode=='idle':
    while not os.path.exists('/w/shared/n-dropped'): time.sleep(.02)
    time.sleep(float(os.environ.get('WAIT','1')))
t0=time.time()
try:
    cur.execute(s); print(f"OK +{time.time()-t0:.2f}s",flush=True)
except Exception as e:
    print(f"+{time.time()-t0:.2f}s {type(e).__name__} {str(e)!r} pgcode={e.pgcode}",flush=True)
EOF
cat > "$W/race.py" <<'EOF'
# Autocommit INSERT loop with a rising id. On every error: reconnect and carry on.
import psycopg2,os,time
H=os.environ['H']; N=int(os.environ['N']); SC=os.environ['SC']
def conn():
    while True:
        try:
            c=psycopg2.connect(f"host={H} user=postgres password=test dbname=postgres application_name=race sslmode=disable options='-c synchronous_commit={SC}'"); c.autocommit=True; return c
        except psycopg2.OperationalError: time.sleep(.05)
c=conn(); cur=c.cursor()
i=0; acked=[]; amb=[]; msgs={}
while len(amb)<N:
    i+=1
    try:
        cur.execute('insert into r(id) values(%s)',(i,)); acked.append(i)
    except psycopg2.Error as e:
        m=type(e).__name__+': '+str(e).split('\n')[0]; msgs[m]=msgs.get(m,0)+1
        amb.append(i); c=conn(); cur=c.cursor()
c=conn(); cur=c.cursor(); cur.execute('select id from r'); rows={r[0] for r in cur.fetchall()}
print(f'synchronous_commit={SC} statements',i,'errors',len(amb),'errored-but-committed',sum(a in rows for a in amb),'errored-not-committed',sum(a not in rows for a in amb))
print('acknowledged',len(acked),'acknowledged-but-gone',sum(a not in rows for a in acked))
print(msgs)
open('/w/shared/race-done','w').close()
EOF
cat > "$W/sqla.py" <<'EOF'
import sys,time,os
from sqlalchemy import create_engine,text
dial,pp,size=sys.argv[1],sys.argv[2]=='1',int(sys.argv[3])
kw=eval(os.environ.get('KW') or '{}')
e=create_engine(f"postgresql+{dial}://postgres:test@{os.environ['H']}/postgres",pool_pre_ping=pp,pool_size=size,max_overflow=0,connect_args=dict(application_name='sa',sslmode='disable',**kw))
cs=[e.connect() for _ in range(size)]
for c in cs: c.execute(text('select 1')); c.commit()
for c in cs: c.close()
open('/w/shared/sa-ready','w').close()
while not os.path.exists('/w/shared/sa-go'): time.sleep(.02)
time.sleep(float(os.environ.get('D','1')))
for i in range(size+2):
    t0=time.time()
    try:
        with e.connect() as c:
            c.execute(text('select 1')); print(f"req{i} OK +{time.time()-t0:.2f}s",flush=True)
    except Exception as x:
        print(f"req{i} +{time.time()-t0:.2f}s {type(x).__name__} invalidated={getattr(x,'connection_invalidated',None)} {str(x).splitlines()[0]!r}",flush=True)
EOF
cat > "$W/old.py" <<'EOF'
# Which psycopg2 keeps the FATAL? The terminate comes 1.5 s in. idle: the query follows it. busy: the query is running.
import psycopg2,os,time
c=psycopg2.connect(f"host={os.environ['H']} user=postgres password=test dbname=postgres application_name=old sslmode=disable"); c.autocommit=True
cur=c.cursor()
try:
    if os.environ['MODE']=='busy': cur.execute('select pg_sleep(20)')
    else:
        time.sleep(3); cur.execute('select 1')
    print('OK')
except Exception as e:
    print(f"psycopg2 {psycopg2.__version__.split()[0]} libpq {psycopg2.__libpq_version__} {os.environ['MODE']}: {type(e).__name__} pgcode={e.pgcode} msg={str(e)!r}")
EOF
cat > "$W/pc.py" <<'EOF'
# INSERT in a transaction, print the xid, COMMIT. KIND=before: a deferred trigger holds COMMIT
# before the commit record. KIND=after: a synchronous standby that does not exist holds it after the flush.
import psycopg2,os,time
c=psycopg2.connect(f"host={os.environ['H']} user=postgres password=test dbname=postgres application_name=c-py sslmode=disable"); c.autocommit=True
cur=c.cursor(); cur.execute('begin')
if os.environ['KIND']=='after': cur.execute('set local synchronous_commit=on')
cur.execute("insert into t(tag) values(%s)",(os.environ['TAG'],))
cur.execute('select pg_current_xact_id()'); print('XID',cur.fetchone()[0],flush=True)
t0=time.time()
try:
    cur.execute('commit'); print('commit OK',flush=True)
except Exception as e:
    print(f"commit +{time.time()-t0:.2f}s {type(e).__name__} pgcode={e.pgcode} msg={str(e)!r} transaction_status={c.info.transaction_status}",flush=True)
EOF
mkdir -p "$W/goc"
cat > "$W/goc/main.go" <<'EOF'
package main
import("context";"fmt";"os";"time";"github.com/jackc/pgx/v5";"github.com/jackc/pgx/v5/pgconn")
func main(){ctx:=context.Background()
 c,err:=pgx.Connect(ctx,fmt.Sprintf("postgres://postgres:test@%s:5432/postgres?application_name=c-pgx&sslmode=disable",os.Getenv("H")))
 if err!=nil{fmt.Println(err);return}
 tx,_:=c.Begin(ctx)
 if os.Getenv("KIND")=="after"{tx.Exec(ctx,"set local synchronous_commit=on")}
 tx.Exec(ctx,"insert into t(tag) values($1)",os.Getenv("TAG"))
 t0:=time.Now();e:=tx.Commit(ctx)
 if e==nil{fmt.Println("commit OK");return}
 fmt.Printf("commit +%.2fs %T safeToRetry=%v msg=%q\n",time.Since(t0).Seconds(),e,pgconn.SafeToRetry(e),e.Error())}
EOF
docker exec "$PREFIX-go" sh -c 'cd /w/goc && (go mod init c >/dev/null 2>&1 || true) && go get github.com/jackc/pgx/v5@v5.11.0 >/dev/null 2>&1 && CGO_ENABLED=0 go build -o /w/pgxc . '
echo "ready"