# lib.sh -- ephemeral primary/standby pairs for the PG19 promotion-hang repro.
# Source it; then `pair_up 19` or `pair_up 18`.  Everything lives under
# $RUN (tmpfs), on high ports, and is torn down by pair_down.  Never touches
# any other cluster on the machine.
#
#   VER=19 -> REL_19_STABLE 0b8bff042a9 (19rc1)    ports 56190/56191
#   VER=18 -> REL_18_6                              ports 56180/56181
set -u
HERE=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
RUN=${RUN:-/tmp/shihao/run}
bin_of()  { case $1 in 19) echo /tmp/shihao/pg19rc1/inst/bin;; 18) echo /tmp/shihao/pg186/inst/bin;; esac; }
pport()   { echo $((56100 + $1 * 10)); }      # primary
sport()   { echo $((56101 + $1 * 10)); }      # standby
cport()   { echo $((56102 + $1 * 10)); }      # cascaded standby
pdir()    { echo "$RUN/$1/primary"; }
sdir()    { echo "$RUN/$1/standby"; }
cdir()    { echo "$RUN/$1/cascade"; }
now_ms()  { date +%s%3N; }

# q VER PORT SQL -- one-shot psql, tuples only, unaligned
q() { "$(bin_of "$1")/psql" -X -h 127.0.0.1 -p "$2" -d postgres -Atq -c "$3"; }

common_conf() {   # $1 datadir  $2 port
  cat >> "$1/postgresql.conf" <<CONF
port = $2
listen_addresses = '127.0.0.1'
unix_socket_directories = '$1'
wal_level = replica
max_wal_senders = 10
hot_standby = on
fsync = off
log_line_prefix = '%m [%p] %b '
log_min_messages = log
logging_collector = off
shared_buffers = 64MB
CONF
}

pair_up() {   # VER [cascade]
  local v=$1 b; b=$(bin_of "$v")
  pair_down "$v" >/dev/null 2>&1
  mkdir -p "$RUN/$v"
  "$b/initdb" -D "$(pdir "$v")" -A trust -U "$USER" --no-sync >"$RUN/$v/initdb.log" 2>&1
  common_conf "$(pdir "$v")" "$(pport "$v")"
  echo "host replication all 127.0.0.1/32 trust" >> "$(pdir "$v")/pg_hba.conf"
  "$b/pg_ctl" -D "$(pdir "$v")" -l "$RUN/$v/primary.log" -w start >/dev/null
  # No replication slot on purpose: the report is about wal_level=replica, no slots.
  "$b/pg_basebackup" -h 127.0.0.1 -p "$(pport "$v")" -D "$(sdir "$v")" -R -X stream --no-sync >/dev/null
  sed -i "s/^port = .*/port = $(sport "$v")/; s#^unix_socket_directories = .*#unix_socket_directories = '$(sdir "$v")'#" "$(sdir "$v")/postgresql.conf"
  "$b/pg_ctl" -D "$(sdir "$v")" -l "$RUN/$v/standby.log" -w start >/dev/null
  if [ "${2:-}" = cascade ]; then
    "$b/pg_basebackup" -h 127.0.0.1 -p "$(sport "$v")" -D "$(cdir "$v")" -R -X stream --no-sync >/dev/null
    sed -i "s/^port = .*/port = $(cport "$v")/; s#^unix_socket_directories = .*#unix_socket_directories = '$(cdir "$v")'#" "$(cdir "$v")/postgresql.conf"
    "$b/pg_ctl" -D "$(cdir "$v")" -l "$RUN/$v/cascade.log" -w start >/dev/null
  fi
  # Fail loudly if a server did not come up: a dead standby looks exactly like a
  # hung promotion to the callers ("still in recovery" never turns false).
  if ! wait_until 30 "$v" "$(pport "$v")" "select count(*) >= 1 from pg_stat_replication where state = 'streaming'" t \
     || [ "$(q "$v" "$(sport "$v")" "select pg_is_in_recovery()" 2>/dev/null)" != t ]; then
    mkdir -p /tmp/shihao/evidence/setup-failed
    cp "$RUN/$v"/*.log /tmp/shihao/evidence/setup-failed/ 2>/dev/null
    echo "pair_up $v: setup failed" >&2
    return 1
  fi
}

pair_down() {
  local v=$1 b d; b=$(bin_of "$v")
  pkill -f "stall_client.py $(sport "$v") " 2>/dev/null
  pkill -f "stall_client.py $(pport "$v") " 2>/dev/null
  for d in "$(cdir "$v")" "$(sdir "$v")" "$(pdir "$v")"; do
    [ -f "$d/postmaster.pid" ] && "$b/pg_ctl" -D "$d" -m immediate -w stop >/dev/null 2>&1
  done
  rm -rf "${RUN:?}/$v"
}

# wait_until SECS VER PORT SQL EXPECTED -- poll every 10 ms; 0 if it matched
wait_until() {
  local deadline=$(( $(now_ms) + $1 * 1000 ))
  while [ "$(now_ms)" -lt "$deadline" ]; do
    [ "$(q "$2" "$3" "$4" 2>/dev/null)" = "$5" ] && return 0
    sleep 0.01
  done
  return 1
}

# stall VER PORT [SQL] -- start a non-reading client, echo its backend pid once
# that backend sits in wait_event ClientWrite
stall() {
  local v=$1 port=$2 sql=${3:-"select repeat('x', 8000) from generate_series(1, 100000000)"} out pid
  out=$(mktemp "$RUN/$v/stall.XXXX")
  python3 "$HERE/stall_client.py" "$port" postgres "$sql" > "$out" 2>&1 &
  for _ in $(seq 500); do grep -q '^SENT' "$out" && break; sleep 0.01; done
  pid=$(awk '/^SENT/{print $2}' "$out")
  [ -n "$pid" ] || { echo "stall client failed: $(cat "$out")" >&2; return 1; }
  wait_until 20 "$v" "$port" "select wait_event from pg_stat_activity where pid = $pid" ClientWrite \
    || { echo "backend $pid never reached ClientWrite" >&2; return 1; }
  # ClientWrite is not enough.  Right after the client's buffer fills, its
  # kernel still reopens the window a few bytes at a time; a barrier that lands
  # then gets absorbed because the backend completes a flush and reaches
  # CHECK_FOR_INTERRUPTS().  Wait until the server socket stops moving data
  # (bytes_acked unchanged for 2 s) so every run starts from a truly stuck backend.
  local lport last="" cur stable=0
  lport=$(awk '/^SENT/{print $3}' "$out")
  for _ in $(seq 60); do
    cur=$(ss -tinH "( sport = :$port and dport = :$lport )" | grep -o 'bytes_acked:[0-9]*')
    if [ -n "$cur" ] && [ "$cur" = "$last" ]; then stable=$((stable + 1)); else stable=0; fi
    [ $stable -ge 4 ] && break
    last=$cur; sleep 0.5
  done
  [ $stable -ge 4 ] || echo "warning: socket of backend $pid never settled" >&2
  # Settled bytes_acked is still not enough.  The kernel can free room in the LOCAL
  # send buffer after the backend's last EAGAIN without reporting the socket as
  # writable (epoll waits for a threshold).  A barrier's SIGUSR1 then wakes the
  # backend, the retried write() fits, the message completes, the next
  # CHECK_FOR_INTERRUPTS() absorbs the barrier, and promotion goes through by luck.
  # Saturate first: wake the backend with a harmless signal until its Send-Q stops
  # growing.  (STALL_SATURATE=0 skips this, to show the luck.)
  if [ "${STALL_SATURATE:-1}" = 1 ]; then
    local sq prev="" same=0
    for _ in $(seq 40); do
      q "$v" "$port" "select pg_log_backend_memory_contexts($pid)" >/dev/null 2>&1
      sleep 0.2
      sq=$(ss -tnH "( sport = :$port and dport = :$lport )" | awk '{print $3}')
      if [ "$sq" = "$prev" ]; then same=$((same + 1)); else same=0; fi
      [ $same -ge 3 ] && break
      prev=$sq
    done
    [ $same -ge 3 ] || echo "warning: send queue of backend $pid kept growing" >&2
  fi
  echo "$pid"
}