#!/bin/bash
# Prints the client messages and server logs behind the "server closed the
# connection unexpectedly" article. Run setup.sh first, with the same PREFIX.
#
#   PREFIX=pgclosed ./run.sh 2>&1 | tee results.txt      # everything except "long"
#   PREFIX=pgclosed ./run.sh race xid                    # chosen sections
#   PREFIX=pgclosed ./run.sh long                        # about 16 minutes
#
# Sections: terminate timeouts crash oom stops network rst serverside commit
#           race xid pool pgbouncer psycopg2 long
# The race and xid sections kill backends on purpose; only the containers from
# setup.sh are touched. Timings and the race percentages are yours, not ours.
set -u
PREFIX=${PREFIX:-pgclosed}
HERE=$(cd "$(dirname "$0")" && pwd)
W="$HERE/work"; SH="$W/shared"; CLI="$PREFIX-cli"
c() { echo "$PREFIX-$1"; }
sq() { docker exec "$(c "$1")" psql -U postgres -Atc "$2"; }
up() { sleep .5; for _ in $(seq 120); do docker exec "$(c "$1")" pg_isready -q -h 127.0.0.1 -U postgres && return; sleep .25; done; }
marker() { echo "go-$(date +%s%N)"; }
hdr() { printf '\n######## %s\n' "$*"; }

terminate_sql="select count(pg_terminate_backend(pid)) from pg_stat_activity where application_name like 'p-%'"
pids_sql="select string_agg(pid::text,' ') from pg_stat_activity where application_name like 'p-%'"

# six clients at once: probe <server> <mode> <seconds before queries> <trigger command> [port]
probe() {
  local s=$1 mode=$2 delay=$3 trigger=$4 port=${5:-5432} host=${6:-$1} g since
  g=$(marker); since=$(date -u +%FT%TZ)
  local E=(-e "H=$host" -e "MODE=$mode" -e "GOF=/w/shared/$g" -e "P=$port" -e "CI=${CI:-}")
  docker exec "${E[@]}" "$CLI" /w/psql.sh &
  docker exec "${E[@]}" "$CLI" python /w/p.py psycopg2 &
  docker exec "${E[@]}" "$CLI" python /w/p.py psycopg3 &
  docker exec "${E[@]}" "$CLI" /w/pgxbin &
  docker exec "${E[@]}" "$(c node)" node /w/p.js &
  docker exec "${E[@]}" "$(c java)" java -cp /pgjdbc/pg.jar /w/P.java &
  sleep "${READY:-6}"
  bash -c "$trigger" >/dev/null
  sleep "$delay"; docker exec "$CLI" touch "/w/shared/$g"
  wait
  [ "$mode" = busy ] || sleep 1
  echo "--- server log ($s)"; docker logs --since "$since" "$(c "${LOGOF:-$s}")" 2>&1 | grep -v checkpoint | head -30
}
killp() { printf 'p=$(docker exec %s psql -U postgres -Atc "%s"); docker exec %s kill -9 $p' "$(c "$1")" "$pids_sql" "$(c "$1")"; }

s_terminate() {
  for v in pg13 pg15 pg17 pg18; do hdr "pg_terminate_backend, idle, $v"; probe $v idle 1 "docker exec $(c $v) psql -U postgres -Atc \"$terminate_sql\""; done
  hdr "pg_terminate_backend, running query, pg17"; probe pg17 busy 0 "docker exec $(c pg17) psql -U postgres -Atc \"$terminate_sql\""
  hdr "node-pg without an 'error' listener"
  docker exec -e H=pg17 -e MODE=idle -e GOF=/w/shared/never -e NOLISTENER=1 "$(c node)" node /w/p.js 2>&1 | grep -E 'READY|Unhandled|^error:' & sleep 3
  sq pg17 "$terminate_sql" >/dev/null; wait
}
s_timeouts() {
  for row in "pg13 idle idle_session_timeout" "pg15 idle idle_session_timeout" "pg17 idle idle_session_timeout" "pg18 idle idle_session_timeout" \
             "pg13 txn idle_in_transaction_session_timeout" "pg17 txn idle_in_transaction_session_timeout" "pg17 txn transaction_timeout" "pg18 txn transaction_timeout"; do
    set -- $row; hdr "$3=2s, $1, mode $2"
    sq $1 "alter role postgres set $3='2s'"
    READY=7 probe $1 $2 2 true
    sq $1 "alter role postgres reset $3"
  done
}
s_crash() {
  for v in pg13 pg15 pg17 pg18; do
    hdr "kill -9 on the client's own backend, $v"; probe $v idle 3 "$(killp $v)"; up $v
    hdr "kill -9 on another backend (bystanders), $v"
    probe $v idle 3 "docker exec -d $(c $v) psql -U postgres -c 'set application_name=victim; select pg_sleep(300)'; sleep 1; docker exec $(c $v) sh -c 'kill -9 \$(psql -U postgres -Atc \"select pid from pg_stat_activity where application_name=\\\$\\\$victim\\\$\\\$\")'"; up $v
  done
}
s_oom() {
  hdr "real OOM kill: --memory=300m, work_mem=2GB, sort of 30 million rows"
  docker exec "$CLI" psql -X -h pgoom -U postgres -c "select count(*) from (select g, md5(g::text) m from generate_series(1,30000000) g order by m) s" 2>&1
  sleep 3; docker logs --since 2m "$(c pgoom)" 2>&1 | grep -E 'signal 9|Failed process|terminating any' | cut -c1-200
  docker inspect -f 'OOMKilled={{.State.OOMKilled}} Running={{.State.Running}}' "$(c pgoom)"
}
s_stops() {
  for v in pg13 pg15 pg17 pg18; do for m in fast immediate; do
    hdr "pg_ctl stop -m $m, idle, $v"
    probe $v idle 2 "docker exec -u postgres $(c $v) sh -c 'pg_ctl -D \$PGDATA stop -m $m'"
    docker start "$(c $v)" >/dev/null; up $v; sleep 1
  done; done
  hdr "docker stop, idle, pg17"; probe pg17 idle 1 "docker stop $(c pg17)"; docker start "$(c pg17)" >/dev/null; up pg17; sleep 1
  hdr "docker kill -s KILL, idle, pg17"; probe pg17 idle 2 "docker kill -s KILL $(c pg17)"; docker start "$(c pg17)" >/dev/null; up pg17; sleep 1
}
drop_rules() { docker exec "$(c pgnet)" sh -c 'iptables -I INPUT -p tcp --dport 5432 -j DROP; iptables -I OUTPUT -p tcp --sport 5432 -j DROP'; }
flush() { docker exec "$(c pgnet)" iptables -F; sq pgnet "select count(pg_terminate_backend(pid)) from pg_stat_activity where application_name like 'n-%' or application_name like 'p-%'" >/dev/null; sleep 1; }
ka='keepalives_idle=3 keepalives_interval=1 keepalives_count=3'; ut='tcp_user_timeout=5000'
hole() { # label mode wait-seconds conninfo
  hdr "black hole: $1 [$4]"
  docker exec "$CLI" sh -c 'rm -f /w/shared/n-ready /w/shared/n-dropped'
  docker exec -e H=pgnet -e WAIT="$3" -e CI="$4" "$CLI" timeout 45 python /w/n.py "$2" & local p=$!
  until [ -e "$SH/n-ready" ]; do sleep .05; done; sleep .7
  drop_rules; docker exec "$CLI" touch /w/shared/n-dropped
  wait $p; [ $? = 124 ] && echo "still hanging at the 45 s cut-off"
  flush
}
s_network() {
  hole "a  idle, query 1 s later"  idle 1  "$ka"
  hole "a2 idle, query 10 s later" idle 10 "$ka"
  hole "b  idle, query 1 s later"  idle 1  "$ut"
  hole "b2 idle, query 10 s later" idle 10 "$ut"
  hole "c  query running"          busy 0  ""
  hole "d  query running"          busy 0  "$ka"
  hole "e  query running"          busy 0  "$ut"
  hole "f  query running"          busy 0  "$ka $ut"
}
s_rst() {
  hdr "RST from the network; postgres keeps running"
  READY=7 probe pgnet idle 1 "docker exec $(c pgnet) iptables -I INPUT -p tcp --dport 5432 -j REJECT --reject-with tcp-reset"
  flush
}
s_serverside() { # the client vanishes while its query runs
  for row in "default||" "check=1s|client_connection_check_interval=1000|" "check=1s+keepalive|client_connection_check_interval=1000|ka"; do
    IFS='|' read -r label chk k <<<"$row"
    hdr "server side, query running, client becomes a black hole: $label"
    local set="set application_name='srv';"
    [ -n "$chk" ] && set="$set set $chk;"
    [ -n "$k" ] && set="$set set tcp_keepalives_idle=3; set tcp_keepalives_interval=1; set tcp_keepalives_count=3;"
    docker exec -d "$CLI" psql -X -h pgnet -U postgres -qc "$set select pg_sleep(30)"
    sleep 1.5; local since; since=$(date -u +%FT%TZ)
    docker exec "$CLI" sh -c 'iptables -I OUTPUT -p tcp --dport 5432 -j DROP; iptables -I INPUT -p tcp --sport 5432 -j DROP'
    local t0=$SECONDS gone=''
    while [ $((SECONDS - t0)) -lt 25 ]; do
      [ "$(sq pgnet "select count(*) from pg_stat_activity where application_name='srv'")" = 0 ] && { gone=$((SECONDS - t0)); break; }
      sleep .5
    done
    echo "backend gone after: ${gone:-more than 25} s"
    docker exec "$CLI" iptables -F
    docker logs --since "$since" "$(c pgnet)" 2>&1 | grep -E 'FATAL|client' | head -4
    sq pgnet "select count(pg_terminate_backend(pid)) from pg_stat_activity where application_name='srv'" >/dev/null; sleep 6
  done
}
prep_commit() {
  sq pgxid "create table if not exists t(id serial primary key, tag text)" >/dev/null
  sq pgxid "create or replace function slow() returns trigger language plpgsql as \$\$ begin perform pg_sleep(30); return null; end \$\$" >/dev/null
  sq pgxid "drop trigger if exists tr on t" >/dev/null
  sq pgxid "create constraint trigger tr after insert on t deferrable initially deferred for each row when (new.tag like 'slow%') execute function slow()" >/dev/null
  sq pgxid "alter system set synchronous_standby_names='nobody'" >/dev/null
  sq pgxid "alter system set synchronous_commit='local'" >/dev/null
  sq pgxid "select pg_reload_conf()" >/dev/null
}
s_commit() {
  prep_commit
  for how in kill9 terminate; do for kind in before after; do for lib in psycopg2 pgx; do
    [ $kind = before ] && tag="slow-$how-$lib" || tag="fast-$how-$lib"
    hdr "COMMIT interrupted: $how, $kind the commit record, $lib"
    if [ $lib = pgx ]; then docker exec -e H=pgxid -e TAG=$tag -e KIND=$kind "$CLI" /w/pgxc > "$W/out/c.txt" 2>&1 &
    else docker exec -e H=pgxid -e TAG=$tag -e KIND=$kind "$CLI" python /w/pc.py > "$W/out/c.txt" 2>&1 & fi
    local p=$! pid=''
    for _ in $(seq 100); do pid=$(sq pgxid "select pid from pg_stat_activity where application_name like 'c-%' and state='active' and query ilike 'commit%'"); [ -n "$pid" ] && break; sleep .1; done
    sleep .3
    if [ $how = kill9 ]; then docker exec "$(c pgxid)" kill -9 "$pid"; else sq pgxid "select pg_terminate_backend($pid)" >/dev/null; fi
    wait $p; cat "$W/out/c.txt"; up pgxid
    echo "rows after recovery: $(sq pgxid "select count(*) from t where tag='$tag'")"
  done; done; done
}
s_race() {
  sq pgnet "create table if not exists r(id bigint primary key)" >/dev/null
  for sc in on off; do
    hdr "100 unscripted kills, synchronous_commit=$sc"
    sq pgnet "truncate r" >/dev/null
    docker exec "$CLI" rm -f /w/shared/race-done
    docker exec -e H=pgnet -e N=100 -e SC=$sc "$CLI" python /w/race.py & local p=$!
    until [ -e "$SH/race-done" ]; do
      sleep 0.9
      pid=$(sq pgnet "select pid from pg_stat_activity where application_name='race'" 2>/dev/null)
      [ -n "$pid" ] && docker exec "$(c pgnet)" kill -9 "$pid" 2>/dev/null
    done
    wait $p
  done
}
s_xid() {
  prep_commit
  hdr "pg_xact_status after kill -9 during COMMIT, 8 rounds"
  for r in 1 2 3 4 5 6 7 8; do
    docker exec -e H=pgxid -e TAG="slow-xid-$r" -e KIND=before "$CLI" python /w/pc.py > "$W/out/x.txt" 2>&1 & local p=$! pid=''
    for _ in $(seq 100); do pid=$(sq pgxid "select pid from pg_stat_activity where application_name='c-py' and state='active' and query ilike 'commit%'"); [ -n "$pid" ] && break; sleep .1; done
    sleep .3; docker exec "$(c pgxid)" kill -9 "$pid"; wait $p; up pgxid
    x=$(grep -oE 'XID [0-9]+' "$W/out/x.txt" | cut -d' ' -f2)
    a=$(sq pgxid "select pg_xact_status('$x'::xid8)" 2>&1)
    n=$(docker exec "$(c pgxid)" psql -U postgres -At -c "begin; insert into t(tag) values('other') returning 0; select pg_current_xact_id(); commit;" 2>&1 | grep -E '^[1-9][0-9]*$')
    b=$(sq pgxid "select pg_xact_status('$x'::xid8)" 2>&1)
    echo "r$r xid=$x right-after='$a' other-xid=$n then='$b' own-rows=$(sq pgxid "select count(*) from t where tag='slow-xid-$r'")"
  done
}
s_pool() {
  sa() { # label dialect pre_ping trigger [KW] [cap]
    hdr "SQLAlchemy $1 $2 pre_ping=$3"
    docker exec "$CLI" sh -c 'rm -f /w/shared/sa-ready /w/shared/sa-go'
    docker exec -e H=pgnet -e KW="${5:-}" "$CLI" timeout "${6:-20}" python /w/sqla.py "$2" "$3" 3 2>&1 | cut -c1-230 & local p=$!
    until [ -e "$SH/sa-ready" ]; do sleep .05; done
    bash -c "$4" >/dev/null; docker exec "$CLI" touch /w/shared/sa-go; wait $p
  }
  local t="docker exec $(c pgnet) psql -U postgres -Atc \"select count(pg_terminate_backend(pid)) from pg_stat_activity where application_name='sa'\""
  local bh="docker exec $(c pgnet) sh -c 'iptables -I INPUT -p tcp --dport 5432 -j DROP; iptables -I OUTPUT -p tcp --sport 5432 -j DROP'"
  for d in psycopg2 psycopg; do for pp in 0 1; do sa terminate $d $pp "$t"; done; done
  sa blackhole psycopg2 1 "$bh" "" 30; echo "(cut off at 30 s if no req lines above)"; flush
  sa "blackhole+tcp_user_timeout" psycopg2 1 "$bh" "dict(tcp_user_timeout=3000)" 25; flush
}
s_pgbouncer() {
  printf '"postgres" "test"\n' > "$SH/ub.txt"
  cat > "$SH/pgb.ini" <<I
[databases]
postgres = host=pgnet port=5432 user=postgres password=test
[pgbouncer]
listen_addr = 0.0.0.0
listen_port = 6432
auth_type = plain
auth_file = /etc/pgbouncer/ub.txt
pool_mode = session
ignore_startup_parameters = extra_float_digits,options
log_disconnections = 1
client_idle_timeout = 2
I
  docker rm -f "$(c bou)" >/dev/null 2>&1
  docker run -d --name "$(c bou)" --network "$PREFIX-net" --hostname bou -v "$SH/pgb.ini:/etc/pgbouncer/pgbouncer.ini:ro" -v "$SH/ub.txt:/etc/pgbouncer/ub.txt:ro" edoburu/pgbouncer:latest >/dev/null
  sleep 2
  hdr "PgBouncer client_idle_timeout=2"
  READY=7 LOGOF=bou probe pgnet idle 1 true 6432 bou
  docker rm -f "$(c bou)" >/dev/null
}
s_psycopg2() {
  hdr "the FATAL per psycopg2 wheel: terminate on an idle and on a busy session, pg17"
  one() { docker exec -e H=pg17 -e MODE="$3" "$1" $2 /w/old.py > "$W/out/o.txt" 2>&1 & local p=$!
    sleep 1.5; sq pg17 "select count(pg_terminate_backend(pid)) from pg_stat_activity where application_name='old'" >/dev/null; wait $p; cat "$W/out/o.txt"; }
  for m in idle busy; do
    for v in 2.8.6 2.9.3 2.9.9; do one "$(c old)" "/v$v/bin/python" $m; done
    one "$CLI" python $m
  done
}
s_long() {
  hdr "black hole, all six clients, nothing configured, query 1 s later (about 16 minutes)"
  READY=8 probe pgnet idle 1 "docker exec $(c pgnet) sh -c 'iptables -I INPUT -p tcp --dport 5432 -j DROP; iptables -I OUTPUT -p tcp --sport 5432 -j DROP'"
  flush
}

sections=${*:-terminate timeouts crash oom stops network rst serverside commit race xid pool pgbouncer psycopg2}
for s in $sections; do "s_$s"; done