#!/usr/bin/env bash
# PostgreSQL 直连 vs 每次新建连接 vs PgBouncer 事务池 演示脚本 v1.0（surinch.com）
# 仅限测试实例！会对目标库做 pgbench 只读压测（-S），会短时间打满连接数。不要对生产库运行。
# 前置：pgbench -i -s 10 已初始化；PgBouncer 已在 BOUNCER_PORT 监听同一个库（pool_mode=transaction）。
# 用法：PGHOST=127.0.0.1 PGUSER=lab PGPASSWORD=... DB=pgconn_lab bash pg-connection-pool-pgbouncer-demo-v1.0.sh
set -u
PGBENCH=${PGBENCH:-pgbench}
DB=${DB:-pgconn_lab}
PGPORT=${PGPORT:-5432}
BOUNCER_PORT=${BOUNCER_PORT:-6432}
DUR=${DUR:-15}
C=${C:-20}
BIG_C=${BIG_C:-120}

psql -p "$PGPORT" -d "$DB" -Atc "SELECT version()"
psql -p "$PGPORT" -d "$DB" -Atc "SELECT 'max_connections='||current_setting('max_connections')||' superuser_reserved='||current_setting('superuser_reserved_connections')"

run() { # $1=标签 $2=端口 其余=pgbench 参数
  local label=$1 port=$2; shift 2
  echo; echo "== $label"
  $PGBENCH -p "$port" -n -S -j 4 -T "$DUR" "$@" "$DB" 2>&1 | grep -E "number of clients|number of transactions actually|failed|latency average|average connection time|tps =|FATAL|error" 
}

run "A 直连长连接（-c $C）" "$PGPORT" -c "$C"
run "B 直连，每个事务新建连接（-c $C -C）" "$PGPORT" -c "$C" -C
run "C 经 PgBouncer，每个事务新建连接（-c $C -C）" "$BOUNCER_PORT" -c "$C" -C
run "D 经 PgBouncer 长连接（-c $C）" "$BOUNCER_PORT" -c "$C"

echo; echo "== E 直连 $BIG_C 个客户端（超过 max_connections）"
# -j 1：单线程建连，避免多线程并发 SCRAM 握手时 libpq/OpenSSL 偶发的客户端报错干扰结果
$PGBENCH -p "$PGPORT" -n -S -j 1 -T 5 -c "$BIG_C" "$DB" 2>&1 | grep -E "FATAL|could not create|tps =" | head -5

echo; echo "== F 经 PgBouncer $BIG_C 个客户端（default_pool_size 决定真正连到 PG 的数量）"
( DUR=10 run "F" "$BOUNCER_PORT" -c "$BIG_C" ) &
sleep 5
echo "-- 压测进行中，PG 侧按 application_name / state 统计后端连接："
psql -p "$PGPORT" -d "$DB" -c "SELECT coalesce(nullif(application_name,''),'(空)') AS app, coalesce(state,'-') AS state, count(*) FROM pg_stat_activity WHERE datname=current_database() AND pid<>pg_backend_pid() GROUP BY 1,2 ORDER BY 1,2"
wait

echo; echo "== G 事务池被 idle in transaction 占满：20 个客户端经 PgBouncer 各开事务后空闲 12 秒，第 21 个客户端执行 select 1"
for i in $(seq 1 20); do ( (echo "BEGIN; SELECT abalance FROM pgbench_accounts WHERE aid=$i;"; sleep 12; echo "ROLLBACK;") | psql -p "$BOUNCER_PORT" -d "$DB" -q >/dev/null 2>&1 ) & done
sleep 3
t0=$(date +%s.%N); psql -p "$BOUNCER_PORT" -d "$DB" -Atc "select 1" >/dev/null 2>&1; t1=$(date +%s.%N)
echo "第 21 个客户端耗时 $(awk -v a=$t0 -v b=$t1 'BEGIN{printf "%.3f", b-a}') 秒（default_pool_size=20 时应接近 9 秒）"
wait
# 可选对照（仅测试环境）：另起一个 auth_type=trust 的 PgBouncer，再跑一次 -C，可看出每次新建连接的耗时主要花在 SCRAM 认证上。