#!/usr/bin/env bash
# Throughput of three variants with pgbench -S -M prepared, interleaved and
# rotated per round so slow drift is spread evenly across them:
#   base   = master without the patch                          (BASE_PORT)
#   pplain = patched server, GUC at its default (plain)        (PATCHED_PORT)
#   prich  = patched server, ready_for_query_message=rich      (PATCHED_PORT, via PGOPTIONS)
# The client is always the PATCHED pgbench: an unpatched libpq cannot connect
# in rich mode (see nocfbot-rfq-old-libpq-client.c.txt).
#
#   PATCHED_BIN=/path/to/patched/bin bash rfq_pgbench.sh [rounds] [seconds] [clients]
#
# Both servers must already be running (see the README).  Optional: CPUS=4-11
# pins every pgbench run to those cores; OUT= sets the output directory.
set -euo pipefail
: "${PATCHED_BIN:?set PATCHED_BIN to the bin/ directory of the patched install}"
SOCK=${SOCK:-/tmp/rfq-sock}
BASE_PORT=${BASE_PORT:-54501}
PATCHED_PORT=${PATCHED_PORT:-54502}
ROUNDS=${1:-10}; SECS=${2:-10}; CLIENTS=${3:-8}
JOBS=$(( CLIENTS < 4 ? CLIENTS : 4 ))
OUT=${OUT:-./rfq-pgbench-out}
PGB=$PATCHED_BIN/pgbench
PIN=(); [ -n "${CPUS:-}" ] && PIN=(taskset -c "$CPUS")
mkdir -p "$OUT"; rm -f "$OUT"/{base,pplain,prich}.txt

for port in "$BASE_PORT" "$PATCHED_PORT"; do
  "$PGB" -h "$SOCK" -p "$port" -U postgres -i -s 10 -q postgres >/dev/null 2>&1
done

variant_port() { [ "$1" = base ] && echo "$BASE_PORT" || echo "$PATCHED_PORT"; }
variant_opts() { [ "$1" = prich ] && echo '-c ready_for_query_message=rich' || true; }

run_one() {  # <variant> <seconds> -> prints tps
  PGOPTIONS="$(variant_opts "$1")" "${PIN[@]}" "$PGB" -h "$SOCK" -p "$(variant_port "$1")" \
    -U postgres -n -S -M prepared -c "$CLIENTS" -j "$JOBS" -T "$2" postgres 2>&1 |
    awk '/^tps/{print $3}'
}

# Sanity check: the effective setting seen by each variant.
for v in base pplain prich; do
  echo "$v: ready_for_query_message=$(PGOPTIONS="$(variant_opts "$v")" "$PATCHED_BIN/psql" -X -At \
    -h "$SOCK" -p "$(variant_port "$v")" -U postgres \
    -c "select coalesce(current_setting('ready_for_query_message', true), '(does not exist)')")"
done

for v in base pplain prich; do run_one "$v" 5 >/dev/null; done   # warm-up

order=(base pplain prich)
for ((r = 0; r < ROUNDS; r++)); do
  for ((k = 0; k < 3; k++)); do
    v=${order[$(( (r + k) % 3 ))]}
    run_one "$v" "$SECS" >> "$OUT/$v.txt"
  done
  echo "round $((r + 1))/$ROUNDS done"
done

python3 - "$OUT" "$CLIENTS" <<'PY'
import statistics as st
import sys

out, clients = sys.argv[1], int(sys.argv[2])
tps = {v: [float(x) for x in open(f"{out}/{v}.txt") if x.strip()]
       for v in ("base", "pplain", "prich")}
for v, xs in tps.items():
    med = st.median(xs)
    # closed loop: mean latency per transaction = clients / tps
    print(f"{v:7s} n={len(xs):2d}  median tps={med:10.1f}  latency={clients * 1e6 / med:7.3f} us"
          f"  all={[round(x) for x in xs]}")
b, p, r = (st.median(tps[v]) for v in ("base", "pplain", "prich"))
print(f"pplain vs base:  {(p / b - 1) * 100:+.2f}% (median tps)")
print(f"prich vs pplain: {(r / p - 1) * 100:+.2f}% (median tps)")
PY