#!/bin/bash
#
# 0003: Use WALReadFromBuffers() for local WAL reads.
#
# Emits WAL with pg_logical_emit_message() while one backend scans it back with
# pg_walinspect, which reads through read_local_xlog_page().
#
# Run once per build, same knobs both times:
#   BUILD=HEAD    PGBIN=~/pg/head/bin   RUNS=3 ./nocfbot-local_wal_read_perf.sh
#   BUILD=PATCHED PGBIN=~/pg/patched/bin RUNS=3 ./nocfbot-local_wal_read_perf.sh
#
# Needs pg_walinspect (make -C contrib/pg_walinspect install) and sysstat.

set -eo pipefail

PGBIN=${PGBIN:?set PGBIN to the bin directory of the build under test}
BUILD=${BUILD:-HEAD}
RUNS=${RUNS:-3}
DURATION=${DURATION:-180}
CLIENTS=${CLIENTS:-16}
DRAIN=${DRAIN:-30}
SCAN_SLEEP=${SCAN_SLEEP:-0}
DIRECT_IO=${DIRECT_IO:-"on off"}
DATA_ROOT=${DATA_ROOT:-/pgdata}
PGDATA=${PGDATA:-$DATA_ROOT/$BUILD/localread-data}
WALDIR=${WALDIR:-}
PORT=${PORT:-5432}
SHARED_BUFFERS=${SHARED_BUFFERS:-8GB}
WAL_BUFFERS=${WAL_BUFFERS:-16MB}
MAX_WAL_SIZE=${MAX_WAL_SIZE:-64GB}
OUTDIR=${OUTDIR:-$HOME/wal-read-perf/localread-$BUILD-$(date -u +%Y%m%d-%H%M%S)}

export PGHOST=/tmp PGPORT=$PORT PGDATABASE=postgres
PATH=$PGBIN:$PATH
mkdir -p "$OUTDIR"
CSV=$OUTDIR/runs.csv
SUMMARY=$OUTDIR/SUMMARY.txt
FMT='%-8s %-14s %-9s %-11s %-18s %-20s %-12s %-13s %-15s %-16s %s\n'

psql_q() { psql -Atq -F' ' -c "$1"; }
pg_start() {
  pg_ctl -D "$PGDATA" -l "$OUTDIR/server-$1.log" -w start >/dev/null ||
    { echo "server did not start, see $OUTDIR/server-$1.log"; exit 1; }
}
pg_stop() { pg_ctl -D "$PGDATA" -m fast -w stop >/dev/null 2>&1 || true; }
BG_PIDS=()
trap 'kill ${BG_PIDS[@]} 2>/dev/null || true; pg_stop' EXIT

# median of the per-run values: integers for counts and large values, more
# digits only where a small value would otherwise round away to nothing
median() { printf '%s\n' "$@" | sort -n |
  awk '{a[NR]=$1} END {m = (NR%2) ? a[(NR+1)/2] : (a[NR/2]+a[NR/2+1])/2
                       fmt = (m == int(m)) ? "%d\n" : (m >= 100) ? "%.0f\n" : (m >= 1) ? "%.1f\n" : "%.2f\n"
                       printf fmt, m}'; }

# peak read and write MB/s, skipping iostat's first since-boot report
iostat_peak() {
  [ -s "$1" ] || { echo "0 0"; return; }
  awk '/^Device/ { hdr++; for (i = 1; i <= NF; i++) { if ($i == "rMB/s") r = i; if ($i == "wMB/s") w = i } next }
       hdr >= 2 && r > 0 && NF > 3 { if ($r > pr) pr = $r; if ($w > pw) pw = $w }
       END { printf "%.2f %.2f\n", pr, pw }' "$1"
}

# assert builds and a build that overrides wal_buffers make the numbers useless
check_server() {
  [ "$(psql_q "SELECT current_setting('debug_assertions')")" = off ] || [ "${ALLOW_CASSERT:-0}" = 1 ] ||
    { echo "$PGBIN is an assert-enabled build, use a release build or set ALLOW_CASSERT=1"; exit 1; }
  [ "$(psql_q "SELECT pg_size_bytes(current_setting('wal_buffers')) = pg_size_bytes('$WAL_BUFFERS')")" = t ] ||
    { echo "wal_buffers is $(psql_q "SELECT current_setting('wal_buffers')"), not the requested $WAL_BUFFERS: this build forces it"; exit 1; }
}

setup_cluster() {
  if [ ! -s "$PGDATA/PG_VERSION" ]; then
    mkdir -p "$(dirname "$PGDATA")" "${WALDIR:-$PGDATA}"
    initdb -D "$PGDATA" --no-sync ${WALDIR:+-X "$WALDIR"} >"$OUTDIR/initdb.log" 2>&1
  fi
  grep -q "include 'bench.conf'" "$PGDATA/postgresql.conf" ||
    echo "include 'bench.conf'" >>"$PGDATA/postgresql.conf"
  [ -f "$(pg_config --pkglibdir)/pg_walinspect.so" ] ||
    { echo "pg_walinspect not installed: make -C contrib/pg_walinspect install"; exit 1; }
}

write_conf() { # $1 = debug_io_direct value ('wal' or '')
  cat >"$PGDATA/bench.conf" <<EOF
port = $PORT
unix_socket_directories = '/tmp'
shared_buffers = $SHARED_BUFFERS
wal_buffers = $WAL_BUFFERS
max_wal_size = $MAX_WAL_SIZE
synchronous_commit = on
wal_level = logical
track_io_timing = on
debug_io_direct = '$1'
EOF
}

detect_wal_dev() {
  local src
  src=$(df --output=source "$PGDATA/pg_wal" | tail -1)
  lsblk -no PKNAME "$src" 2>/dev/null | head -1 | grep . || basename "$src"
}

sample_lag() { # $1 = output file; how far the reader trails the insert point
  while :; do
    psql_q "SELECT round(pg_wal_lsn_diff(pg_current_wal_lsn(), lsn)/1024.0, 1) FROM wal_scan_pos" 2>/dev/null
    sleep 5
  done >"$1"
}

# the reader: one forward step from where the last step stopped to the flush point
setup_reader() {
  psql_q "CREATE EXTENSION IF NOT EXISTS pg_walinspect" >/dev/null
  psql_q "CREATE UNLOGGED TABLE IF NOT EXISTS wal_scan_pos (lsn pg_lsn)" >/dev/null
  psql_q "CREATE OR REPLACE FUNCTION wal_scan_step() RETURNS bigint LANGUAGE plpgsql AS \$fn\$
    DECLARE s pg_lsn; e pg_lsn; n bigint := 0;
    BEGIN
      SELECT lsn INTO s FROM wal_scan_pos FOR UPDATE;
      e := pg_current_wal_flush_lsn();
      IF e > s THEN
        -- the flush point can land mid record, in which case there is no
        -- complete record after s yet; leave the cursor and retry later
        BEGIN
          SELECT coalesce(sum(count), 0) INTO n FROM pg_get_wal_stats(s, e);
          UPDATE wal_scan_pos SET lsn = e;
        EXCEPTION WHEN others THEN
          PERFORM pg_sleep(0.001);
        END;
      ELSE
        PERFORM pg_sleep(0.001);   -- nothing new yet, do not spin
      END IF;
      RETURN n;
    END \$fn\$" >/dev/null
  psql_q "TRUNCATE wal_scan_pos" >/dev/null
  psql_q "INSERT INTO wal_scan_pos VALUES (pg_current_wal_flush_lsn())" >/dev/null
}

run_one() { # $1 = direct io label (on|off), $2 = run number
  local dio=$1 run=$2 tag="$BUILD-dio$1-run$2" start_lsn ios= dec lag

  setup_reader
  psql_q "CHECKPOINT" >/dev/null
  psql_q "SELECT pg_stat_reset_shared('io')" >/dev/null
  start_lsn=$(psql_q "SELECT pg_current_wal_lsn()")

  if [ "$WAL_DEV" != none ]; then
    iostat -dxm 1 "$WAL_DEV" >"$OUTDIR/iostat-$tag.log" 2>/dev/null &
    ios=$!; disown $ios 2>/dev/null || true
  fi
  sample_lag "$OUTDIR/lag-$tag.log" &
  lag=$!; disown $lag 2>/dev/null || true
  # the reader keeps scanning for DRAIN seconds after the emitters stop
  pgbench -n -f "$OUTDIR/scan.sql" -c 1 -j 1 -T $((DURATION + DRAIN)) \
    >"$OUTDIR/scan-$tag.txt" 2>&1 &
  dec=$!
  BG_PIDS=($ios $lag $dec)
  sleep 2

  pgbench -n -f "$OUTDIR/emit.sql" -c "$CLIENTS" -j "$CLIENTS" -T "$DURATION" -P 10 \
    >"$OUTDIR/pgbench-$tag.txt" 2>&1

  local wal_mb
  wal_mb=$(psql_q "SELECT round(pg_wal_lsn_diff(pg_current_wal_lsn(), '$start_lsn')/1048576.0, 1)")
  wait $dec 2>/dev/null || true

  kill $ios $lag 2>/dev/null || true
  BG_PIDS=()
  sleep 3 # let the exited backends' stats land

  local tps dtps read_mb hits reads pr pw lag_kb
  tps=$(awk '/^tps =/ {print $3}' "$OUTDIR/pgbench-$tag.txt" | tail -1)
  dtps=$(awk '/^tps =/ {print $3}' "$OUTDIR/scan-$tag.txt" | tail -1)
  read -r read_mb hits reads <<<"$(psql_q "SELECT round(coalesce(sum(read_bytes), 0)/1048576.0, 1), coalesce(sum(hits), 0), coalesce(sum(reads), 0) FROM pg_stat_io WHERE backend_type = 'client backend' AND object = 'wal'")"
  read -r pr pw <<<"$(iostat_peak "$OUTDIR/iostat-$tag.log")"
  lag_kb=$(sort -n "$OUTDIR/lag-$tag.log" 2>/dev/null | tail -1)
  : "${lag_kb:=0}"

  echo "$BUILD,$dio,$run,$tps,$dtps,$wal_mb,$read_mb,$reads,$hits,$pr,$pw,$lag_kb" | tee -a "$CSV"

  R_TPS+=("$tps"); R_DTPS+=("$dtps"); R_WAL+=("$wal_mb"); R_READ+=("$read_mb")
  R_READS+=("$reads"); R_HITS+=("$hits"); R_DR+=("$pr"); R_DW+=("$pw"); R_LAG+=("$lag_kb")
}

echo "SELECT pg_logical_emit_message(true, 'bench', repeat('x', 350));" >"$OUTDIR/emit.sql"
echo "SELECT wal_scan_step();" >"$OUTDIR/scan.sql"
[ "$SCAN_SLEEP" = 0 ] || echo "SELECT pg_sleep($SCAN_SLEEP);" >>"$OUTDIR/scan.sql"

setup_cluster
WAL_DEV=${WAL_DEV:-$(detect_wal_dev)}
[ -n "$WAL_DEV" ] || { echo "could not detect the pg_wal device, set WAL_DEV"; exit 1; }
[ "$WAL_DEV" = none ] || command -v iostat >/dev/null ||
  { echo "iostat not found: install sysstat, or set WAL_DEV=none to skip disk sampling"; exit 1; }
echo "build=$BUILD pgdata=$PGDATA wal_dev=$WAL_DEV runs=$RUNS duration=${DURATION}s clients=$CLIENTS out=$OUTDIR"

echo "build,direct_io,run,emit_tps,scans_per_sec,wal_generated_mb,reader_read_mb,reads,hits,disk_read_peak_mbs,disk_write_peak_mbs,lag_kb" >"$CSV"
{
  echo "local WAL reads via pg_walinspect forward scan, build=$BUILD"
  echo "shared_buffers=$SHARED_BUFFERS wal_buffers=$WAL_BUFFERS max_wal_size=$MAX_WAL_SIZE synchronous_commit=on"
  echo "workload: pgbench -n -f emit.sql -c $CLIENTS -j $CLIENTS -T $DURATION, reader: pg_get_wal_stats() forward scan (1 client, +${DRAIN}s drain)"
  echo "runs=$RUNS (median reported), wal device=$WAL_DEV"
  echo "binary: $PGBIN, $(pg_config --version)"
  echo
  printf "$FMT" build "WAL direct IO" TPS "scans/s" "WAL generated MB" "reader read MB" \
    "file reads" "buffer hits" "WAL-disk reads" "WAL-disk writes" "reader lag"
} >"$SUMMARY"

for dio in $DIRECT_IO; do
  if [ "$dio" = on ]; then write_conf wal; else write_conf ''; fi
  pg_stop
  pg_start "$BUILD-dio$dio"
  check_server
  # report what the server actually used, a build may override these
  psql_q "SELECT name || ' = ' || current_setting(name) FROM pg_settings WHERE name IN ('shared_buffers', 'wal_buffers', 'max_wal_size', 'synchronous_commit', 'wal_level', 'debug_io_direct', 'summarize_wal', 'wal_writer_delay')" >"$OUTDIR/settings-$dio.txt"
  echo "direct IO $dio: effective wal_buffers=$(psql_q "SELECT current_setting('wal_buffers')") debug_io_direct='$(psql_q "SELECT current_setting('debug_io_direct')")'"

  R_TPS=(); R_DTPS=(); R_WAL=(); R_READ=(); R_READS=(); R_HITS=(); R_DR=(); R_DW=(); R_LAG=()
  for run in $(seq 1 "$RUNS"); do run_one "$dio" "$run"; done

  printf "$FMT" "$BUILD" "$dio" "$(median "${R_TPS[@]}")" "$(median "${R_DTPS[@]}")" \
    "$(median "${R_WAL[@]}")" "$(median "${R_READ[@]}")" "$(median "${R_READS[@]}")" \
    "$(median "${R_HITS[@]}")" "$(median "${R_DR[@]}") MB/s" "$(median "${R_DW[@]}") MB/s" \
    "$(median "${R_LAG[@]}") KB" >>"$SUMMARY"
  pg_stop
done

cat "$SUMMARY"
echo "per-run rows and effective settings in $OUTDIR"