#!/usr/bin/env bash
# Replay hint-bit FPIs, using the workload from the August 2026 results.
# Usage: bash wal-zstd-replay-bench.sh /path/to/install/bin [zstd|lz4]
# Run separately with the unpatched and patched binaries, repeating both.
# lz4 is an unchanged control; the patch only optimizes zstd.
# Build with -O2 and without assertions for performance measurements.
#
# Each invocation creates a new cluster and keeps its data and logs.
# BENCH_TMPDIR selects the storage (default /tmp; not necessarily tmpfs).
# BENCH_ROWS and BENCH_SHARED_BUFFERS can reduce the size for a smoke test.
# On a shared benchmark host, acquire its benchmark lock before running.
set -euo pipefail
export LANG=C LC_ALL=C

if [[ $# -lt 1 || $# -gt 2 ]]; then
	echo "Usage: $0 /path/to/install/bin [zstd|lz4]" >&2
	exit 1
fi
bench_bin=$(cd "$1" && pwd)
bench_method=${2:-zstd}
bench_rows=${BENCH_ROWS:-6000000}
bench_buffers=${BENCH_SHARED_BUFFERS:-8GB}
case "$bench_method" in
	zstd|lz4) ;;
	*) echo "Expected zstd or lz4" >&2; exit 1 ;;
esac
if [[ ! "$bench_rows" =~ ^[1-9][0-9]*$ ]]; then
	echo "BENCH_ROWS must be a positive integer" >&2
	exit 1
fi

bench_dir=$(mktemp -d "${BENCH_TMPDIR:-/tmp}/wal-zstd-replay.XXXXXX")
bench_data="$bench_dir/data"
# Keep the socket path short even when the data directory path is long.
bench_socket=$(mktemp -d /tmp/wal-zstd-socket.XXXXXX)
bench_port=54996
echo "Data and logs: $bench_dir" >&2
cleanup()
{
	if [[ -f "$bench_data/postmaster.pid" ]]; then
		"$bench_bin/pg_ctl" -D "$bench_data" stop -m immediate -w \
			>> "$bench_dir/pg_ctl.log" 2>&1 ||
			echo "Could not stop cluster at $bench_data" >&2
	fi
	rmdir "$bench_socket" 2>/dev/null || true
}
trap cleanup EXIT
trap 'exit 130' INT
trap 'exit 143' TERM

"$bench_bin/pg_config" --configure > "$bench_dir/build.txt"
"$bench_bin/postgres" --version >> "$bench_dir/build.txt"
"$bench_bin/initdb" -D "$bench_data" --no-sync --locale=C \
	--username=wal_bench > "$bench_dir/initdb.log" 2>&1
cat >> "$bench_data/postgresql.conf" <<EOF
listen_addresses = ''
port = $bench_port
unix_socket_directories = '$bench_socket'
fsync = off
synchronous_commit = off
shared_buffers = $bench_buffers
max_wal_size = 64GB
checkpoint_timeout = 1h
autovacuum = off
wal_compression = $bench_method
wal_log_hints = on
EOF

bench_psql=("$bench_bin/psql" -X -qAt -v ON_ERROR_STOP=1
	-h "$bench_socket" -p "$bench_port" -U wal_bench -d postgres)
"$bench_bin/pg_ctl" -D "$bench_data" -l "$bench_dir/generate.log" \
	start -w > "$bench_dir/pg_ctl.log" 2>&1
"${bench_psql[@]}" -c "SELECT name, setting FROM pg_settings WHERE name IN
	('debug_assertions', 'data_checksums', 'shared_buffers',
	 'max_parallel_workers_per_gather', 'wal_compression') ORDER BY name" \
	> "$bench_dir/settings.txt"
"${bench_psql[@]}" -c "CREATE TABLE t AS
	SELECT g AS id, repeat(md5(g::text), 5) AS pad
	FROM generate_series(1, $bench_rows) g" > /dev/null
"${bench_psql[@]}" -c 'CREATE TABLE flush_marker (id int)' > /dev/null
"${bench_psql[@]}" -c CHECKPOINT > /dev/null
bench_lsn0=$("${bench_psql[@]}" -c 'SELECT pg_current_wal_insert_lsn()')
"${bench_psql[@]}" -c 'SELECT count(*) FROM t' > "$bench_dir/count-before.txt"
bench_lsn1=$("${bench_psql[@]}" -c 'SELECT pg_current_wal_insert_lsn()')
bench_wal=$("${bench_psql[@]}" -c \
	"SELECT '$bench_lsn1'::pg_lsn - '$bench_lsn0'::pg_lsn")
if [[ "$bench_wal" -le 0 ]]; then
	echo "The scan generated no WAL" >&2
	exit 1
fi
# Read-only queries do not flush hint-bit WAL at commit. Make sure the
# scan's WAL is written before immediate shutdown, even in a short run.
"${bench_psql[@]}" -c 'SET synchronous_commit = on' \
	-c 'INSERT INTO flush_marker VALUES (1)' > /dev/null
printf 'start_lsn=%s end_lsn=%s\n' "$bench_lsn0" "$bench_lsn1" \
	> "$bench_dir/wal-range.txt"
"$bench_bin/pg_ctl" -D "$bench_data" stop -m immediate -w \
	>> "$bench_dir/pg_ctl.log" 2>&1
"$bench_bin/pg_ctl" -D "$bench_data" -l "$bench_dir/recovery.log" \
	start -w >> "$bench_dir/pg_ctl.log" 2>&1
bench_elapsed=$(sed -n \
	's/.*redo done at .*elapsed: \([0-9.]*\) s.*/\1/p' \
	"$bench_dir/recovery.log")
if [[ ! "$bench_elapsed" =~ ^[0-9]+([.][0-9]+)?$ ]]; then
	echo "No unique redo duration in $bench_dir/recovery.log" >&2
	exit 1
fi
"${bench_psql[@]}" -c 'SELECT count(*) FROM t' > "$bench_dir/count-after.txt"
cmp "$bench_dir/count-before.txt" "$bench_dir/count-after.txt"
if [[ $("${bench_psql[@]}" -c 'SELECT count(*) FROM flush_marker') != 1 ]]; then
	echo "Recovery did not reach the flush marker" >&2
	exit 1
fi
printf 'method=%s rows=%s wal_bytes=%s redo_elapsed_s=%s\n' \
	"$bench_method" "$bench_rows" "$bench_wal" "$bench_elapsed" \
	| tee "$bench_dir/result.txt"