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

package stress

import (
	"fmt"
	"strings"
)

// ============================
// 脚本生成（根据指标生成远程执行的 shell 脚本）
// ============================

// BuildScript 根据指标列表和配置生成 self-contained shell 脚本
// 脚本输出结构化文本，便于 ParseResults 解析
func BuildScript(metrics []Metric, cfg Config) string {
	var sb strings.Builder

	// 脚本头部
	sb.WriteString("#!/bin/bash\n")
	sb.WriteString("set -e\n\n")

	// 系统信息探测
	sb.WriteString("echo '===SYSTEM_INFO==='\n")
	sb.WriteString("echo hostname:$(hostname)\n")
	sb.WriteString("echo cpu:$(lscpu 2>/dev/null | grep 'Model name' | sed 's/Model name:\\s*//' || echo unknown)\n")
	sb.WriteString("echo cores:$(nproc 2>/dev/null || echo 0)\n")
	sb.WriteString("echo memory:$(free -h 2>/dev/null | awk '/Mem:/{print $2}' || echo unknown)\n")
	sb.WriteString("echo kernel:$(uname -r 2>/dev/null || echo unknown)\n")
	sb.WriteString("echo disk:$(lsblk -d -o NAME,SIZE 2>/dev/null | head -3 | tr '\\n' ' ')\n")
	sb.WriteString("echo ''\n\n")

	// 启动温度监控（后台）
	hasTemp := false
	for _, m := range metrics {
		if m.Name == "temp" && m.Enabled {
			hasTemp = true
			break
		}
	}
	if hasTemp {
		tempInterval := int(cfg.TempLogInt.Seconds())
		if tempInterval <= 0 {
			tempInterval = 10
		}
		sb.WriteString(fmt.Sprintf("# 温度监控（后台）\n"))
		sb.WriteString("TEMP_LOG=/tmp/auto-check-temp.log\n")
		sb.WriteString("> $TEMP_LOG\n")
		sb.WriteString(fmt.Sprintf("(while true; do sensors 2>/dev/null | grep -i 'temp\\|core\\|cpu' | head -5 >> $TEMP_LOG; sleep %d; done) &\n", tempInterval))
		sb.WriteString("TEMP_PID=$!\n\n")
	}

	// 逐项执行测试
	for _, m := range metrics {
		if m.IsMonitor || !m.Enabled {
			continue
		}
		sb.WriteString(fmt.Sprintf("echo '===TEST:%s==='\n", m.Name))
		sb.WriteString("START_TIME=$(date +%%s%%N)\n")
		sb.WriteString("set +e\n")

		cmd := buildTestCommand(m, cfg)
		sb.WriteString(fmt.Sprintf("OUTPUT=$(%s 2>&1)\n", cmd))
		sb.WriteString("EXIT_CODE=$?\n")
		sb.WriteString("set -e\n")
		sb.WriteString("END_TIME=$(date +%%s%%N)\n")
		sb.WriteString("DURATION=$(( (END_TIME - START_TIME) / 1000000 ))\n")

		// 判定结果
		sb.WriteString("if [ $EXIT_CODE -eq 0 ]; then\n")
		sb.WriteString("  echo 'status:pass'\n")
		sb.WriteString("else\n")
		sb.WriteString("  echo 'status:fail'\n")
		sb.WriteString("fi\n")
		sb.WriteString("echo \"duration:${DURATION}ms\"\n")
		sb.WriteString("echo \"output:${OUTPUT}\"\n")
		sb.WriteString("echo ''\n\n")
	}

	// 收集温度监控结果
	if hasTemp {
		sb.WriteString("echo '===MONITOR:temp==='\n")
		sb.WriteString("if [ -f $TEMP_LOG ] && [ -s $TEMP_LOG ]; then\n")
		sb.WriteString("  echo 'status:pass'\n")
		sb.WriteString("  TEMP_MAX=$(grep -oP '\\d+\\.?\\d*°C' $TEMP_LOG | sort -t. -k1 -n | tail -1 || echo 'unknown')\n")
		sb.WriteString("  echo \"samples:$(wc -l < $TEMP_LOG)\"\n")
		sb.WriteString("  echo \"max:$TEMP_MAX\"\n")
		sb.WriteString("  echo \"output:$(tail -10 $TEMP_LOG)\"\n")
		sb.WriteString("else\n")
		sb.WriteString("  echo 'status:skip'\n")
		sb.WriteString("fi\n")
		sb.WriteString("echo ''\n")
		sb.WriteString("kill $TEMP_PID 2>/dev/null || true\n\n")
	}

	// 收集 dmesg 结果
	sb.WriteString("echo '===MONITOR:dmesg==='\n")
	sb.WriteString("DMESG_ERR=$(dmesg --level=err,crit,alert,emerg 2>/dev/null | tail -30 || true)\n")
	sb.WriteString("if [ -n \"$DMESG_ERR\" ]; then\n")
	sb.WriteString("  echo 'status:fail'\n")
	sb.WriteString("  echo \"output:$DMESG_ERR\"\n")
	sb.WriteString("else\n")
	sb.WriteString("  echo 'status:pass'\n")
	sb.WriteString("  echo 'output:无硬件相关内核报错'\n")
	sb.WriteString("fi\n")
	sb.WriteString("echo ''\n")
	sb.WriteString("echo '===END==='\n")

	return sb.String()
}

// buildTestCommand 根据指标生成测试命令
func buildTestCommand(m Metric, cfg Config) string {
	sec := int(cfg.Duration.Seconds())
	threads := cfg.Threads

	switch m.Name {
	case "cpu":
		return fmt.Sprintf("stress-ng --cpu %d --cpu-method all --timeout %ds --metrics-brief --temp-path /tmp", threads, sec)

	case "memory":
		workers := threads
		if workers > 4 {
			workers = 4
		}
		return fmt.Sprintf("stress-ng --vm %d --vm-bytes 256M --vm-method all --timeout %ds --metrics-brief", workers, sec)

	case "disk":
		testDir := cfg.DiskDir
		if testDir == "" {
			testDir = "/tmp/stress-disk-test"
		}
		diskMB := cfg.DiskSizeMB
		if diskMB <= 0 {
			diskMB = 1024
		}
		return fmt.Sprintf("mkdir -p %s && stress-ng --iomix 2 --iomix-bytes %dM --timeout %ds --metrics-brief && rm -rf %s",
			testDir, diskMB, sec, testDir)

	case "memnative":
		memMB := cfg.MemSizeMB
		if memMB <= 0 {
			memMB = 0 // 脚本内自动计算
		}
		if memMB > 0 {
			return fmt.Sprintf("stressapptest -s %d -M %d -f 0 -v", sec, memMB)
		}
		// 自动计算可用内存的 60%
		return fmt.Sprintf("M=$(free -m 2>/dev/null | awk '/Mem:/{print int($7*0.6)}' || echo 512) && stressapptest -s %d -M $M -f 0 -v", sec)

	case "full":
		return fmt.Sprintf("stress-ng --cpu %d --vm 2 --vm-bytes 128M --iomix 1 --iomix-bytes 256M --timeout %ds --metrics-brief",
			threads, sec)

	default:
		return "echo '未知测试类型'"
	}
}