topfans/backend/scripts/loadgen/scripts/prod_loadtest.sh
2026-06-16 23:00:07 +08:00

347 lines
11 KiB
Bash
Executable File

#!/bin/bash
# ===================================================================
# prod_loadtest.sh — 本地一键压测生产环境
# ===================================================================
# 设计目的: 不用 SSH 上生产机,本地直接调 seed/loadgen/cleanup
# 打生产 DB(走 SSH 端口转发)和生产网关(走公网)
#
# 架构:
# 1. SSH 端口转发: 本地 25432 → 生产机 docker 5432
# (绕开 host pg_dump 17.4 vs 容器 PG 18.3 版本不匹配问题)
# 2. seed/loadgen/cleanup: 直接在本地跑,通过环境变量/flag 传生产凭据
#
# 用法:
# ./prod_loadtest.sh <command> [args...]
#
# Commands:
# up 建立 SSH 端口转发 (本地 25432 → 生产 5432)
# down 关闭 SSH 端口转发
# status 查看隧道/DB/容器状态
# backup 备份生产 DB 到 /opt/topfans/backups/
# preflight 跑 7 项 preflight 检查
# seed 灌 1000 测试用户 + 23k 行数据
# loadgen [args...] 跑压测,所有参数透传给 ./bin/loadgen
# 例: ./prod_loadtest.sh loadgen --scenarios=S1,S2,S4 \
# --stage=step --step-schedule='5,10,20' \
# --duration=60s
# report 把 reports/ 下的 S*.json 合成 final-report.md
# cleanup 清理压测数据(保留 users + stars,适合多轮压测)
# cleanup-all 全清(含 users + stars,适合彻底收尾)
# pipeline [args...] 一键: up → backup → seed → preflight → loadgen
# → report → cleanup-all → down(loadgen args 透传)
# help 显示本帮助
#
# 环境变量(可覆盖):
# PROD_HOST 生产 SSH 地址 (默认 root@101.132.250.62)
# TUNNEL_PORT 本地转发端口 (默认 25432)
# BACKEND_DIR backend 绝对路径 (默认脚本所在目录的 ../../.. 即 backend/)
#
# 典型工作流(交互):
# ./prod_loadtest.sh up
# ./prod_loadtest.sh backup
# ./prod_loadtest.sh seed
# ./prod_loadtest.sh preflight
# ./prod_loadtest.sh loadgen --scenarios=S1,S2,S4 --stage=step \
# --step-schedule='5,10,20' --duration=60s
# ./prod_loadtest.sh report
# ./prod_loadtest.sh cleanup-all
# ./prod_loadtest.sh down
#
# 一键完成:
# ./prod_loadtest.sh pipeline --scenarios=S1,S2,S4 --stage=step \
# --step-schedule='5,10,20' --duration=60s
# ===================================================================
set -euo pipefail
# ===== 路径与默认值 =====
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
BACKEND_DIR="${BACKEND_DIR:-$(cd "$SCRIPT_DIR/../../.." && pwd)}"
REPO_ROOT="$(cd "$BACKEND_DIR/.." && pwd)"
PROD_HOST="${PROD_HOST:-root@101.132.250.62}"
TUNNEL_PORT="${TUNNEL_PORT:-25432}"
PROD_DB_NAME="topfans"
PROD_DB_USER="postgres"
ENV_PROD="$REPO_ROOT/docker/.env.prod"
# ===== 颜色输出 =====
if [[ -t 1 ]]; then
RED=$'\033[0;31m'; GREEN=$'\033[0;32m'; YELLOW=$'\033[0;33m'
BLUE=$'\033[0;34m'; BOLD=$'\033[1m'; RESET=$'\033[0m'
else
RED=""; GREEN=""; YELLOW=""; BLUE=""; BOLD=""; RESET=""
fi
info() { echo "${BLUE}>>>${RESET} $*"; }
ok() { echo "${GREEN}${RESET} $*"; }
warn() { echo "${YELLOW}⚠️${RESET} $*"; }
err() { echo "${RED}${RESET} $*" >&2; }
fatal() { err "$*"; exit 1; }
# ===== 通用工具 =====
read_env() {
# 读 docker/.env.prod 里的 key,值通过 stdout 返回
local key="$1"
if [[ ! -f "$ENV_PROD" ]]; then
fatal "$ENV_PROD 不存在"
fi
grep "^${key}=" "$ENV_PROD" | cut -d= -f2
}
ensure_built() {
# 跑命令前确保 bin/seed bin/loadgen 已编译
if [[ ! -x "$BACKEND_DIR/bin/seed" || ! -x "$BACKEND_DIR/bin/loadgen" ]]; then
info "编译 seed + loadgen → $BACKEND_DIR/bin/"
(cd "$BACKEND_DIR" && make loadgen-build >/dev/null)
fi
}
require_tunnel() {
if ! lsof -iTCP:"$TUNNEL_PORT" -sTCP:LISTEN >/dev/null 2>&1; then
fatal "本地 $TUNNEL_PORT 端口未监听,先跑: $0 up"
fi
}
# ===== 子命令:up =====
cmd_up() {
if lsof -iTCP:"$TUNNEL_PORT" -sTCP:LISTEN >/dev/null 2>&1; then
ok "隧道已存在 (PID: $(lsof -t -iTCP:$TUNNEL_PORT -sTCP:LISTEN), port $TUNNEL_PORT)"
return 0
fi
info "建立 SSH 隧道: 本地 $TUNNEL_PORT$PROD_HOST:5432"
ssh -o StrictHostKeyChecking=no \
-o ServerAliveInterval=30 \
-L "$TUNNEL_PORT":127.0.0.1:5432 \
-N -f "$PROD_HOST"
sleep 1
if lsof -iTCP:"$TUNNEL_PORT" -sTCP:LISTEN >/dev/null 2>&1; then
ok "隧道已建立 (PID: $(lsof -t -iTCP:$TUNNEL_PORT -sTCP:LISTEN))"
else
fatal "隧道建立失败"
fi
}
# ===== 子命令:down =====
cmd_down() {
if ! lsof -iTCP:"$TUNNEL_PORT" -sTCP:LISTEN >/dev/null 2>&1; then
ok "隧道已关闭 (或没建过)"
return 0
fi
local pid
pid=$(lsof -t -iTCP:"$TUNNEL_PORT" -sTCP:LISTEN)
kill "$pid" 2>/dev/null || true
sleep 1
if lsof -iTCP:"$TUNNEL_PORT" -sTCP:LISTEN >/dev/null 2>&1; then
warn "普通 kill 没关掉,试 kill -9"
kill -9 "$pid" 2>/dev/null || true
fi
ok "隧道已关闭 (PID: $pid)"
}
# ===== 子命令:status =====
cmd_status() {
echo "${BOLD}== SSH 隧道 ==${RESET}"
if lsof -iTCP:"$TUNNEL_PORT" -sTCP:LISTEN >/dev/null 2>&1; then
ok "本地 $TUNNEL_PORT 端口转发在跑 (PID: $(lsof -t -iTCP:$TUNNEL_PORT -sTCP:LISTEN))"
else
warn "本地 $TUNNEL_PORT 端口转发没建"
fi
echo ""
echo "${BOLD}== 生产 DB 容器 ==${RESET}"
ssh "$PROD_HOST" "docker ps --filter name=topfans-postgres --format 'table {{.Names}}\t{{.Status}}\t{{.Ports}}'"
echo ""
echo "${BOLD}== 最近 3 次备份 ==${RESET}"
ssh "$PROD_HOST" "ls -lh /opt/topfans/backups/ 2>/dev/null | tail -3" || warn "无备份或 backups 目录不存在"
echo ""
echo "${BOLD}== 本地 reports/ ==${RESET}"
ls -lh "$BACKEND_DIR/reports/"*.md "$BACKEND_DIR/reports/"*.json 2>/dev/null | tail -5 || warn "无 reports"
}
# ===== 子命令:backup =====
cmd_backup() {
local ts
ts=$(date +%Y%m%d-%H%M)
local backup_path="/opt/topfans/backups/pre-loadtest-${ts}.sql"
info "备份生产 DB → $backup_path"
local db_pass
db_pass=$(read_env DB_PASSWORD)
ssh "$PROD_HOST" "DB_PASSWORD=\$(grep '^DB_PASSWORD=' /opt/topfans/docker/.env.prod | cut -d= -f2)
mkdir -p /opt/topfans/backups
docker exec -e PGPASSWORD=\"\$DB_PASSWORD\" topfans-postgres \
pg_dump -U postgres -d topfans -Fc > $backup_path
ls -lh $backup_path"
ok "备份完成"
}
# ===== 子命令:seed =====
cmd_seed() {
require_tunnel
ensure_built
# bcrypt 文件(seed 用)
local bcrypt_file="$BACKEND_DIR/loadtest_bcrypt.txt"
if [[ ! -f "$bcrypt_file" ]]; then
info "生成 $bcrypt_file (匹配 tokens.go 硬编码密码 'Test@123')"
python3 -c "import bcrypt; print(bcrypt.hashpw(b'Test@123', bcrypt.gensalt(rounds=10)).decode())" > "$bcrypt_file"
fi
info "seed: 1000 users + 23k 行 (写生产 DB 走 SSH 隧道)"
local db_pass
db_pass=$(read_env DB_PASSWORD)
(
cd "$BACKEND_DIR"
DB_PASSWORD="$db_pass" \
JWT_SECRET=$(read_env JWT_SECRET) \
./bin/seed \
--db-host=127.0.0.1 --db-port="$TUNNEL_PORT" \
--db-name="$PROD_DB_NAME" --db-user="$PROD_DB_USER" \
--db-password="$db_pass"
)
ok "seed 完成"
}
# ===== 子命令:preflight =====
cmd_preflight() {
ensure_built
local db_pass
db_pass=$(read_env DB_PASSWORD)
info "跑 preflight (7 项检查)"
(
cd "$BACKEND_DIR"
DB_PASSWORD="$db_pass" \
JWT_SECRET=$(read_env JWT_SECRET) \
./bin/loadgen --cmd=preflight \
--target=http://101.132.250.62:8080 \
--prod-ssh="$PROD_HOST"
)
}
# ===== 子命令:loadgen =====
cmd_loadgen() {
require_tunnel
ensure_built
if [[ $# -eq 0 ]]; then
fatal "loadgen 需要参数,例: $0 loadgen --scenarios=S1,S2,S4 --stage=step"
fi
local db_pass
db_pass=$(read_env DB_PASSWORD)
info "loadgen ${*}"
(
cd "$BACKEND_DIR"
DB_PASSWORD="$db_pass" \
JWT_SECRET=$(read_env JWT_SECRET) \
./bin/loadgen --cmd=run \
--target=http://101.132.250.62:8080 \
--prod-ssh="$PROD_HOST" \
"$@"
)
}
# ===== 子命令:report =====
cmd_report() {
ensure_built
info "生成 final-report.md"
(
cd "$BACKEND_DIR"
./bin/loadgen --cmd=report \
--input=./reports \
--output=./reports/final-report.md
)
ok "final-report.md 已更新: $BACKEND_DIR/reports/final-report.md"
}
# ===== 子命令:cleanup =====
cmd_cleanup() {
require_tunnel
ensure_built
local db_pass
db_pass=$(read_env DB_PASSWORD)
warn "清理压测数据(保留 users + stars,适合多轮压测)"
(
cd "$BACKEND_DIR"
DB_PASSWORD="$db_pass" \
./bin/seed --cleanup \
--db-host=127.0.0.1 --db-port="$TUNNEL_PORT" \
--db-name="$PROD_DB_NAME" --db-user="$PROD_DB_USER" \
--db-password="$db_pass"
)
ok "cleanup 完成"
}
# ===== 子命令:cleanup-all =====
cmd_cleanup_all() {
require_tunnel
ensure_built
local db_pass
db_pass=$(read_env DB_PASSWORD)
warn "⚠️ 全清压测数据(含 users + stars,序列同步重置)"
(
cd "$BACKEND_DIR"
DB_PASSWORD="$db_pass" \
./bin/seed --cleanup --full \
--db-host=127.0.0.1 --db-port="$TUNNEL_PORT" \
--db-name="$PROD_DB_NAME" --db-user="$PROD_DB_USER" \
--db-password="$db_pass"
)
ok "cleanup-all 完成"
}
# ===== 子命令:pipeline =====
cmd_pipeline() {
info "=== pipeline: 1/8 up ==="
cmd_up
info "=== pipeline: 2/8 backup ==="
cmd_backup
info "=== pipeline: 3/8 seed ==="
cmd_seed
info "=== pipeline: 4/8 preflight ==="
cmd_preflight || warn "preflight 有 FAIL,继续 (③ backup <50MB 预期会 FAIL)"
info "=== pipeline: 5/8 loadgen ==="
cmd_loadgen "$@" || warn "loadgen 出错,继续到 cleanup"
info "=== pipeline: 6/8 report ==="
cmd_report || warn "report 生成失败,继续"
info "=== pipeline: 7/8 cleanup-all ==="
cmd_cleanup_all || warn "cleanup 失败"
info "=== pipeline: 8/8 down ==="
cmd_down
ok "pipeline 完成"
}
# ===== 帮助 =====
cmd_help() {
sed -n '3,50p' "$0"
}
# ===== 入口 =====
if [[ $# -eq 0 ]]; then
cmd_help
exit 1
fi
cmd="${1:-help}"
shift || true
case "$cmd" in
up) cmd_up ;;
down) cmd_down ;;
status) cmd_status ;;
backup) cmd_backup ;;
seed) cmd_seed "$@" ;;
preflight) cmd_preflight ;;
loadgen) cmd_loadgen "$@" ;;
report) cmd_report ;;
cleanup) cmd_cleanup ;;
cleanup-all) cmd_cleanup_all ;;
pipeline) cmd_pipeline "$@" ;;
help|-h|--help) cmd_help ;;
*) err "未知子命令: $cmd"; cmd_help; exit 2 ;;
esac