389 lines
9.7 KiB
Go
389 lines
9.7 KiB
Go
package stress
|
||
|
||
import (
|
||
"fmt"
|
||
"net"
|
||
"os"
|
||
"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
|
||
}
|
||
|
||
// NewRunner 创建执行器（检测工具 → 构建指标 → 生成脚本 → 执行）
|
||
func NewRunner(cfg Config, types []TestType, client *ssh.Client) *Runner {
|
||
tools := DetectTools(client)
|
||
metrics := BuildMetrics(types, tools)
|
||
return &Runner{cfg: cfg, types: types, client: client, tools: tools, metrics: metrics}
|
||
}
|
||
|
||
// 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()
|
||
|
||
// 显示将执行的指标
|
||
fmt.Printf(" 指标: ")
|
||
for _, m := range r.metrics {
|
||
if m.Enabled {
|
||
fmt.Printf("%s ", m.Name)
|
||
}
|
||
}
|
||
fmt.Println()
|
||
|
||
// 按需读取各测试脚本，拼装成完整脚本
|
||
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)
|
||
}
|
||
|
||
// 生成采样折线图
|
||
for _, res := range report.Results {
|
||
if len(res.Samples) > 0 {
|
||
chartPath, err := WriteChartHTML(ip, res, "reports")
|
||
if err != nil {
|
||
fmt.Printf(" [图表] 生成失败: %v\n", err)
|
||
} else {
|
||
fmt.Printf(" [图表] %s\n", chartPath)
|
||
}
|
||
}
|
||
}
|
||
|
||
report.EndTime = time.Now()
|
||
report.Duration = report.EndTime.Sub(report.StartTime)
|
||
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
|
||
inSamples := false
|
||
|
||
flush := func() {
|
||
if currentTest == "" {
|
||
return
|
||
}
|
||
status := StatusPass
|
||
switch currentStatus {
|
||
case "fail":
|
||
status = StatusFail
|
||
case "skip":
|
||
status = StatusSkip
|
||
case "error":
|
||
status = StatusError
|
||
}
|
||
duration, _ := time.ParseDuration(currentDuration)
|
||
report.AddResult(Result{
|
||
Type: TestType(currentTest),
|
||
Status: status,
|
||
Output: strings.Join(currentOutput, "\n"),
|
||
Duration: duration,
|
||
Samples: currentSamples,
|
||
})
|
||
currentTest = ""
|
||
currentStatus = ""
|
||
currentDuration = ""
|
||
currentOutput = nil
|
||
currentSamples = nil
|
||
}
|
||
|
||
for _, line := range lines {
|
||
line = strings.TrimSpace(line)
|
||
|
||
// 采样数据块
|
||
if line == "===SAMPLES===" {
|
||
inSamples = true
|
||
continue
|
||
}
|
||
if line == "===END_SAMPLES===" {
|
||
inSamples = false
|
||
continue
|
||
}
|
||
if inSamples {
|
||
parts := strings.Split(line, ",")
|
||
if len(parts) >= 3 {
|
||
cpu, _ := strconv.ParseFloat(strings.TrimSpace(parts[1]), 64)
|
||
temp, _ := strconv.ParseFloat(strings.TrimSpace(parts[2]), 64)
|
||
currentSamples = append(currentSamples, Sample{
|
||
Time: strings.TrimSpace(parts[0]),
|
||
CPU: cpu,
|
||
Temp: temp,
|
||
})
|
||
}
|
||
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, "output:") {
|
||
currentOutput = append(currentOutput, strings.TrimPrefix(line, "output:"))
|
||
continue
|
||
}
|
||
if strings.HasPrefix(line, "samples:") || strings.HasPrefix(line, "max:") {
|
||
currentOutput = append(currentOutput, line)
|
||
continue
|
||
}
|
||
|
||
// 普通输出行
|
||
if currentTest != "" && line != "" {
|
||
currentOutput = append(currentOutput, line)
|
||
}
|
||
}
|
||
flush()
|
||
|
||
// 从系统信息输出解析主机信息
|
||
for _, line := range lines {
|
||
if strings.HasPrefix(line, "hostname:") {
|
||
// 可扩展：存入 report 的系统信息字段
|
||
}
|
||
}
|
||
}
|
||
|
||
// assembleScript 根据 metrics 读取对应的独立脚本文件，拼装成完整脚本
|
||
func (r *Runner) assembleScript() string {
|
||
var sb strings.Builder
|
||
|
||
// shebang + 环境变量导出
|
||
sb.WriteString("#!/bin/bash\n")
|
||
sb.WriteString("set -e\n\n")
|
||
|
||
// 导出参数变量
|
||
sb.WriteString(fmt.Sprintf("export AUTOCHECK_DURATION=%d\n", int(r.cfg.Duration.Seconds())))
|
||
sb.WriteString(fmt.Sprintf("export AUTOCHECK_THREADS=%d\n", r.cfg.Threads))
|
||
sb.WriteString(fmt.Sprintf("export AUTOCHECK_DISK_SIZE_MB=%d\n", r.cfg.DiskSizeMB))
|
||
sb.WriteString(fmt.Sprintf("export AUTOCHECK_MEM_SIZE_MB=%d\n", r.cfg.MemSizeMB))
|
||
sb.WriteString(fmt.Sprintf("export AUTOCHECK_TEMP_INTERVAL=%d\n", int(r.cfg.TempLogInt.Seconds())))
|
||
sb.WriteString("export AUTOCHECK_SAMPLE_INTERVAL=2\n")
|
||
if r.cfg.DiskDir != "" {
|
||
sb.WriteString(fmt.Sprintf("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")
|
||
}
|
||
|
||
// 逐个测试脚本
|
||
for _, m := range r.metrics {
|
||
if m.IsMonitor || !m.Enabled {
|
||
continue
|
||
}
|
||
data := r.readScript(m.Name)
|
||
if data == nil {
|
||
sb.WriteString(fmt.Sprintf("echo '===TEST:%s==='\n", m.Name))
|
||
sb.WriteString("echo 'status:skip'\n")
|
||
sb.WriteString(fmt.Sprintf("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
|
||
}
|
||
|
||
// ============================
|
||
// 业务入口（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)
|
||
|
||
conn, err := sshclient.Connect(ip, sshCfg)
|
||
if err != nil {
|
||
return NewSSHFailReport(ip, err)
|
||
}
|
||
defer conn.Close()
|
||
|
||
// 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)
|
||
}
|
||
|
||
stressCfg := Config{
|
||
Duration: cfg.Stress.Duration,
|
||
Threads: cfg.Stress.Threads,
|
||
DiskSizeMB: 1024,
|
||
TempLogInt: 10 * time.Second,
|
||
}
|
||
return NewRunner(stressCfg, ParseTypes(cfg.Stress.Types), conn).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
|
||
}
|
||
|
||
// 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
|
||
}
|