membank/auto-check/pkg/stress/runner.go

package stress

import (
	"fmt"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
)

// Runner 压力测试调度器
type Runner struct {
	Config  Config
	Types   []TestType
	Executor Executor
}

// NewRunner 创建调度器
func NewRunner(cfg Config, types []TestType, executor Executor) *Runner {
	return &Runner{
		Config:  cfg,
		Types:   types,
		Executor: executor,
	}
}

// Run 执行压力测试（单台主机）
func (r *Runner) Run(client *ssh.Client, ip string) *Report {
	report := &Report{IP: ip, StartTime: time.Now()}
	fmt.Printf("\n════════ [%s] 压力测试开始 ════════\n", ip)

	// 系统信息探测
	info := ProbeSystem(client, r.Executor)
	fmt.Printf("  主机: %s | CPU: %s (%s核) | 内存: %s | 内核: %s\n",
		info.Hostname, info.CPUModel, info.CPUCores, info.MemTotal, info.KernelVer)

	// 启动温度监控
	tempMon := NewTempMonitor(client, r.Executor, r.Config.TempLogInt)
	tempMon.Start()

	// 启动 dmesg 监控
	dmesgMon := NewDmesgMonitor(client, r.Executor)
	dmesgMon.Start()

	// 逐项执行测试
	for _, t := range r.Types {
		fmt.Printf("\n  ── 开始: %s ──\n", t)
		var res Result

		switch t {
		case TestCPU:
			res = CPUStress(client, r.Executor, r.Config)
		case TestMemory:
			res = MemoryStress(client, r.Executor, r.Config)
		case TestDiskIO:
			res = DiskIOStress(client, r.Executor, r.Config)
		case TestMemNative:
			res = MemNativeStress(client, r.Executor, r.Config)
		case TestFull:
			res = FullStress(client, r.Executor, r.Config)
		default:
			res = Result{Type: t, Status: "skip", Error: fmt.Sprintf("未知测试类型: %s", t)}
		}

		report.AddResult(res)
		fmt.Printf("  ── 完成: %s [%s] %s ──\n", t, res.Status, res.Duration.Round(time.Millisecond))
	}

	// 停止监控，收集结果
	tempLogs := tempMon.Stop()
	dmesgErrors := dmesgMon.Stop()

	// 添加温度报告
	if len(tempLogs) > 0 {
		maxTemp := ""
		for _, line := range tempLogs {
			// 提取温度值
			parts := strings.Fields(line)
			for _, p := range parts {
				if strings.HasSuffix(p, "°C") || strings.HasSuffix(p, "C") {
					maxTemp = p
				}
			}
		}
		tempSummary := fmt.Sprintf("采样 %d 次", len(tempLogs))
		if maxTemp != "" {
			tempSummary += ", 最高温度: " + maxTemp
		}
		status := "pass"
		if strings.Contains(strings.ToLower(strings.Join(tempLogs, " ")), "throttl") {
			status = "fail"
			tempSummary += " [检测到降频!]"
		}
		report.AddResult(Result{
			Type:     MonitorTemp,
			Status:   status,
			Output:   tempSummary,
			Duration: time.Since(report.StartTime),
		})
	}

	// 添加 dmesg 报告
	if len(dmesgErrors) > 0 {
		report.AddResult(Result{
			Type:     MonitorDmesg,
			Status:   "fail",
			Output:   fmt.Sprintf("检测到 %d 条硬件相关报错:\n%s", len(dmesgErrors), strings.Join(dmesgErrors[:min(10, len(dmesgErrors))], "\n")),
			Duration: time.Since(report.StartTime),
		})
	} else {
		report.AddResult(Result{
			Type:     MonitorDmesg,
			Status:   "pass",
			Output:   "无硬件相关内核报错",
			Duration: time.Since(report.StartTime),
		})
	}

	report.EndTime = time.Now()
	report.Duration = report.EndTime.Sub(report.StartTime)

	fmt.Printf("\n════════ [%s] 压力测试完成 ════════\n", ip)
	fmt.Println(report.ToText())
	return report
}

// RunAll 对多台主机执行压力测试
func (r *Runner) RunAll(clients map[string]*ssh.Client) []*Report {
	var reports []*Report
	for ip, client := range clients {
		reports = append(reports, r.Run(client, ip))
	}
	return reports
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}