#!/bin/bash
#
# What WAL replay's prefetcher costs when the records ahead of it carry FPIs.
#
#   case1  insert-only load into a table with no index, crash recovery
#   case2  pgbench, checkpoint, update of 5% of pgbench_accounts
#   case3  vacuum of a table with a dead row on every page
#   case4  table at fillfactor 10, every row updated in place
#   case5  case2 with three more indexes on pgbench_accounts
#   case6  pgbench TPC-B run, random rows, 8 clients
#   case7  case6 with full_page_writes off: the feature's best case
#   case8  case7 replayed warm: the CPU cost of the look-ahead on its own
#
# servers: shared_buffers 128MB, fsync off, full_page_writes on (off in 7 and
# 8), wal_level replica, checkpoint_timeout 1h, max_wal_size 32GB, autovacuum
# off; replaying server adds recovery_prefetch off/on, hot_standby on, and
# pauses at the target LSN.  maintenance_io_concurrency is left at the
# server default, so the look-ahead is 4 x 16 = 64 block references.
#
set -euo pipefail

BIN=${BIN:?set BIN to the bin directory of the unpatched build}
V3_BIN=${V3_BIN:-}; HT_BIN=${HT_BIN:-}
SCALE=${SCALE:-1}; REPEAT=${REPEAT:-3}; PORT=${PORT:-5599}
WORK=${WORK:-$PWD/prefetch_cases_work}
SHARED_BUFFERS=${SHARED_BUFFERS:-128MB}
MIO=${MIO:-}
CASES=${CASES:-case1 case2 case3 case4 case5 case6 case7 case8}

die() { echo "error: $*" >&2; exit 1; }
[ "$(id -u)" != 0 ] || die "run as a non-root user"
for b in "$BIN" ${V3_BIN:+"$V3_BIN"} ${HT_BIN:+"$HT_BIN"}; do
  [ -x "$b/postgres" ] || die "no postgres in $b"
done
[ ! -e "/tmp/.s.PGSQL.$PORT" ] || die "port $PORT is already in use"
mkdir -p "$WORK"; WORK=$(cd "$WORK" && pwd)
ROWS=$WORK/rows.txt; : > "$ROWS"
export PGUSER=postgres PGDATABASE=postgres PGPORT=$PORT PGHOST=/tmp
export PATH=$BIN:$PATH
trap '[ -f "$WORK/primary/postmaster.pid" ] && pg_ctl -D "$WORK/primary" -m immediate stop >/dev/null 2>&1; true' EXIT

q() { psql -qc "$@"; }
v() { psql -Atc "$@"; }

FPW=on                              # cases 7 and 8 turn full_page_writes off
EVICT=yes                           # case 8 replays warm
MODE=standby                        # case 1 uses crash recovery
CASE="" A="" B="" START_LSN="" END_LSN=""

prep() {   # $1 = case: a fresh primary archiving into $WORK/archive_$1
  CASE=$1; A=$WORK/archive_$1; B=$WORK/base_$1
  rm -rf "$A" "$B" "$WORK/primary"; mkdir -p "$A"
  initdb -D "$WORK/primary" -A trust -U postgres >/dev/null
  cat >> "$WORK/primary/postgresql.conf" <<EOF
port = $PORT
unix_socket_directories = '/tmp'
listen_addresses = ''
shared_buffers = $SHARED_BUFFERS
fsync = off
full_page_writes = $FPW
wal_level = replica
archive_mode = on
archive_command = 'cp %p $A/%f'
checkpoint_timeout = 1h
max_wal_size = 32GB
autovacuum = off
${MIO:+maintenance_io_concurrency = $MIO}
EOF
  pg_ctl -D "$WORK/primary" -l "$WORK/primary.log" -w start >/dev/null
}

mark() {   # checkpoint, take the base backup, remember where replay starts
  q "CHECKPOINT"
  [ "$MODE" = crash ] || pg_basebackup -D "$B" -X none -c fast 2>/dev/null
  START_LSN=$(v "select pg_current_wal_lsn()")
}

fin() {   # $1 = description: archive the WAL, then replay it with every build
  local seg i wal_mb
  END_LSN=$(v "select pg_current_wal_lsn()")
  wal_mb=$(v "select (pg_wal_lsn_diff('$END_LSN', '$START_LSN') / 1048576)::int")
  seg=$(v "select pg_walfile_name('$END_LSN')")
  v "select pg_switch_wal()" >/dev/null
  until [ -f "$A/$seg" ]; do sleep 0.2; done
  sleep 1
  if [ "$MODE" = crash ]; then
    pg_ctl -D "$WORK/primary" -m immediate stop >/dev/null
    mv "$WORK/primary" "$B"           # replay the crashed directory itself
  else
    pg_ctl -D "$WORK/primary" -m fast stop >/dev/null
  fi
  echo "$CASE: $1, $wal_mb MB of WAL"
  for i in $(seq "$REPEAT"); do
    replay "$BIN" off off
    replay "$BIN" on  on
    [ -z "$V3_BIN" ] || replay "$V3_BIN" on v3
    [ -z "$HT_BIN" ] || replay "$HT_BIN" on HT
  done
}

replay() {   # $1 = bindir, $2 = off|on, $3 = label
  local bindir=$1 D=$WORK/replay spid="" mb=0 secs counters t0 t1 waited=0
  rm -rf "$D" "$D.log"; cp -a "$B" "$D"; rm -f "$D/postmaster.pid"
  cat > "$D/postgresql.auto.conf" <<EOF
port = $PORT
unix_socket_directories = '/tmp'
listen_addresses = ''
shared_buffers = $SHARED_BUFFERS
recovery_prefetch = $2
hot_standby = on
archive_mode = off
log_min_messages = log
log_timezone = 'UTC'
${MIO:+maintenance_io_concurrency = $MIO}
EOF
  local marker='redo done at'
  if [ "$MODE" = standby ]; then
    cat >> "$D/postgresql.auto.conf" <<EOF
restore_command = 'cat $A/%f > %p'
recovery_target_lsn = '$END_LSN'
recovery_target_action = 'pause'
EOF
    touch "$D/recovery.signal"
    marker='pausing at the end of recovery'
  fi
  # keep the WAL warm so what is measured is data reads.  cat, not cp: on a
  # reflink filesystem cp shares extents and never fills the page cache, and
  # the WAL reads would then land in the numbers below.
  cat "$A"/0* > /dev/null
  sync
  if [ "$EVICT" = yes ]; then   # drop the data files, keep pg_wal
    find "$D" -path "$D/pg_wal" -prune -o -type f \
      -exec dd if={} iflag=nocache count=0 status=none \;
  else
    find "$D" -type f -exec cat {} + > /dev/null
  fi

  "$bindir/pg_ctl" -D "$D" -l "$D.log" -W start >/dev/null
  until [ -f "$D/postmaster.pid" ]; do
    sleep 0.05; waited=$((waited + 1))
    [ $waited -lt 600 ] || die "server did not start, see $D.log"
  done
  until grep -q "$marker" "$D.log" 2>/dev/null; do
    # the startup process is a child of this postmaster: do not match other
    # servers on the machine, whose /proc we may not even be able to read
    [ -n "$spid" ] || spid=$(pgrep -P "$(head -1 "$D/postmaster.pid")" -f startup | head -1 || true)
    [ -f "$D/postmaster.pid" ] || die "server died during replay, see $D.log"
    sleep 0.1
  done
  [ -z "$spid" ] || mb=$(awk '/^read_bytes/{printf "%d", $2/1048576}' "/proc/$spid/io" 2>/dev/null || echo 0)
  if [ "$MODE" = standby ]; then
    t0=$(grep -m1 'redo starts at' "$D.log" | cut -c1-23)
    t1=$(grep -m1 "$marker" "$D.log" | cut -c1-23)
    secs=$(awk -v a="$(date -d "$t0" +%s.%N)" -v b="$(date -d "$t1" +%s.%N)" \
             'BEGIN {printf "%.2f", b - a}')
  else   # the crash case logs its own elapsed time when redo finishes
    secs=$(grep -o 'redo done at .*' "$D.log" | tail -1 \
             | sed -n 's/.*elapsed: \([0-9.]*\) s.*/\1/p')
  fi
  until "$bindir/pg_isready" -q; do sleep 0.2; done
  counters=$("$bindir/psql" -Atc "select prefetch||' '||hit||' '||skip_fpw||' '||skip_init||' '||skip_new||' '||skip_rep from pg_stat_recovery_prefetch")
  "$bindir/pg_ctl" -D "$D" -m fast stop >/dev/null
  rm -rf "$D"

  printf '%s %s %s %s %s\n' "$CASE" "$3" "$secs" "$mb" "$counters" >> "$ROWS"
  printf '  %-6s %-4s redo %7ss  read %6s MB  prefetch/hit/fpw/init/new/rep %s\n' \
    "$CASE" "$3" "$secs" "$mb" "$counters"
}

# the cases
PGB=$((25 * SCALE))                 # pgbench scale
TXNS=$((20000 * SCALE / 8))         # TPC-B transactions per client, 8 clients
UPD5="UPDATE pgbench_accounts SET abalance = abalance + 1 WHERE aid % 20 = 0"

run_case() {
  case $1 in case7 | case8) FPW=off ;; esac
  [ "$1" != case8 ] || EVICT=no
  prep "$1"; FPW=on
  case $1 in
    case1)
      # crash recovery: the data directory already holds the grown table, so
      # the blocks the WAL revisits exist and get prefetched.  From a base
      # backup taken before the inserts they would all be new, and skipped.
      MODE=crash
      q "CREATE TABLE test_noindex(id bigint, val1 int, val2 int, payload text)"
      mark
      while [ "$(v "select (pg_wal_lsn_diff(pg_current_wal_lsn(), '$START_LSN') / 1048576)::int")" -lt $((250 * SCALE)) ]; do
        q "INSERT INTO test_noindex SELECT g, (g*7+13)%100000, (g*31+17)%100000, repeat(chr(65+(g%26)),60) FROM generate_series(1, 500000) g"
      done
      fin "inserts into a table with no index" ;;
    case2)
      pgbench -i -s $PGB -q >/dev/null 2>&1
      mark; q "$UPD5"
      fin "pgbench scale $PGB, 5% update" ;;
    case3)
      q "CREATE TABLE chunks(id int, payload text);
         INSERT INTO chunks SELECT g, repeat('x', 100) FROM generate_series(1, $((1500000 * SCALE))) g;
         DELETE FROM chunks WHERE id % 20 = 0"
      mark; q "VACUUM chunks"
      fin "vacuum with a dead row on every page" ;;
    case4)
      q "CREATE TABLE sparse(id int, payload text) WITH (fillfactor = 10);
         INSERT INTO sparse SELECT g, repeat('x', 100) FROM generate_series(1, $((250000 * SCALE))) g"
      mark; q "UPDATE sparse SET id = id + 1"
      fin "rows at fillfactor 10 updated in place" ;;
    case5)
      pgbench -i -s $PGB -q >/dev/null 2>&1
      q "CREATE INDEX ON pgbench_accounts (bid);
         CREATE INDEX ON pgbench_accounts (abalance);
         CREATE INDEX ON pgbench_accounts (filler)"
      mark; q "$UPD5"
      fin "pgbench scale $PGB with 4 indexes, 5% update" ;;
    case6 | case7 | case8)
      pgbench -i -s $PGB -q >/dev/null 2>&1
      mark; pgbench -c 8 -j 2 -n -t $TXNS >/dev/null 2>&1
      fin "pgbench scale $PGB, TPC-B on 8 clients, full_page_writes $FPW" ;;
    *) die "unknown case: $1" ;;
  esac
  EVICT=yes; MODE=standby
}

for c in $CASES; do
  echo "$(date +%H:%M:%S) generating $c" >&2
  run_case "$c"
done

# the table
LABELS="off on"
[ -z "$V3_BIN" ] || LABELS="$LABELS v3"
[ -z "$HT_BIN" ] || LABELS="$LABELS HT"

med() {   # $1 = case, $2 = label, $3 = field: 3 = redo seconds, 5 = prefetches
  local n vals
  vals=$(awk -v c="$1" -v l="$2" -v f="$3" '$1==c && $2==l {print $f}' "$ROWS" | sort -n)
  n=$(echo "$vals" | grep -c .)
  [ "$n" -gt 0 ] && echo "$vals" | sed -n "$(( n / 2 + 1 ))p" || echo "-"
}

echo
echo "medians of $REPEAT runs; off/on = recovery_prefetch off/on, same build"
printf '\n%*s%s%*s%s\n' 23 "" "redo, s" 26 "" "prefetches"
printf '%-13s' "case"
for l in $LABELS; do printf '%7s' "$l"; done
printf '  '
for l in $LABELS; do printf '%9s' "$l"; done
echo
for c in $CASES; do
  printf '%-13s' "$c"
  for l in $LABELS; do printf '%7s' "$(med "$c" "$l" 3)"; done
  printf '  '
  for l in $LABELS; do printf '%9s' "$(med "$c" "$l" 5)"; done
  echo
done
echo
echo "raw rows: $ROWS"