#!/bin/bash
# repro_parallel_apply.sh -- backlog-drain benchmark for non-streamed parallel apply (v27 thread)
#
# Starts a publisher and a subscriber on this host, creates the schema, loads it, and for each
# point <kind>:<workers> times how long the subscriber needs to apply a backlog of
# transactions with max_parallel_apply_workers_per_subscription = <workers> (0 = serial).
# Prints one DRAIN line per drain and a summary: median txn/s and speedup vs workers=0.
#
# Usage:  PGBIN=/path/to/install/bin ./repro_parallel_apply.sh [rounds] ["points"]
#   rounds   default 3 (median over rounds; point order is reversed in even rounds)
#   points   default "hot10:0 hot10:4 indep:0 indep:4 indepK1000:0 indepK1000:4"
#   PGBIN    directory with initdb, pg_ctl, psql, pgbench (default: whatever is in PATH)
#   WORKDIR  scratch dir for both clusters (default: mktemp -d); removed at exit unless KEEP=1
#   N        transactions per drain for hot<H> / indep (default 30000)
#   TOT      row changes per drain for indepK<K>: TOT/K txns, at most N (default 600000)
#   CLIENTS  pgbench clients generating the backlog (default 8)
#   STAB     1 (default): VACUUM + CHECKPOINT on both nodes before every drain
#   PORT_P / PORT_S  default 56351 / 56352 (unix socket only, trust auth)
#
# Workloads (each pgbench script is one transaction):
#   hot<H>     50% UPDATE counters SET version = version + 1 WHERE id = random(1, H)
#              50% INSERT INTO events (new row)
#   indep      50% UPDATE rows SET version = version + 1 WHERE id = nextval(...): a row no
#              other txn touches (nextval runs as its own read-only statement: no xid, not
#              replicated), 50% INSERT INTO events
#   indepK<K>  one UPDATE of K consecutive rows of "rows"; ranges are disjoint (they start at
#              a multiple of K, wrapped modulo 3,999,000), no INSERTs
#
# Method: the subscription is DISABLED and the worker count is set (ALTER SYSTEM + reload; the
# leader only picks it up when it restarts, hence the DISABLE/ENABLE); the backlog is committed
# on the publisher, then a marker row; ENABLE; time until the marker is visible on the
# subscriber (commit order is preserved, so everything before it has been applied). Polling is
# every 10 ms, so very short (serial) drains are approximate. maxPA = parallel apply workers
# seen during the drain. Our runs pinned the whole script to 8 CPUs (taskset -c ...).
# Drains shorter than ~0.5 s (serial, or a build without the 1 ms sleep, at N=30000) show a
# start-up jitter of ~0.3 s on either side; for those use N=300000 (and TOT=3000000).
set -eu -o pipefail
rounds=${1:-3}
points=${2:-"hot10:0 hot10:4 indep:0 indep:4 indepK1000:0 indepK1000:4"}
N=${N:-30000}; TOT=${TOT:-600000}; CLIENTS=${CLIENTS:-8}; STAB=${STAB:-1}
PP=${PORT_P:-56351}; PS=${PORT_S:-56352}
if [ -n "${PGBIN:-}" ]; then export PATH=$PGBIN:$PATH; fi
for b in initdb pg_ctl postgres psql pgbench; do
  command -v $b >/dev/null || { echo "$b not found (set PGBIN)" >&2; exit 1; }
done
R=${WORKDIR:-$(mktemp -d)}; S=$R/sock; mkdir -p "$S"
export PGOPTIONS="-c client_min_messages=warning"
q() { local p=$1; shift; psql -X -q -At -h "$S" -p "$p" -d postgres -v ON_ERROR_STOP=1 "$@"; }
stop() {
  for n in P S; do pg_ctl -D "$R/$n" -m immediate stop >/dev/null 2>&1 || true; done
  if [ "${KEEP:-0}" = 1 ]; then echo "# kept $R"; else rm -rf "$R"; fi
}
trap stop EXIT

for n in P S; do
  p=$PP; [ $n = S ] && p=$PS
  initdb -A trust -N -D "$R/$n" >/dev/null
  cat >>"$R/$n/postgresql.conf" <<EOF
port = $p
unix_socket_directories = '$S'
listen_addresses = ''
shared_buffers = 256MB
fsync = off
wal_level = logical
max_logical_replication_workers = 12
max_worker_processes = 24
max_wal_size = 4GB
EOF
  pg_ctl -D "$R/$n" -l "$R/$n.log" -w start >/dev/null
done
echo "# $(postgres --version); $(nproc) CPUs; rounds=$rounds N=$N TOT=$TOT CLIENTS=$CLIENTS"

for p in $PP $PS; do q $p -c "
  CREATE TABLE counters (id int PRIMARY KEY, version bigint NOT NULL);
  CREATE TABLE rows (id bigint PRIMARY KEY, version bigint NOT NULL);
  CREATE TABLE events (id bigserial PRIMARY KEY, c int);
  CREATE TABLE marker (id int PRIMARY KEY)"; done
q $PP -c "
  INSERT INTO counters SELECT g, 0 FROM generate_series(1, 10000) g;
  INSERT INTO rows SELECT g, 0 FROM generate_series(1, 4000000) g;
  CREATE SEQUENCE s_rows; CREATE SEQUENCE s_blk;
  CREATE PUBLICATION pub FOR TABLE counters, rows, events, marker"
q $PS -c "CREATE SUBSCRIPTION sub CONNECTION 'host=$S port=$PP dbname=postgres' PUBLICATION pub"
until [ "$(q $PS -c "SELECT count(*) FROM pg_subscription_rel WHERE srsubstate <> 'r'")" = 0 ]
do sleep 0.5; done

cat >"$R/hot.pgb" <<'EOF'
\set h random(1, :hot)
UPDATE counters SET version = version + 1 WHERE id = :h;
EOF
cat >"$R/indep.pgb" <<'EOF'
SELECT nextval('s_rows') AS k \gset
UPDATE rows SET version = version + 1 WHERE id = :k;
EOF
cat >"$R/ins.pgb" <<'EOF'
\set h random(1, 10)
INSERT INTO events (c) VALUES (:h);
EOF
cat >"$R/blk.pgb" <<'EOF'
SELECT (nextval('s_blk') * :K) % 3999000 + 1 AS b \gset
UPDATE rows SET version = version + 1 WHERE id BETWEEN :b AND :b + :K - 1;
EOF

mk=0; pts=($points); out=$R/drains.txt; : >"$out"
for r in $(seq "$rounds"); do
  order=("${pts[@]}")
  if [ $((r % 2)) = 0 ]; then
    order=(); for ((i = ${#pts[@]} - 1; i >= 0; i--)); do order+=("${pts[$i]}"); done
  fi
  for pt in "${order[@]}"; do
    kind=${pt%%:*}; w=${pt##*:}
    q $PS -c "ALTER SUBSCRIPTION sub DISABLE"
    q $PS -c "ALTER SYSTEM SET max_parallel_apply_workers_per_subscription = $w"
    q $PS -c "SELECT pg_reload_conf()" >/dev/null
    until [ "$(q $PP -c "SELECT count(*) FROM pg_replication_slots
                         WHERE slot_name = 'sub' AND active")" = 0 ]; do sleep 0.1; done
    if [ "$STAB" = 1 ]; then
      for p in $PP $PS; do
        q $p -c "VACUUM rows" -c "VACUUM events" -c "VACUUM counters" -c "CHECKPOINT"
      done
    fi
    ntx=$N
    case $kind in
      indep) f=(-f "$R/indep.pgb" -f "$R/ins.pgb");;
      indepK*) K=${kind#indepK}; ntx=$((TOT / K)); [ $ntx -gt "$N" ] && ntx=$N
               f=(-D K="$K" -f "$R/blk.pgb");;
      hot*) f=(-D hot="${kind#hot}" -f "$R/hot.pgb" -f "$R/ins.pgb");;
      *) echo "unknown kind $kind" >&2; exit 1;;
    esac
    ntx=$(( (ntx / CLIENTS) * CLIENTS ))
    pgbench -n -h "$S" -p $PP -c "$CLIENTS" -j "$CLIENTS" -t $((ntx / CLIENTS)) "${f[@]}" \
      postgres >"$R/pgb.out" 2>&1 || { echo "pgbench failed" >&2; cat "$R/pgb.out" >&2; exit 1; }
    mk=$((mk + 1)); q $PP -c "INSERT INTO marker VALUES ($mk)"
    t0=$(date +%s.%N); q $PS -c "ALTER SUBSCRIPTION sub ENABLE"; maxpa=0
    while [ "$(q $PS -c "SELECT count(*) FROM marker WHERE id = $mk")" != 1 ]; do
      a=$(q $PS -c "SELECT count(*) FROM pg_stat_subscription
                    WHERE worker_type = 'parallel apply'")
      [ "$a" -gt $maxpa ] && maxpa=$a
      sleep 0.01
    done
    t1=$(date +%s.%N)
    awk -v r="$r" -v k="$kind" -v w="$w" -v a="$t0" -v b="$t1" -v n="$ntx" -v m=$maxpa 'BEGIN {
      printf "DRAIN round=%d kind=%s workers=%d txns=%d secs=%.3f txn_per_s=%.0f maxPA=%d\n",
             r, k, w, n, b - a, n / (b - a), m }' | tee -a "$out"
  done
done

ck="SELECT sum(version) FROM rows UNION ALL SELECT sum(version) FROM counters
    UNION ALL SELECT count(*) FROM events"
if [ "$(q $PP -c "$ck")" = "$(q $PS -c "$ck")" ]; then echo "# check: P and S identical"
else echo "# check: P and S DIFFER"; fi
echo "# summary (median over rounds; speedup = median / median at workers=0)"
awk '
  { for (i = 1; i <= NF; i++) { split($i, kv, "="); v[kv[1]] = kv[2] }
    key = v["kind"] SUBSEP v["workers"]; c[key]++; x[key, c[key]] = v["txn_per_s"] + 0
    if (!(key in seen)) { seen[key] = 1; keys[++nk] = key } }
  function med(key,   i, j, t, n) {
    n = c[key]
    for (i = 1; i <= n; i++) for (j = i + 1; j <= n; j++)
      if (x[key, j] < x[key, i]) { t = x[key, i]; x[key, i] = x[key, j]; x[key, j] = t }
    return n % 2 ? x[key, (n + 1) / 2] : (x[key, n / 2] + x[key, n / 2 + 1]) / 2 }
  END {
    for (i = 1; i <= nk; i++) m[keys[i]] = med(keys[i])
    for (i = 1; i <= nk; i++) {
      split(keys[i], p, SUBSEP); base = p[1] SUBSEP 0
      sp = (base in m) ? sprintf("%.4fx", m[keys[i]] / m[base]) : "n/a"
      printf "%-12s workers=%-2s median=%9.0f txn/s  n=%d  vs serial=%s\n",
             p[1], p[2], m[keys[i]], c[keys[i]], sp } }' "$out"