#!/usr/bin/env bash
set -euo pipefail

# 仅执行 pgbench 的脚本（不安装 PostgreSQL，不 initdb，不启停服务，不改配置）。
#
# 前提条件：
# 1) PostgreSQL 已经在运行。
# 2) 测试库已经存在（默认 benchdb）。
# 3) PostgreSQL 已经在外部配置为禁用 mmap 共享内存：
#    shared_memory_type = sysv
#    dynamic_shared_memory_type = sysv
#
# 关于 Malloc0：
# - 当前后端是内存虚拟设备，容量有限。
# - 默认参数故意设置得较小，避免一次灌入过多数据。
#
# 关于 LD_PRELOAD：
# - USE_LD_PRELOAD_INIT=1：初始化阶段（pgbench -i）启用 LD_PRELOAD
# - USE_LD_PRELOAD_RUN=1 ：压测阶段启用 LD_PRELOAD
# - 设为 0 即可关闭对应阶段的 LD_PRELOAD
#
# 用法：
#   bash codex/run_pgbench_no_mmap.sh
#
# 可选环境变量（含义）：
#   PG_HOST=127.0.0.1
#     PostgreSQL 服务器地址。
#   PG_PORT=5432
#     PostgreSQL 服务器端口（默认改为 5432）。
#   PG_DB=benchdb
#     压测数据库名。
#   PG_SCALE=2
#     pgbench 初始化规模因子（-s），越大初始数据越多。
#   PG_TIME=20
#     压测持续时间（秒，pgbench -T）。
#   PG_CLIENTS=2
#     并发客户端数（pgbench -c）。
#   PG_JOBS=2
#     工作线程数（pgbench -j）。
#   PG_SUPERUSER=postgres
#     执行 pgbench 的系统用户（通常是 postgres）。
#   LD_PRELOAD_PATH=/home/lian/try/zvfs/src/libzvfs.so
#     LD_PRELOAD 目标库路径（你的 zvfs hook so）。
#   PG_BIN_DIR=/usr/lib/postgresql/16/bin
#     pgbench 所在目录；不填时自动从 PATH 查找。
#   USE_LD_PRELOAD_INIT=1
#     初始化阶段（pgbench -i）是否启用 LD_PRELOAD：1=启用，0=关闭。
#   USE_LD_PRELOAD_RUN=1
#     压测阶段是否启用 LD_PRELOAD：1=启用，0=关闭。

PG_HOST="${PG_HOST:-127.0.0.1}"
PG_PORT="${PG_PORT:-5432}"
PG_DB="${PG_DB:-benchdb}"
PG_SCALE="${PG_SCALE:-2}"
PG_TIME="${PG_TIME:-20}"
PG_CLIENTS="${PG_CLIENTS:-2}"
PG_JOBS="${PG_JOBS:-2}"
PG_SUPERUSER="${PG_SUPERUSER:-postgres}"
LD_PRELOAD_PATH="${LD_PRELOAD_PATH:-/home/lian/try/zvfs/src/libzvfs.so}"
PG_BIN_DIR="${PG_BIN_DIR:-$(dirname "$(command -v pgbench 2>/dev/null || true)")}"
USE_LD_PRELOAD_INIT="${USE_LD_PRELOAD_INIT:-1}"
USE_LD_PRELOAD_RUN="${USE_LD_PRELOAD_RUN:-1}"

if [[ -z "${PG_BIN_DIR}" || ! -x "${PG_BIN_DIR}/pgbench" ]]; then
  echo "未找到 pgbench，请设置 PG_BIN_DIR 或把 pgbench 放到 PATH 中。" >&2
  exit 1
fi

run_pgbench_cmd() {
  local use_preload="$1"
  shift
  if [[ "${use_preload}" == "1" ]]; then
    sudo -u "${PG_SUPERUSER}" env LD_PRELOAD="${LD_PRELOAD_PATH}" "$@"
  else
    sudo -u "${PG_SUPERUSER}" "$@"
  fi
}

echo "当前参数："
echo "  host=${PG_HOST} port=${PG_PORT} db=${PG_DB}"
echo "  scale=${PG_SCALE} clients=${PG_CLIENTS} jobs=${PG_JOBS} time=${PG_TIME}s"
echo "  preload_init=${USE_LD_PRELOAD_INIT} preload_run=${USE_LD_PRELOAD_RUN}"

echo "[1/2] 初始化数据（pgbench -i）"
run_pgbench_cmd "${USE_LD_PRELOAD_INIT}" \
  "${PG_BIN_DIR}/pgbench" -h "${PG_HOST}" -p "${PG_PORT}" -i -s "${PG_SCALE}" "${PG_DB}"

echo "[2/2] 执行压测（pgbench -T）"
run_pgbench_cmd "${USE_LD_PRELOAD_RUN}" \
  "${PG_BIN_DIR}/pgbench" -h "${PG_HOST}" -p "${PG_PORT}" \
  -c "${PG_CLIENTS}" -j "${PG_JOBS}" -T "${PG_TIME}" -P 5 "${PG_DB}"