#!/usr/bin/env bash
# xmhls-burnin 1.0.0 — 超频服务器分阶段烤机编排（会施加满负载）
#
# 用法: sudo bash xmhls-burnin-1.0.0.sh [选项]
#   --phase P        cpu | mem | latency | all（默认 all，按 cpu → mem → latency 依次执行）
#   --duration D     每阶段时长，如 90s / 30m / 2h（默认 30m）
#   --max-temp C     温度上限 °C，超过立即终止并判定失败（默认 90）
#   --cpus LIST      latency 阶段 cyclictest 使用的核（默认除 CPU0 外全部；建议填隔离核，如 2-15）
#   --lat-max US     cyclictest 最大延迟判定上限 µs（默认 100；调优后的隔离核通常应远低于此值）
#   --interval S     监控采样间隔秒数（默认 10）
#   --outdir DIR     结果目录（默认 ./xmhls-burnin-<主机名>-<UTC 时间>）
#   --allow-no-temp  读不到温度传感器时仍继续（默认拒绝：超温保护不可用）
#   --dry-run        只打印计划与命令，不执行
#   --yes            跳过开始前的确认
# 工具: cpu 需 stress-ng；mem 需 stressapptest 或 stress-ng；latency 需 cyclictest（可选 hwlatdetect），且需 root。
# 安全: 超频硬件满负载存在损坏风险，请值守并同时观察 BMC / 水冷状态；不要在生产交易时段运行。
#
# MIT License · Copyright (c) 2026 西门华莱硕投资咨询有限公司 · https://xmhls.com/guide/#script-burnin
set -uo pipefail
VERSION=1.0.0
PHASE=all; DURATION=30m; MAX_TEMP=90; CPUS=""; LAT_MAX=100; INTERVAL=10; OUTDIR=""
ALLOW_NO_TEMP=0; DRY=0; YES=0
HWMON_DIR="${XMHLS_HWMON_DIR:-/sys/class/hwmon}"   # 测试钩子：可指向模拟的 hwmon 目录

usage() { sed -n '2,21p' "$0" | sed 's/^# \{0,1\}//'; }
while [ $# -gt 0 ]; do
    case "$1" in
        --phase) PHASE="${2:?}"; shift 2 ;;
        --duration) DURATION="${2:?}"; shift 2 ;;
        --max-temp) MAX_TEMP="${2:?}"; shift 2 ;;
        --cpus) CPUS="${2:?}"; shift 2 ;;
        --lat-max) LAT_MAX="${2:?}"; shift 2 ;;
        --interval) INTERVAL="${2:?}"; shift 2 ;;
        --outdir) OUTDIR="${2:?}"; shift 2 ;;
        --allow-no-temp) ALLOW_NO_TEMP=1; shift ;;
        --dry-run) DRY=1; shift ;;
        --yes) YES=1; shift ;;
        -h|--help) usage; exit 0 ;;
        *) echo "未知参数: $1" >&2; usage >&2; exit 2 ;;
    esac
done

die() { echo "错误: $*" >&2; exit 2; }
have() { command -v "$1" >/dev/null 2>&1; }
to_seconds() {
    [[ "$1" =~ ^([0-9]+)([smhd]?)$ ]] || return 1
    local n="${BASH_REMATCH[1]}"
    case "${BASH_REMATCH[2]}" in m) n=$((n * 60)) ;; h) n=$((n * 3600)) ;; d) n=$((n * 86400)) ;; esac
    echo "$n"
}
SECS="$(to_seconds "$DURATION")" || die "--duration 格式应为 数字[s|m|h|d]，如 30m"
[ "$SECS" -ge 30 ] || die "--duration 至少 30s"
[[ "$MAX_TEMP" =~ ^[0-9]+$ && "$LAT_MAX" =~ ^[0-9]+$ && "$INTERVAL" =~ ^[0-9]+$ ]] || die "--max-temp / --lat-max / --interval 须为整数"
case "$PHASE" in cpu|mem|latency) PHASES=("$PHASE") ;; all) PHASES=(cpu mem latency) ;; *) die "--phase 只能是 cpu|mem|latency|all" ;; esac
NCPU="$(nproc)"
[ -n "$CPUS" ] || CPUS="$([ "$NCPU" -gt 1 ] && echo "1-$((NCPU - 1))" || echo 0)"
[[ "$CPUS" =~ ^[0-9,-]+$ ]] || die "--cpus 格式应为 CPU 列表，如 2-15 或 2,4,6"
count_cpus() { local n=0 p; IFS=',' read -ra ps <<<"$1"; for p in "${ps[@]}"; do if [[ "$p" == *-* ]]; then n=$((n + ${p#*-} - ${p%-*} + 1)); else n=$((n + 1)); fi; done; echo "$n"; }
NTHREADS="$(count_cpus "$CPUS")"
HOST="$(hostname)"; TS="$(date -u +%Y%m%dT%H%M%SZ)"
[ -n "$OUTDIR" ] || OUTDIR="./xmhls-burnin-$HOST-$TS"

# ---------- 传感器 ----------
max_temp() {   # 输出整数 °C；无传感器输出 NA。只取 CPU 相关芯片（coretemp / k10temp / zenpower / cpu_thermal）
    local best="" d name f v
    for d in "$HWMON_DIR"/hwmon*; do
        name="$(cat "$d/name" 2>/dev/null)"
        case "$name" in coretemp|k10temp|zenpower|cpu_thermal|x86_pkg_temp) ;; *) continue ;; esac
        for f in "$d"/temp*_input; do
            [ -r "$f" ] || continue
            v=$(( $(cat "$f") / 1000 ))
            [ -z "$best" ] || [ "$v" -gt "$best" ] && best="$v"
        done
    done
    echo "${best:-NA}"
}
avg_mhz() { awk -F: '/^cpu MHz/{s+=$2; n++} END{if (n) printf "%.0f", s/n; else print "NA"}' /proc/cpuinfo; }
edac_sum() {   # 输出 "ce ue"
    local ce=0 ue=0 mc
    for mc in /sys/devices/system/edac/mc/mc*; do
        [ -d "$mc" ] || continue
        ce=$((ce + $(cat "$mc/ce_count"))); ue=$((ue + $(cat "$mc/ue_count")))
    done
    echo "$ce $ue"
}
mce_lines() { journalctl -k -b --no-pager 2>/dev/null | grep -ciE 'mce:|machine check|hardware error' || true; }

# ---------- 计划 ----------
plan_cmd() {   # plan_cmd <阶段> <秒数>：打印该阶段将执行的命令（每行一条）
    local s="$2"
    case "$1" in
        cpu)
            echo "stress-ng --cpu 0 --cpu-method all --verify --tz --metrics-brief -t $((s * 2 / 3))s"
            echo "stress-ng --matrix 0 --tz --metrics-brief -t $((s - s * 2 / 3))s" ;;
        mem)
            local avail_mb=$(( $(awk '/^MemAvailable:/{print $2}' /proc/meminfo) * 3 / 4 / 1024 ))
            if have stressapptest; then echo "stressapptest -s $s -M $avail_mb -W"
            else echo "stress-ng --vm $NCPU --vm-bytes 75% --verify --metrics-brief -t ${s}s"; fi ;;
        latency)
            local cy="$s"
            if have hwlatdetect; then   # 先在空闲状态测硬件 / 固件停顿（占 1/3 时长），余下时间给 cyclictest
                echo "hwlatdetect --duration=$((s / 3))s --threshold=10 --report=$OUTDIR/hwlatdetect-report.txt"
                cy=$((s - s / 3))
            fi
            echo "cyclictest -m -p95 -t $NTHREADS -a $CPUS -i 200 -h 400 -q -D $cy" ;;
    esac
}

echo "xmhls-burnin $VERSION · host=$HOST · CPU×$NCPU · 每阶段 $DURATION · 温度上限 ${MAX_TEMP}°C · 结果目录 $OUTDIR"
T0="$(max_temp)"
echo "当前最高 CPU 温度: ${T0}°C · 平均频率: $(avg_mhz) MHz · EDAC(ce ue): $(edac_sum)"
missing=0
for ph in "${PHASES[@]}"; do
    echo; echo "【阶段 $ph】"
    case "$ph" in
        cpu) have stress-ng || { echo "  缺少 stress-ng（下载：https://xmhls.com/downloads/#stability）"; missing=1; } ;;
        mem) have stressapptest || have stress-ng || { echo "  缺少 stressapptest / stress-ng"; missing=1; } ;;
        latency)
            have cyclictest || { echo "  缺少 cyclictest（rt-tests，下载：https://xmhls.com/downloads/#latency）"; missing=1; }
            [ "$(id -u)" -eq 0 ] || { echo "  latency 阶段需要 root（实时优先级与 hwlatdetect）"; missing=1; } ;;
    esac
    plan_cmd "$ph" "$SECS" | sed 's/^/  $ /'
done
[ "$DRY" = 1 ] && { echo; echo "[DRY_RUN] 仅显示计划，未执行任何负载。"; exit 0; }
[ "$missing" = 0 ] || die "缺少依赖，未开始"
if [ "$T0" = NA ] && [ "$ALLOW_NO_TEMP" = 0 ]; then
    die "读不到 CPU 温度传感器（coretemp / k10temp），超温保护不可用。确认已加载传感器驱动，或加 --allow-no-temp 并人工全程监控"
fi
if [ "$YES" = 0 ]; then
    [ -t 0 ] || die "非交互环境请加 --yes 确认"
    read -r -p $'\n确认开始？此操作将使 CPU / 内存满负载运行。输入 yes 继续: ' ans
    [ "$ans" = yes ] || { echo "已取消"; exit 1; }
fi

# ---------- 执行 ----------
mkdir -p "$OUTDIR" || die "无法创建 $OUTDIR"
ABORT="$OUTDIR/.abort"; CUR="$OUTDIR/.current_pid"; rm -f "$ABORT" "$CUR"
{
    echo "xmhls-burnin $VERSION  start=$(date -u '+%F %T UTC')  host=$HOST  args: phase=$PHASE duration=$DURATION max-temp=$MAX_TEMP cpus=$CPUS lat-max=$LAT_MAX"
    uname -a; grep -m1 'model name' /proc/cpuinfo; echo "microcode: $(grep -m1 '^microcode' /proc/cpuinfo | awk '{print $3}')"
} > "$OUTDIR/env.txt"
echo "utc,phase,max_temp_c,avg_mhz,edac_ce,edac_ue" > "$OUTDIR/monitor.csv"
echo idle > "$OUTDIR/.phase"

stop_current() {   # 终止当前负载进程及其子进程（stress-ng 收到 TERM 也会自行回收 worker）
    local pid; pid="$(cat "$CUR" 2>/dev/null)"
    [ -n "$pid" ] || return 0
    pkill -TERM -P "$pid" 2>/dev/null; kill -TERM "$pid" 2>/dev/null
}
monitor() {
    while :; do
        local t m e ph
        t="$(max_temp)"; m="$(avg_mhz)"; e="$(edac_sum)"; ph="$(cat "$OUTDIR/.phase" 2>/dev/null)"
        echo "$(date -u +%FT%TZ),$ph,$t,$m,${e% *},${e#* }" >> "$OUTDIR/monitor.csv"
        if [ "$t" != NA ] && [ "$t" -ge "$MAX_TEMP" ]; then
            echo "$(date -u +%FT%TZ) 温度 ${t}°C ≥ 上限 ${MAX_TEMP}°C（阶段 $ph）" > "$ABORT"
            stop_current
        fi
        sleep "$INTERVAL"
    done
}
monitor & MON=$!
cleanup() {
    stop_current
    kill "$MON" 2>/dev/null; wait "$MON" 2>/dev/null
}
trap 'echo; echo "收到中断，正在停止负载…"; cleanup; exit 130' INT TERM

run_step() {   # run_step <日志文件> <命令...>：后台运行并记录 PID，供超温 / 中断时终止
    local log="$1"; shift
    "$@" > "$log" 2>&1 &
    local pid=$!
    echo "$pid" > "$CUR"
    wait "$pid"; local rc=$?
    : > "$CUR"
    return "$rc"
}

declare -A RESULT NOTE
CE0_ALL="$(edac_sum)"; MCE0_ALL="$(mce_lines)"
for ph in "${PHASES[@]}"; do
    echo "$ph" > "$OUTDIR/.phase"
    e0="$(edac_sum)"; m0="$(mce_lines)"; ok=1; note=""
    echo; echo "▶ 阶段 $ph 开始 $(date -u '+%T UTC')"
    i=0
    while IFS= read -r cmd; do
        i=$((i + 1)); log="$OUTDIR/$ph-$i.log"
        echo "  \$ $cmd"
        # shellcheck disable=SC2086
        run_step "$log" $cmd; rc=$?
        [ -f "$ABORT" ] && { ok=0; note="超温终止：$(cat "$ABORT")"; break; }
        case "$cmd" in
            stressapptest*) grep -q 'Status: PASS' "$log" || { ok=0; note+="stressapptest 未 PASS；"; } ;;
            hwlatdetect*)
                n="$(grep -oE 'Samples exceeding threshold: *[0-9]+' "$log" | grep -oE '[0-9]+$')"
                [ "${n:-0}" = 0 ] || { ok=0; note+="hwlatdetect 超阈值样本 $n 个；"; } ;;
            cyclictest*)
                mx="$(grep -E '^# Max Latencies:' "$log" | tr -s ' ' '\n' | grep -E '^[0-9]+$' | sort -n | tail -n 1)"
                if [ -n "$mx" ]; then note+="cyclictest 最大 $((10#$mx)) µs；"; [ "$((10#$mx))" -le "$LAT_MAX" ] || ok=0
                else ok=0; note+="cyclictest 无结果；"; fi ;;
            *) [ "$rc" = 0 ] || { ok=0; note+="$(cut -d' ' -f1 <<<"$cmd") 退出码 $rc；"; } ;;
        esac
    done < <(plan_cmd "$ph" "$SECS")
    e1="$(edac_sum)"; m1="$(mce_lines)"
    dce=$(( ${e1% *} - ${e0% *} )); due=$(( ${e1#* } - ${e0#* } )); dm=$(( m1 - m0 ))
    [ "$dce" -eq 0 ] && [ "$due" -eq 0 ] || { ok=0; note+="EDAC 新增 可纠正 $dce / 不可纠正 $due；"; }
    [ "$dm" -eq 0 ] || { ok=0; note+="内核日志新增 MCE / 硬件错误 $dm 行；"; }
    RESULT[$ph]=$([ "$ok" = 1 ] && echo PASS || echo FAIL); NOTE[$ph]="${note:-无异常}"
    echo "◀ 阶段 $ph：${RESULT[$ph]}（${NOTE[$ph]}）"
    [ -f "$ABORT" ] && break
done
cleanup; trap - INT TERM

{
    echo; echo "========== 烤机结果 =========="
    for ph in "${PHASES[@]}"; do printf '  %-8s %-5s %s\n' "$ph" "${RESULT[$ph]:-未执行}" "${NOTE[$ph]:-}"; done
    echo "  峰值温度: $(awk -F, 'NR>1 && $3!="NA"{if ($3>m) m=$3} END{print (m==""?"NA":m"°C")}' "$OUTDIR/monitor.csv")"
    echo "  EDAC 全程: 开始 $CE0_ALL → 结束 $(edac_sum)（ce ue）；内核 MCE 行: $MCE0_ALL → $(mce_lines)"
    echo "  结果目录: $OUTDIR（monitor.csv 为温度 / 频率 / EDAC 采样，*.log 为各工具原始输出）"
} | tee "$OUTDIR/summary.txt"

for ph in "${PHASES[@]}"; do [ "${RESULT[$ph]:-}" = PASS ] || exit 1; done
exit 0