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

package stress

import (
	"fmt"
	"strings"
	"time"

	"auto-check/pkg/config"
	"auto-check/pkg/sshclient"

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

// Results 测试结果表（IP → 报告），包级公开
var Results = make(map[string]*Report)

// 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
}

// IsPassed 报告是否通过
func IsPassed(rpt *Report) bool {
	return rpt != nil && rpt.Failed == 0 && rpt.Errors == 0
}

// Status 报告状态字符串：pass / fail / untested
func Status(rpt *Report) string {
	if rpt == nil {
		return "untested"
	}
	if IsPassed(rpt) {
		return "pass"
	}
	return "fail"
}

// 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
}

// TestDevice 对单台设备执行 SSH 登录 + 压力测试（业务入口）
func TestDevice(ip string, cfg config.Config) *Report {
	client := sshclient.NewClient(cfg.SSH.User, cfg.SSH.Password, cfg.SSH.KeyFile, cfg.SSH.Port, cfg.Scan.Timeout)

	conn, err := client.Connect(ip)
	if err != nil {
		fmt.Printf("  [测试] %s SSH 连接失败: %v\n", ip, err)
		return NewSSHFailReport(ip, err)
	}
	defer conn.Close()

	return NewRunner(
		Config{Duration: cfg.Stress.Duration, Threads: cfg.Stress.Threads, DiskSizeMB: 1024, TempLogInt: 10 * time.Second},
		ParseTypes(cfg.Stress.Types),
		client,
	).Run(conn, ip)
}

// NewSSHFailReport 构造 SSH 连接失败报告
func NewSSHFailReport(ip string, err error) *Report {
	return &Report{
		IP:        ip,
		StartTime: time.Now(),
		Results:   []Result{{Type: "ssh", Status: "fail", Error: err.Error()}},
		Failed:    1,
	}
}

// min 取较小值
func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}