#!/usr/bin/env bash
# pgbench select-only on a partitioned pgbench_accounts, one installed tree.
# Set BUILD to the prefix (directory that contains bin/postgres).
#
# A generic plan locks every partition at executor startup.  Relation locks
# beyond the backend's fast-path slots (max_locks_per_transaction, default 64)
# go to the shared lock table, taking lock-manager partition LWLocks in
# exclusive mode (16 padded on base, 128 packed 8 per cache line on patched).
set -euo pipefail

BUILD="${BUILD:?set BUILD to the install prefix, e.g. .../builds/lwlock-aligned/patched}"
NCPU="$(sysctl -n hw.ncpu 2>/dev/null || nproc)"
CLIENTS="${CLIENTS:-1 $NCPU}"
DURATION="${DURATION:-15}"
ROUNDS="${ROUNDS:-3}"
SCALE="${SCALE:-10}"
PARTITIONS="${PARTITIONS:-128}"
PORT="${PORT:-55440}"
DATA="${DATA:-$(dirname "$BUILD")/data-pgbench-s$SCALE-p$PARTITIONS}"
LOG="${LOG:-$(dirname "$BUILD")/$(basename "$BUILD").log}"
RESULTS="${RESULTS:-$(dirname "$BUILD")/results-pgbench-$(basename "$BUILD")-$(date +%Y%m%d-%H%M%S).csv}"

[[ -x $BUILD/bin/postgres ]] || { echo "no $BUILD/bin/postgres" >&2; exit 1; }

start() {
  "$BUILD/bin/pg_ctl" -D "$DATA" -l "$LOG" -w -o \
    "-p $PORT -k /tmp -c listen_addresses= -c shared_buffers=1GB -c fsync=off \
     -c autovacuum=off -c jit=off -c max_parallel_workers_per_gather=0" start >/dev/null
}

stop() {
  "$BUILD/bin/pg_ctl" -D "$DATA" -m fast stop >/dev/null 2>&1 || true
}

trap stop EXIT

if [[ ! -d $DATA ]]; then
  "$BUILD/bin/initdb" -D "$DATA" -A trust >/dev/null
  start
  "$BUILD/bin/createdb" -h /tmp -p "$PORT" bench
  "$BUILD/bin/pgbench" -i -q -s "$SCALE" --partitions="$PARTITIONS" -h /tmp -p "$PORT" bench
  stop
fi

export PGOPTIONS="-c plan_cache_mode=force_generic_plan"

echo "build=$(basename "$BUILD") scale=$SCALE partitions=$PARTITIONS"
echo "round,clients,tps" >"$RESULTS"
for ((r = 1; r <= ROUNDS; r++)); do
  start
  # prewarm: one full pass over the accounts so the timed runs read from shared buffers
  "$BUILD/bin/psql" -qXAt -h /tmp -p "$PORT" -d bench -c "SELECT count(*) FROM pgbench_accounts" >/dev/null
  for c in $CLIENTS; do
    tps=$("$BUILD/bin/pgbench" -n -S -M prepared -c "$c" -j "$c" -T "$DURATION" \
      -h /tmp -p "$PORT" bench 2>&1 |
      sed -n 's/^tps = \([0-9.]*\).*/\1/p')
    echo "$r,$c,$tps" | tee -a "$RESULTS"
  done
  stop
done

echo
echo "== median tps over $ROUNDS rounds ($RESULTS)"
printf "%7s %12s\n" clients tps
awk -F, 'NR > 1 { v[$2] = v[$2] " " $3 }
  function median(s,  a, n, i, j, t) {
    n = split(s, a, " ")
    for (i = 2; i <= n; i++) for (j = i; j > 1 && a[j - 1] + 0 > a[j] + 0; j--) { t = a[j]; a[j] = a[j - 1]; a[j - 1] = t }
    return n % 2 ? a[(n + 1) / 2] : (a[n / 2] + a[n / 2 + 1]) / 2
  }
  END { for (c in v) printf "%7s %12.1f\n", c, median(v[c]) }' "$RESULTS" | sort -k1,1n