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

package stress

import (
	"fmt"
	"strings"
	"time"

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

// truncate 截断过长输出
func truncate(s string, maxLen int) string {
	s = strings.TrimSpace(s)
	if len(s) > maxLen {
		return s[:maxLen] + "\n  ... (输出已截断)"
	}
	return s
}

// ============================
// stress-ng 压测
// ============================

// CPUStress stress-ng CPU 压力测试
func CPUStress(client *ssh.Client, ex Executor, cfg Config) Result {
	start := time.Now()
	fmt.Printf("    [CPU] stress-ng CPU 压力 (%d线程, %v)...\n", cfg.Threads, cfg.Duration)

	info, err := EnsureTool(client, ex, "stress-ng", "apt install stress-ng / opkg install stress-ng")
	if err != nil {
		return Result{Type: TestCPU, Status: "skip", Error: err.Error(), Duration: time.Since(start)}
	}
	fmt.Printf("    [CPU] 使用 %s (%s)\n", info.Path, info.Version)

	cmd := fmt.Sprintf("%s --cpu %d --cpu-method all --timeout %v --metrics-brief --temp-path /tmp 2>&1",
		info.Path, cfg.Threads, cfg.Duration)
	out, err := exec(client, ex, cmd)
	duration := time.Since(start)

	if err != nil {
		return Result{Type: TestCPU, Status: "error", Output: truncate(out, 600), Error: err.Error(), Duration: duration}
	}
	return Result{Type: TestCPU, Status: "pass", Output: truncate(out, 600), Duration: duration}
}

// MemoryStress stress-ng 内存压力测试
func MemoryStress(client *ssh.Client, ex Executor, cfg Config) Result {
	start := time.Now()
	fmt.Printf("    [Memory] stress-ng 内存压力 (%d线程, %v)...\n", cfg.Threads, cfg.Duration)

	info, err := EnsureTool(client, ex, "stress-ng", "apt install stress-ng / opkg install stress-ng")
	if err != nil {
		return Result{Type: TestMemory, Status: "skip", Error: err.Error(), Duration: time.Since(start)}
	}

	memWorkers := cfg.Threads
	if memWorkers > 4 {
		memWorkers = 4
	}
	cmd := fmt.Sprintf("%s --vm %d --vm-bytes 256M --vm-method all --timeout %v --metrics-brief 2>&1",
		info.Path, memWorkers, cfg.Duration)
	out, err := exec(client, ex, cmd)
	duration := time.Since(start)

	if err != nil {
		return Result{Type: TestMemory, Status: "error", Output: truncate(out, 600), Error: err.Error(), Duration: duration}
	}
	return Result{Type: TestMemory, Status: "pass", Output: truncate(out, 600), Duration: duration}
}

// DiskIOStress stress-ng 磁盘IO压力测试
func DiskIOStress(client *ssh.Client, ex Executor, cfg Config) Result {
	start := time.Now()
	fmt.Printf("    [DiskIO] stress-ng 磁盘IO压力 (%v)...\n", cfg.Duration)

	info, err := EnsureTool(client, ex, "stress-ng", "apt install stress-ng / opkg install stress-ng")
	if err != nil {
		return Result{Type: TestDiskIO, Status: "skip", Error: err.Error(), Duration: time.Since(start)}
	}

	testDir := cfg.DiskDir
	if testDir == "" {
		testDir = "/tmp/stress-disk-test"
	}
	_, _ = exec(client, ex, fmt.Sprintf("mkdir -p %s", testDir))

	cmd := fmt.Sprintf("%s --iomix 2 --iomix-bytes %dM --timeout %v --metrics-brief 2>&1",
		info.Path, cfg.DiskSizeMB, cfg.Duration)
	out, err := exec(client, ex, cmd)
	duration := time.Since(start)

	_, _ = exec(client, ex, fmt.Sprintf("rm -rf %s", testDir))

	if err != nil {
		return Result{Type: TestDiskIO, Status: "error", Output: truncate(out, 600), Error: err.Error(), Duration: duration}
	}
	return Result{Type: TestDiskIO, Status: "pass", Output: truncate(out, 600), Duration: duration}
}

// ============================
// stressapptest 内存精压
// ============================

// MemNativeStress stressapptest 内存稳定性精压
func MemNativeStress(client *ssh.Client, ex Executor, cfg Config) Result {
	start := time.Now()
	fmt.Printf("    [MemNative] stressapptest 内存精压 (%v)...\n", cfg.Duration)

	info, err := EnsureTool(client, ex, "stressapptest", "apt install stressapptest / opkg install stressapptest")
	if err != nil {
		return Result{Type: TestMemNative, Status: "skip", Error: err.Error(), Duration: time.Since(start)}
	}
	fmt.Printf("    [MemNative] 使用 %s\n", info.Path)

	memSize := cfg.MemSizeMB
	if memSize <= 0 {
		out, _ := exec(client, ex, "free -m 2>/dev/null | awk '/Mem:/{print int($7*0.6)}'")
		out = strings.TrimSpace(out)
		if out != "" {
			fmt.Sscanf(out, "%d", &memSize)
		}
		if memSize <= 0 {
			memSize = 512
		}
	}

	cmd := fmt.Sprintf("%s -s %d -M %d -f 0 -v 2>&1",
		info.Path, int(cfg.Duration.Seconds()), memSize)
	out, err := exec(client, ex, cmd)
	duration := time.Since(start)

	output := truncate(out, 800)
	if strings.Contains(strings.ToUpper(output), "PASS") {
		return Result{Type: TestMemNative, Status: "pass", Output: output, Duration: duration}
	}
	if strings.Contains(strings.ToUpper(output), "FAIL") {
		return Result{Type: TestMemNative, Status: "fail", Output: output, Duration: duration}
	}
	if err != nil {
		return Result{Type: TestMemNative, Status: "error", Output: output, Error: err.Error(), Duration: duration}
	}
	return Result{Type: TestMemNative, Status: "pass", Output: output, Duration: duration}
}

// ============================
// 综合压测
// ============================

// FullStress 三合一综合压测 (CPU + 内存 + 磁盘IO 同时进行)
func FullStress(client *ssh.Client, ex Executor, cfg Config) Result {
	start := time.Now()
	fmt.Printf("    [Full] 三合一综合压测 (CPU %d线程 + 内存 + 磁盘IO, %v)...\n", cfg.Threads, cfg.Duration)

	info, err := EnsureTool(client, ex, "stress-ng", "apt install stress-ng / opkg install stress-ng")
	if err != nil {
		return Result{Type: TestFull, Status: "skip", Error: err.Error(), Duration: time.Since(start)}
	}

	cmd := fmt.Sprintf("%s --cpu %d --vm 2 --vm-bytes 128M --iomix 1 --iomix-bytes 256M --timeout %v --metrics-brief 2>&1",
		info.Path, cfg.Threads, cfg.Duration)
	out, err := exec(client, ex, cmd)
	duration := time.Since(start)

	if err != nil {
		return Result{Type: TestFull, Status: "error", Output: truncate(out, 600), Error: err.Error(), Duration: duration}
	}
	return Result{Type: TestFull, Status: "pass", Output: truncate(out, 600), Duration: duration}
}