#!/usr/bin/env bash

set -euo pipefail

case "${1:-}" in
	--syncdatadirectory) test_case=$1; metric="startup" ;;
	--checkpointer) test_case=$1; metric="checkpoint" ;;
	--slru) test_case=$1; metric="SLRU checkpoint" ;;
	*)
		echo "usage: $0 --syncdatadirectory | --checkpointer | --slru"
		exit 1
		;;
esac

BINDIR=${BINDIR:-$(pg_config --bindir)}
BENCH_DIR=${BENCH_DIR:-$HOME/aio-fsync-repro}
NTABLES=${NTABLES:-4000}
SLRU_XACTS=${SLRU_XACTS:-640000}
CLIENTS=${CLIENTS:-16}
PORT=${PGPORT:-55439}

PGDATA=$BENCH_DIR/data
SOCKET_DIR=$BENCH_DIR/socket

PG_CTL=$BINDIR/pg_ctl
INITDB=$BINDIR/initdb
PSQL=$BINDIR/psql
PGBENCH=$BINDIR/pgbench

stop_server()
{
	if [ -s "$PGDATA/postmaster.pid" ]; then
		"$PG_CTL" -D "$PGDATA" -m "${1:-fast}" -w stop >/dev/null
	fi
}

start_server()
{
	local method=$1
	local fsync=$2

	"$PG_CTL" -D "$PGDATA" -l "$BENCH_DIR/server.log" -w start \
		-o "-c io_method=$method -c fsync=$fsync" >/dev/null
}

psql()
{
	"$PSQL" -X -q -v ON_ERROR_STOP=1 -h "$SOCKET_DIR" -p "$PORT" \
		-d postgres "$@"
}

time_ms()
{
	local start end

	start=$(date +%s%N)
	"$@" >/dev/null
	end=$(date +%s%N)
	awk -v start="$start" -v end="$end" \
		'BEGIN {printf "%.1f", (end - start) / 1000000}'
}

dirty_relations()
{
	psql -c 'UPDATE bench SET value = value + 1' >/dev/null
}

create_slru_segments()
{
	PGOPTIONS='-c synchronous_commit=off' \
		"$PGBENCH" -n -M prepared -c "$CLIENTS" -j "$CLIENTS" \
		-t "$((SLRU_XACTS / CLIENTS))" \
		-f <(echo 'SELECT pg_current_xact_id();') \
		-h "$SOCKET_DIR" -p "$PORT" postgres >/dev/null
}

[ ! -e "$BENCH_DIR" ] || {
	echo "$BENCH_DIR already exists" >&2
	exit 1
}
mkdir -p "$SOCKET_DIR"
trap 'stop_server immediate || true; rm -rf -- "$BENCH_DIR"' EXIT
"$INITDB" -D "$PGDATA" -N >/dev/null

cat >> "$PGDATA/postgresql.conf" <<EOF
port = $PORT
unix_socket_directories = '$SOCKET_DIR'
checkpoint_timeout = '1h'
max_wal_size = '32GB'
checkpoint_flush_after = 0
bgwriter_lru_maxpages = 0
autovacuum = off
EOF

if [ "$test_case" = "--slru" ]; then
	echo "track_commit_timestamp = on" >> "$PGDATA/postgresql.conf"
fi

start_server sync on
if [ "$test_case" = "--slru" ]; then
	psql -c 'SELECT pg_current_xact_id()' >/dev/null
	psql -c CHECKPOINT >/dev/null
	echo "$SLRU_XACTS transactions per method"
else
	psql >/dev/null <<SQL
CREATE TABLE bench (id int, value int) PARTITION BY RANGE (id);
DO \$\$
BEGIN
	FOR i IN 1..$NTABLES LOOP
		EXECUTE format(
			'CREATE TABLE bench_%s PARTITION OF bench FOR VALUES FROM (%s) TO (%s)',
			i, i, i + 1);
	END LOOP;
END
\$\$;
INSERT INTO bench SELECT i, 0 FROM generate_series(1, $NTABLES) AS i;
CHECKPOINT;
SQL
	echo "$NTABLES relation files"
fi
stop_server

echo
printf '%-10s | %20s | %15s\n' "io_method" "$metric (ms)" "speedup vs sync"
printf '%-10s-+-%20s-+-%15s\n' "----------" "--------------------" "---------------"

for method in sync worker io_uring; do
	if [ "$test_case" = "--syncdatadirectory" ]; then
		# Write dirty buffers without syncing them, then crash.
		start_server sync off
		dirty_relations
		psql -c CHECKPOINT >/dev/null
		stop_server immediate

		elapsed=$(time_ms start_server "$method" on)
		stop_server
	elif [ "$test_case" = "--checkpointer" ]; then
		start_server "$method" on
		dirty_relations
		elapsed=$(time_ms psql -c CHECKPOINT)
		stop_server
	else
		start_server "$method" on
		create_slru_segments
		elapsed=$(time_ms psql -c CHECKPOINT)
		stop_server
	fi

	if [ "$method" = "sync" ]; then
		sync_elapsed=$elapsed
	fi

	speedup=$(awk -v baseline="$sync_elapsed" -v elapsed="$elapsed" \
		'BEGIN {printf "%.2fx", baseline / elapsed}')
	printf '%-10s | %20.1f | %15s\n' \
		"$method" "$elapsed" "$speedup"
done