#!/bin/bash
# =============================================================================
# 单库并发能力实测 v2：使用持久连接（pgbench 风格）
#
# 【为什么必须用持久连接】
#   v1 每次查询都新建 psql 进程，测到的是「进程启动开销」而非数据库能力，
#   TPS 只有 615，严重低估。生产用连接池（HikariCP），连接复用。
#
# 本脚本用单连接循环 N 次，再乘以并发数，更接近生产形态。
# =============================================================================
set -u
DB=rental_bench
ITER=${1:-200}
CONCURRENCY=${2:-16}

echo "=============================================="
echo " 单库并发能力实测（持久连接）"
echo " 每连接迭代: $ITER   并发连接: $CONCURRENCY"
echo "=============================================="
echo ""

# 用单个 psql 会话执行 N 次循环，测「连接内」真实吞吐
bench_persistent() {
  local name="$1"; local sql="$2"
  local tmpdir=$(mktemp -d)

  for i in $(seq 1 "$CONCURRENCY"); do
    (
      # psql 的 \timing + 循环：单连接内连续执行
      {
        echo "\\timing off"
        for n in $(seq 1 "$ITER"); do echo "$sql;"; done
      } | psql -d "$DB" -tA -q > /dev/null 2>&1
      echo "done" > "$tmpdir/$i"
    ) &
  done

  local start=$(python3 -c 'import time;print(time.time())')
  wait
  local end=$(python3 -c 'import time;print(time.time())')

  local elapsed=$(python3 -c "print(f'{$end-$start:.2f}')")
  local total=$(( ITER * CONCURRENCY ))
  local tps=$(python3 -c "print(int($total/$elapsed))")
  printf "  %-38s %7d 次 / %6ss  =>  %6d TPS\n" "$name" "$total" "$elapsed" "$tps"
  rm -rf "$tmpdir"
}

echo "---- 场景 1：主键点查 ----"
bench_persistent "按 ID 查房源" "SELECT * FROM room WHERE id=500000"

echo ""
echo "---- 场景 2：房源列表页 ----"
bench_persistent "房源列表(租户+状态)" "SELECT * FROM room WHERE tenant_id=1 AND status='vacant' LIMIT 20"

echo ""
echo "---- 场景 3：逾期扫描（Index Only Scan）----"
bench_persistent "逾期账单统计" "SELECT count(*) FROM bill WHERE tenant_id=1 AND status IN ('pending','partial','overdue')"

echo ""
echo "---- 场景 4：聚合统计（看板）----"
bench_persistent "出租率统计" "SELECT status,count(*) FROM room WHERE tenant_id=1 GROUP BY status"

echo ""
echo "=============================================="
echo " 环境说明"
echo "=============================================="
echo "  本机 PG 17.11 默认配置（shared_buffers=128MB）"
echo "  并发连接数: $CONCURRENCY（max_connections=100）"
echo "  阿里云生产（8C32G，shared_buffers=8GB）预计 2-5 倍"