#!/usr/bin/env bash

set -euo pipefail

usage()
{
	cat <<'USAGE'
Usage:
  pgstat_xact_macro_bench.sh [options]

Options:
  --pg-bin DIR              Run one PostgreSQL install.
  --patched-pg-bin DIR      Run and label this install as patched.
  --unpatched-pg-bin DIR    Run and label this install as unpatched.
  --tables 5000             Number of tenant tables to create.
  --rows-per-table 1        Rows inserted into each tenant table.
  --clients 1               pgbench clients. Use 1 for one long-lived backend.
  --jobs 1                  pgbench worker threads.
  --time 60                 Measured pgbench duration in seconds.
  --warmup 10               Warmup duration in seconds.
  --port-base 65440         First temporary server port.
  --keep                    Keep temporary cluster directories.

Environment:
  PG_BIN                    Same as --pg-bin.

Examples:
  ./pgstat_xact_macro_bench.sh --pg-bin ./inst/bin

  ./pgstat_xact_macro_bench.sh \
      --unpatched-pg-bin /tmp/pg-unpatched/inst/bin \
      --patched-pg-bin /tmp/pg-patched/inst/bin

When both --patched-pg-bin and --unpatched-pg-bin are provided, runs are
executed in the same order as the options appear on the command line.  Use
that to check for order bias.

This benchmark models a realistic bad case for transaction-boundary pgstat
work: many tenant tables, one long-lived connection, many short transactions,
and pending stats entries retained in the backend while the workload keeps
running.  The pgbench script is:

  \set id random(1, N)
  BEGIN;
  SELECT count(*) FROM tenant_:id;
  COMMIT;
USAGE
}

pg_bin="${PG_BIN:-}"
patched_pg_bin=""
unpatched_pg_bin=""
run_labels=()
run_bins=()
tables=5000
rows_per_table=1
clients=1
jobs=1
duration=60
warmup=10
port_base=65440
keep=0

while [ "$#" -gt 0 ]; do
	case "$1" in
		--pg-bin)
			pg_bin="$2"
			shift 2
			;;
		--patched-pg-bin)
			patched_pg_bin="$2"
			run_labels+=("patched")
			run_bins+=("$2")
			shift 2
			;;
		--unpatched-pg-bin)
			unpatched_pg_bin="$2"
			run_labels+=("unpatched")
			run_bins+=("$2")
			shift 2
			;;
		--tables)
			tables="$2"
			shift 2
			;;
		--rows-per-table)
			rows_per_table="$2"
			shift 2
			;;
		--clients)
			clients="$2"
			shift 2
			;;
		--jobs)
			jobs="$2"
			shift 2
			;;
		--time)
			duration="$2"
			shift 2
			;;
		--warmup)
			warmup="$2"
			shift 2
			;;
		--port-base)
			port_base="$2"
			shift 2
			;;
		--keep)
			keep=1
			shift
			;;
		--help|-h)
			usage
			exit 0
			;;
		*)
			echo "unknown option: $1" >&2
			usage >&2
			exit 2
			;;
	esac
done

if [ -n "$pg_bin" ]; then
	if [ -n "$patched_pg_bin" ] || [ -n "$unpatched_pg_bin" ]; then
		echo "--pg-bin cannot be combined with --patched-pg-bin or --unpatched-pg-bin" >&2
		exit 1
	fi
	run_labels=("build")
	run_bins=("$pg_bin")
fi

if [ "${#run_bins[@]}" -eq 0 ]; then
	echo "set PG_BIN or pass --pg-bin/--patched-pg-bin/--unpatched-pg-bin" >&2
	exit 1
fi

for bin in "${run_bins[@]}"; do
	for prog in postgres initdb pg_ctl psql pgbench; do
		if [ ! -x "$bin/$prog" ]; then
			echo "missing executable: $bin/$prog" >&2
			exit 1
		fi
	done
done

pg_tool()
{
	local bin="$1"
	local prog="$2"
	local libdir

	shift 2
	libdir="$(cd "$bin/../lib" 2>/dev/null && pwd || true)"
	if [ -z "$libdir" ]; then
		"$bin/$prog" "$@"
		return
	fi
	case "$(uname -s)" in
		Darwin)
			DYLD_LIBRARY_PATH="$libdir${DYLD_LIBRARY_PATH:+:$DYLD_LIBRARY_PATH}" \
				"$bin/$prog" "$@"
			;;
		*)
			LD_LIBRARY_PATH="$libdir${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" \
				"$bin/$prog" "$@"
			;;
	esac
}

tmpdirs=()
datadirs=()
pg_bins=()
created_tmpdir=""
created_script_file=""
created_conninfo=""

cleanup()
{
	set +e
	local i

	for i in "${!datadirs[@]}"; do
		if [ -d "${datadirs[$i]}" ]; then
			pg_tool "${pg_bins[$i]}" pg_ctl -D "${datadirs[$i]}" -m fast -w stop \
				>/dev/null 2>&1
		fi
	done

	if [ "$keep" -eq 0 ]; then
		for i in "${!tmpdirs[@]}"; do
			rm -rf "${tmpdirs[$i]}"
		done
	else
		for i in "${!tmpdirs[@]}"; do
			echo "kept temporary directory: ${tmpdirs[$i]}"
		done
	fi
}
trap cleanup EXIT
trap 'trap - EXIT; cleanup; exit 130' INT
trap 'trap - EXIT; cleanup; exit 143' TERM

create_cluster()
{
	local label="$1"
	local bin="$2"
	local port="$3"
	local tmpdir datadir logfile setup_sql script_file conninfo

	tmpdir="$(mktemp -d "${TMPDIR:-/tmp}/pgstat-xact-macro-${label}.XXXXXX")"
	datadir="$tmpdir/data"
	logfile="$tmpdir/postgres.log"
	setup_sql="$tmpdir/setup.sql"
	script_file="$tmpdir/tenant.sql"

	tmpdirs+=("$tmpdir")
	datadirs+=("$datadir")
	pg_bins+=("$bin")

	pg_tool "$bin" initdb -D "$datadir" --auth trust --no-sync --no-instructions \
		--lc-messages=C >/dev/null
	cat >>"$datadir/postgresql.conf" <<EOF
listen_addresses = ''
port = $port
unix_socket_directories = '$tmpdir'
fsync = off
synchronous_commit = off
autovacuum = off
track_functions = 'none'
max_prepared_transactions = 0
EOF
	pg_tool "$bin" pg_ctl -D "$datadir" -l "$logfile" -w start >/dev/null

	{
		echo "SET client_min_messages = warning;"
		echo "DROP SCHEMA IF EXISTS pgstat_xact_macro CASCADE;"
		echo "CREATE SCHEMA pgstat_xact_macro;"
		echo "SET search_path = pgstat_xact_macro;"
		for i in $(seq 1 "$tables"); do
			printf "CREATE TABLE tenant_%d (a int);\n" "$i"
			printf "INSERT INTO tenant_%d SELECT generate_series(1, %d);\n" \
				"$i" "$rows_per_table"
		done
	} >"$setup_sql"

	pg_tool "$bin" psql -X -q -v ON_ERROR_STOP=1 \
		-h "$tmpdir" -p "$port" -d postgres -f "$setup_sql" >/dev/null

	cat >"$script_file" <<EOF
\\set id random(1, $tables)
BEGIN;
SELECT count(*) FROM pgstat_xact_macro.tenant_:id;
COMMIT;
EOF

	conninfo="host=$tmpdir port=$port dbname=postgres"
	created_tmpdir="$tmpdir"
	created_script_file="$script_file"
	created_conninfo="$conninfo"
}

run_pgbench()
{
	local label="$1"
	local bin="$2"
	local conninfo="$3"
	local script_file="$4"
	local seconds="$5"
	local out_file="$6"

	pg_tool "$bin" pgbench -n -r -M simple -c "$clients" -j "$jobs" -T "$seconds" \
		-f "$script_file" "$conninfo" >"$out_file"
}

extract_metric()
{
	local file="$1"
	local pattern="$2"

	awk -v pat="$pattern" '
		$0 ~ pat {
			for (i = 1; i <= NF; i++)
			{
				if ($i ~ /^[0-9]+([.][0-9]+)?$/)
				{
					print $i;
					exit;
				}
			}
		}
	' "$file"
}

printf "tables=%s rows_per_table=%s clients=%s jobs=%s warmup=%s time=%s\n" \
	"$tables" "$rows_per_table" "$clients" "$jobs" "$warmup" "$duration"
printf "%-12s %12s %14s %14s\n" \
	"label" "tps" "avg_lat_ms" "select_avg_ms"

for idx in "${!run_bins[@]}"; do
	label="${run_labels[$idx]}"
	bin="${run_bins[$idx]}"
	port=$((port_base + idx))
	create_cluster "$label" "$bin" "$port"
	tmpdir="$created_tmpdir"
	script_file="$created_script_file"
	conninfo="$created_conninfo"
	warmup_out="$tmpdir/warmup.out"
	run_out="$tmpdir/pgbench.out"

	if [ "$warmup" -gt 0 ]; then
		run_pgbench "$label" "$bin" "$conninfo" "$script_file" "$warmup" \
			"$warmup_out"
	fi

	run_pgbench "$label" "$bin" "$conninfo" "$script_file" "$duration" \
		"$run_out"

	tps="$(extract_metric "$run_out" "^tps =")"
	avg_lat="$(extract_metric "$run_out" "^latency average =")"
	select_avg="$(awk '
		/SELECT count\(\*\) FROM pgstat_xact_macro[.]tenant_:id/ {
			print $1;
			exit;
		}
	' "$run_out")"

	if [ -z "$tps" ]; then
		tps="?"
	fi
	if [ -z "$avg_lat" ]; then
		avg_lat="?"
	fi
	if [ -z "$select_avg" ]; then
		select_avg="?"
	fi

	printf "%-12s %12s %14s %14s\n" \
		"$label" "$tps" "$avg_lat" "$select_avg"

	if [ "$keep" -eq 1 ]; then
		echo "temporary directory for $label: $tmpdir"
		echo "pgbench output for $label: $run_out"
	fi
done