515 lines
14 KiB
Go
515 lines
14 KiB
Go
package stress
|
||
|
||
import (
|
||
"fmt"
|
||
"net"
|
||
"os"
|
||
"sort"
|
||
"strconv"
|
||
"strings"
|
||
"time"
|
||
|
||
"auto-check/pkg/config"
|
||
"auto-check/pkg/discovery"
|
||
"auto-check/pkg/sshclient"
|
||
|
||
"golang.org/x/crypto/ssh"
|
||
)
|
||
|
||
// Results 测试结果表（IP → 报告），包级公开
|
||
var Results = make(map[string]*Report)
|
||
|
||
// ============================
|
||
// 执行器
|
||
// ============================
|
||
|
||
// Runner 压力测试执行器
|
||
type Runner struct {
|
||
cfg Config
|
||
types []TestType
|
||
client *ssh.Client
|
||
tools ToolSet
|
||
metrics []Metric
|
||
installFailed []string // 自动安装仍失败的工具
|
||
}
|
||
|
||
// NewRunner 创建执行器（检查/在线安装 → 构建指标 → 生成脚本 → 执行）
|
||
func NewRunner(cfg Config, types []TestType, client *ssh.Client, autoInstall bool, sudoPwd string) *Runner {
|
||
// 1. 检测工具，缺失时按发行版在线安装（可配置国内镜像源加速）
|
||
tools, failed := EnsureTools(client, requiredTools, autoInstall, sudoPwd, cfg.UseMirror)
|
||
metrics := BuildMetrics(types, tools)
|
||
return &Runner{cfg: cfg, types: types, client: client, tools: tools, metrics: metrics, installFailed: failed}
|
||
}
|
||
|
||
// Run 执行压力测试（生成脚本 → 远程执行 → 解析结果）
|
||
func (r *Runner) Run(ip string) *Report {
|
||
report := &Report{IP: ip, StartTime: time.Now()}
|
||
fmt.Printf("\n════════ [%s] 压力测试开始 ════════\n", ip)
|
||
|
||
// 显示可用工具
|
||
fmt.Printf(" 工具: ")
|
||
for name, info := range r.tools.Tools {
|
||
fmt.Printf("%s(%s) ", name, info.Version)
|
||
}
|
||
fmt.Println()
|
||
if len(r.installFailed) > 0 {
|
||
fmt.Printf(" [安装] 以下工具自动安装失败，相关测试将跳过: %v\n", r.installFailed)
|
||
}
|
||
|
||
// 显示将执行的指标及被跳过的原因
|
||
fmt.Printf(" 指标: ")
|
||
for _, m := range r.metrics {
|
||
if m.Enabled {
|
||
fmt.Printf("%s ", m.Name)
|
||
}
|
||
}
|
||
fmt.Println()
|
||
for _, m := range r.metrics {
|
||
if !m.Enabled && !m.IsMonitor {
|
||
fmt.Printf(" [跳过] %s: %s\n", m.Name, m.DisableReason)
|
||
}
|
||
}
|
||
|
||
// 按需读取各测试脚本，拼装成完整脚本
|
||
script := r.assembleScript()
|
||
fmt.Printf(" [脚本] 拼装完成，%d 字节\n", len(script))
|
||
|
||
// 通过 heredoc 传到远程执行
|
||
scriptCmd := fmt.Sprintf("cat <<'AUTOCHECKSCRIPT' > /tmp/auto-check.sh\n%s\nAUTOCHECKSCRIPT\nchmod +x /tmp/auto-check.sh && bash /tmp/auto-check.sh 2>&1", script)
|
||
output, err := sshclient.RunCommand(r.client, scriptCmd)
|
||
|
||
if err != nil {
|
||
fmt.Printf(" [执行] 脚本执行出错: %v\n", err)
|
||
}
|
||
|
||
// 解析结果
|
||
if output != "" {
|
||
r.parseResults(report, output)
|
||
}
|
||
|
||
report.EndTime = time.Now()
|
||
report.Duration = report.EndTime.Sub(report.StartTime)
|
||
|
||
// 生成综合 HTML 报告（跑分卡片 + 曲线图）
|
||
if len(report.Results) > 0 {
|
||
reportPath, err := WriteReportHTML(ip, report, "reports")
|
||
if err != nil {
|
||
fmt.Printf(" [报告] 生成失败: %v\n", err)
|
||
} else {
|
||
fmt.Printf(" [报告] %s\n", reportPath)
|
||
}
|
||
}
|
||
|
||
fmt.Printf("\n════════ [%s] 压力测试完成 ════════\n", ip)
|
||
fmt.Println(report.ToText())
|
||
return report
|
||
}
|
||
|
||
// parseResults 解析脚本的结构化输出
|
||
func (r *Runner) parseResults(report *Report, output string) {
|
||
lines := strings.Split(output, "\n")
|
||
var currentTest string
|
||
var currentStatus string
|
||
var currentDuration string
|
||
var currentOutput []string
|
||
var currentSamples []Sample
|
||
var currentScore float64
|
||
var currentMetrics map[string]float64
|
||
var sampleHeader []string
|
||
inSamples := false
|
||
inSysInfo := false
|
||
|
||
flush := func() {
|
||
if currentTest == "" {
|
||
return
|
||
}
|
||
status := StatusPass
|
||
switch currentStatus {
|
||
case "fail":
|
||
status = StatusFail
|
||
case "skip":
|
||
status = StatusSkip
|
||
case "error":
|
||
status = StatusError
|
||
}
|
||
// 温度超限判定（仅对 temp 监控）
|
||
if TestType(currentTest) == MonitorTemp && status == StatusPass && r.cfg.TempLimit > 0 {
|
||
if maxT, ok := currentMetrics["max_temp"]; ok && maxT > r.cfg.TempLimit {
|
||
status = StatusFail
|
||
currentOutput = append(currentOutput,
|
||
fmt.Sprintf("温度超限: 最高 %.1f°C > 阈值 %.1f°C", maxT, r.cfg.TempLimit))
|
||
}
|
||
}
|
||
duration, _ := time.ParseDuration(currentDuration)
|
||
report.AddResult(Result{
|
||
Type: TestType(currentTest),
|
||
Status: status,
|
||
Output: strings.Join(currentOutput, "\n"),
|
||
Duration: duration,
|
||
Samples: currentSamples,
|
||
Score: currentScore,
|
||
Metrics: currentMetrics,
|
||
})
|
||
currentTest = ""
|
||
currentStatus = ""
|
||
currentDuration = ""
|
||
currentOutput = nil
|
||
currentSamples = nil
|
||
currentScore = 0
|
||
currentMetrics = nil
|
||
sampleHeader = nil
|
||
}
|
||
|
||
for _, line := range lines {
|
||
line = strings.TrimSpace(line)
|
||
|
||
// 采样数据块
|
||
if line == "===SAMPLES===" {
|
||
inSamples = true
|
||
sampleHeader = nil
|
||
continue
|
||
}
|
||
if line == "===END_SAMPLES===" {
|
||
inSamples = false
|
||
continue
|
||
}
|
||
if inSamples {
|
||
parts := strings.Split(line, ",")
|
||
if sampleHeader == nil {
|
||
// 首行为表头：time,指标1,指标2,...
|
||
for _, p := range parts {
|
||
sampleHeader = append(sampleHeader, strings.TrimSpace(p))
|
||
}
|
||
continue
|
||
}
|
||
if len(parts) < 1 || strings.TrimSpace(parts[0]) == "" {
|
||
continue
|
||
}
|
||
s := Sample{Time: strings.TrimSpace(parts[0]), Values: map[string]float64{}}
|
||
for i := 1; i < len(parts) && i < len(sampleHeader); i++ {
|
||
name := sampleHeader[i]
|
||
if name == "" || name == "time" {
|
||
continue
|
||
}
|
||
if v, err := strconv.ParseFloat(strings.TrimSpace(parts[i]), 64); err == nil {
|
||
s.Values[name] = v
|
||
}
|
||
}
|
||
currentSamples = append(currentSamples, s)
|
||
continue
|
||
}
|
||
|
||
if line == "===SYSTEM_INFO===" {
|
||
inSysInfo = true
|
||
continue
|
||
}
|
||
if inSysInfo {
|
||
if line == "" {
|
||
inSysInfo = false
|
||
continue
|
||
}
|
||
if i := strings.Index(line, ":"); i > 0 {
|
||
k := strings.TrimSpace(line[:i])
|
||
v := strings.TrimSpace(line[i+1:])
|
||
if report.SysInfo == nil {
|
||
report.SysInfo = map[string]string{}
|
||
}
|
||
report.SysInfo[k] = v
|
||
}
|
||
continue
|
||
}
|
||
|
||
if strings.HasPrefix(line, "===TEST:") {
|
||
flush()
|
||
currentTest = strings.TrimSuffix(strings.TrimPrefix(line, "===TEST:"), "===")
|
||
continue
|
||
}
|
||
|
||
if strings.HasPrefix(line, "===MONITOR:") {
|
||
flush()
|
||
currentTest = strings.TrimSuffix(strings.TrimPrefix(line, "===MONITOR:"), "===")
|
||
continue
|
||
}
|
||
|
||
if line == "===END===" {
|
||
flush()
|
||
continue
|
||
}
|
||
|
||
// 解析 key:value
|
||
if strings.HasPrefix(line, "status:") {
|
||
currentStatus = strings.TrimPrefix(line, "status:")
|
||
continue
|
||
}
|
||
if strings.HasPrefix(line, "duration:") {
|
||
durStr := strings.TrimPrefix(line, "duration:")
|
||
ms, _ := strconv.ParseInt(strings.TrimSuffix(durStr, "ms"), 10, 64)
|
||
currentDuration = fmt.Sprintf("%dms", ms)
|
||
continue
|
||
}
|
||
if strings.HasPrefix(line, "score:") {
|
||
currentScore, _ = strconv.ParseFloat(strings.TrimPrefix(line, "score:"), 64)
|
||
continue
|
||
}
|
||
if strings.HasPrefix(line, "metric:") {
|
||
kv := strings.TrimPrefix(line, "metric:")
|
||
if i := strings.Index(kv, "="); i > 0 {
|
||
name := strings.TrimSpace(kv[:i])
|
||
val, _ := strconv.ParseFloat(strings.TrimSpace(kv[i+1:]), 64)
|
||
if currentMetrics == nil {
|
||
currentMetrics = map[string]float64{}
|
||
}
|
||
currentMetrics[name] = val
|
||
}
|
||
continue
|
||
}
|
||
if strings.HasPrefix(line, "output:") {
|
||
currentOutput = append(currentOutput, strings.TrimPrefix(line, "output:"))
|
||
continue
|
||
}
|
||
|
||
// 普通输出行
|
||
if currentTest != "" && line != "" {
|
||
currentOutput = append(currentOutput, line)
|
||
}
|
||
}
|
||
flush()
|
||
}
|
||
|
||
// assembleScript 根据 metrics 读取对应的独立脚本文件，拼装成完整脚本
|
||
func (r *Runner) assembleScript() string {
|
||
var sb strings.Builder
|
||
|
||
// shebang + 环境变量导出
|
||
sb.WriteString("#!/bin/bash\n")
|
||
sb.WriteString("set -e\n\n")
|
||
|
||
// 导出参数变量
|
||
fmt.Fprintf(&sb, "export AUTOCHECK_DURATION=%d\n", int(r.cfg.Duration.Seconds()))
|
||
fmt.Fprintf(&sb, "export AUTOCHECK_THREADS=%d\n", r.cfg.Threads)
|
||
fmt.Fprintf(&sb, "export AUTOCHECK_DISK_SIZE_MB=%d\n", r.cfg.DiskSizeMB)
|
||
fmt.Fprintf(&sb, "export AUTOCHECK_MEM_SIZE_MB=%d\n", r.cfg.MemSizeMB)
|
||
fmt.Fprintf(&sb, "export AUTOCHECK_TEMP_INTERVAL=%d\n", int(r.cfg.TempLogInt.Seconds()))
|
||
sb.WriteString("export AUTOCHECK_SAMPLE_INTERVAL=2\n")
|
||
if r.cfg.DiskDir != "" {
|
||
fmt.Fprintf(&sb, "export AUTOCHECK_DISK_DIR=%s\n", r.cfg.DiskDir)
|
||
}
|
||
sb.WriteString("\n")
|
||
|
||
// 系统信息
|
||
sb.Write(r.readScript("sysinfo"))
|
||
sb.WriteString("\n")
|
||
|
||
// 温度监控启动
|
||
hasTemp := false
|
||
for _, m := range r.metrics {
|
||
if m.Name == "temp" && m.Enabled {
|
||
hasTemp = true
|
||
break
|
||
}
|
||
}
|
||
if hasTemp {
|
||
sb.Write(r.readScript("temp_start"))
|
||
sb.WriteString("\n")
|
||
}
|
||
|
||
// 按"先单项、后综合"排序测试脚本
|
||
ordered := make([]Metric, len(r.metrics))
|
||
copy(ordered, r.metrics)
|
||
sort.SliceStable(ordered, func(i, j int) bool {
|
||
return testPriority(ordered[i]) < testPriority(ordered[j])
|
||
})
|
||
|
||
// 逐个测试脚本
|
||
for _, m := range ordered {
|
||
if m.IsMonitor {
|
||
continue
|
||
}
|
||
if !m.Enabled {
|
||
// 缺少依赖时仍输出一个 skip 结果，让报告完整展示原因
|
||
fmt.Fprintf(&sb, "echo '===TEST:%s==='\n", m.Name)
|
||
sb.WriteString("echo 'status:skip'\n")
|
||
fmt.Fprintf(&sb, "echo 'output:%s'\n", m.DisableReason)
|
||
sb.WriteString("echo ''\n")
|
||
continue
|
||
}
|
||
data := r.readScript(m.Name)
|
||
if data == nil {
|
||
fmt.Fprintf(&sb, "echo '===TEST:%s==='\n", m.Name)
|
||
sb.WriteString("echo 'status:skip'\n")
|
||
fmt.Fprintf(&sb, "echo 'output:脚本文件 scripts/%s.sh 不存在'\n", m.Name)
|
||
sb.WriteString("echo ''\n")
|
||
continue
|
||
}
|
||
sb.Write(data)
|
||
sb.WriteString("\n")
|
||
}
|
||
|
||
// 温度监控收尾
|
||
if hasTemp {
|
||
sb.Write(r.readScript("temp_end"))
|
||
sb.WriteString("\n")
|
||
}
|
||
|
||
// dmesg + 结束标记
|
||
sb.Write(r.readScript("dmesg"))
|
||
|
||
return sb.String()
|
||
}
|
||
|
||
// readScript 读取 scripts/ 目录下的脚本文件
|
||
func (r *Runner) readScript(name string) []byte {
|
||
data, err := os.ReadFile(fmt.Sprintf("scripts/%s.sh", name))
|
||
if err != nil {
|
||
return nil
|
||
}
|
||
return data
|
||
}
|
||
|
||
// testPriority 测试执行顺序：单项测试在前，综合测试(full)最后，监控项最后
|
||
func testPriority(m Metric) int {
|
||
if m.IsMonitor {
|
||
return 1000
|
||
}
|
||
switch TestType(m.Name) {
|
||
case TestCPU:
|
||
return 0
|
||
case TestMemory:
|
||
return 1
|
||
case TestDiskIO:
|
||
return 2
|
||
case TestMemNative:
|
||
return 3
|
||
case TestFull:
|
||
return 100
|
||
default:
|
||
return 50
|
||
}
|
||
}
|
||
|
||
// ============================
|
||
// 业务入口（workflow 调用）
|
||
// ============================
|
||
|
||
// TestDevice 对单台设备执行 SSH 登录 + 压力测试（业务入口）
|
||
func TestDevice(ip string, cfg config.Config) *Report {
|
||
sshCfg := sshclient.NewConfig(cfg.SSH.User, cfg.SSH.Password, cfg.SSH.KeyFile, cfg.SSH.Port, cfg.Scan.Timeout)
|
||
sshCfg.SudoPwd = cfg.SSH.SudoPwd
|
||
|
||
conn, err := sshclient.Connect(ip, sshCfg)
|
||
if err != nil {
|
||
return NewSSHFailReport(ip, err)
|
||
}
|
||
defer conn.Close()
|
||
|
||
// fnOS 的 admin 用户 HOME 目录可能不存在，登录会输出
|
||
// "Could not chdir to home directory" 横幅污染命令输出，先建好 HOME
|
||
sshclient.EnsureHome(conn)
|
||
|
||
// sudo 密码：优先用专属 sudo_password；为空时退回 SSH 登录密码（飞牛 admin 同密码）
|
||
sudoPwd := cfg.SSH.SudoPwd
|
||
if sudoPwd == "" {
|
||
sudoPwd = cfg.SSH.Password
|
||
}
|
||
|
||
// SSH 登录成功后回填远端 MAC（发现阶段只做了 ping，暂无 MAC）
|
||
if mac := queryRemoteMAC(conn); mac != "" {
|
||
if dev, ok := discovery.Devices[ip]; ok {
|
||
dev.MAC = mac
|
||
discovery.Devices[ip] = dev
|
||
}
|
||
fmt.Printf(" [MAC] %s -> %s\n", ip, mac)
|
||
}
|
||
|
||
// 由默认配置 + yaml 映射构造（避免硬编码，便于后续扩展字段）
|
||
stressCfg := DefaultConfig()
|
||
stressCfg.Duration = cfg.Stress.Duration
|
||
stressCfg.Threads = cfg.Stress.Threads
|
||
stressCfg.MemSizeMB = cfg.Stress.MemSizeMB
|
||
stressCfg.DiskSizeMB = cfg.Stress.DiskSizeMB
|
||
stressCfg.DiskDir = cfg.Stress.DiskDir
|
||
stressCfg.TempLogInt = cfg.Stress.TempInterval
|
||
stressCfg.TempLimit = cfg.Stress.TempLimit
|
||
stressCfg.UseMirror = cfg.Stress.UseMirror
|
||
|
||
return NewRunner(stressCfg, ParseTypes(cfg.Stress.Types), conn, cfg.Stress.AutoInstall, sudoPwd).Run(ip)
|
||
}
|
||
|
||
// NewSSHFailReport 构造 SSH 连接失败报告
|
||
func NewSSHFailReport(ip string, err error) *Report {
|
||
return &Report{
|
||
IP: ip,
|
||
StartTime: time.Now(),
|
||
Results: []Result{{Type: "ssh", Status: StatusFail, Error: err.Error()}},
|
||
Failed: 1,
|
||
}
|
||
}
|
||
|
||
// queryRemoteMAC 查询远端主机的 MAC 地址（取第一个有效的单播地址）。
|
||
// 失败或无有效地址时返回空字符串，不阻断主流程。
|
||
func queryRemoteMAC(client *ssh.Client) string {
|
||
out, err := sshclient.RunCommand(client, "cat /sys/class/net/*/address 2>/dev/null")
|
||
if err != nil {
|
||
return ""
|
||
}
|
||
for _, line := range strings.Split(out, "\n") {
|
||
mac := strings.TrimSpace(line)
|
||
if mac == "" {
|
||
continue
|
||
}
|
||
hw, e := net.ParseMAC(mac)
|
||
if e != nil || len(hw) != 6 {
|
||
continue
|
||
}
|
||
// 排除零地址、广播、组播
|
||
if hw[0] == 0 && hw[1] == 0 && hw[2] == 0 && hw[3] == 0 && hw[4] == 0 && hw[5] == 0 {
|
||
continue
|
||
}
|
||
if hw[0]&0x01 != 0 {
|
||
continue
|
||
}
|
||
return hw.String()
|
||
}
|
||
return ""
|
||
}
|
||
|
||
// ============================
|
||
// 工具函数
|
||
// ============================
|
||
|
||
// IsPassed 报告是否通过
|
||
func IsPassed(rpt *Report) bool {
|
||
return rpt != nil && rpt.Failed == 0 && rpt.Errors == 0
|
||
}
|
||
|
||
// IsSSHFail 报告是否为 SSH 连接失败（连不上，无任何有效测试项）
|
||
func IsSSHFail(rpt *Report) bool {
|
||
if rpt == nil || len(rpt.Results) != 1 {
|
||
return false
|
||
}
|
||
r := rpt.Results[0]
|
||
return r.Type == "ssh" && (r.Status == StatusFail || r.Status == StatusError)
|
||
}
|
||
|
||
// Status 报告状态字符串
|
||
func Status(rpt *Report) string {
|
||
if rpt == nil {
|
||
return StatusUnknown
|
||
}
|
||
if IsPassed(rpt) {
|
||
return StatusPass
|
||
}
|
||
return StatusFail
|
||
}
|
||
|
||
// ParseTypes 逗号分隔字符串 → TestType 列表
|
||
func ParseTypes(s string) []TestType {
|
||
var types []TestType
|
||
for _, t := range strings.Split(s, ",") {
|
||
if t = strings.TrimSpace(t); t != "" {
|
||
types = append(types, TestType(t))
|
||
}
|
||
}
|
||
return types
|
||
}
|