upd

internal/sysinfo/sysinfo.go

@@ -0,0 +1,151 @@
+// Package sysinfo 提供系统资源检测和 worker 数量自动计算功能。
+package sysinfo
+
+import (
+	"os/exec"
+	"runtime"
+	"strconv"
+	"strings"
+)
+
+// SystemInfo 系统资源信息
+type SystemInfo struct {
+	CPUCores int     // CPU 逻辑核心数
+	MemoryGB float64 // 内存大小（GB）
+	MemoryMB int     // 内存大小（MB）
+}
+
+// GetSystemInfo 获取系统资源信息
+func GetSystemInfo() SystemInfo {
+	info := SystemInfo{
+		CPUCores: runtime.NumCPU(),
+	}
+
+	// 获取内存信息
+	info.MemoryMB = getMemoryMB()
+	info.MemoryGB = float64(info.MemoryMB) / 1024.0
+
+	return info
+}
+
+// getMemoryMB 获取系统内存大小（MB）
+func getMemoryMB() int {
+	switch runtime.GOOS {
+	case "windows":
+		return getMemoryWindows()
+	case "linux":
+		return getMemoryLinux()
+	case "darwin":
+		return getMemoryDarwin()
+	default:
+		return 0
+	}
+}
+
+// getMemoryWindows 使用 wmic 获取 Windows 内存
+func getMemoryWindows() int {
+	cmd := exec.Command("wmic", "computersystem", "get", "totalphysicalmemory", "/value")
+	output, err := cmd.Output()
+	if err != nil {
+		return 0
+	}
+
+	// 解析输出：TotalPhysicalMemory=1234567890
+	for _, line := range strings.Split(string(output), "\n") {
+		line = strings.TrimSpace(line)
+		if strings.HasPrefix(line, "TotalPhysicalMemory=") {
+			value := strings.TrimPrefix(line, "TotalPhysicalMemory=")
+			if bytes, err := strconv.ParseInt(value, 10, 64); err == nil {
+				return int(bytes / 1024 / 1024)
+			}
+		}
+	}
+	return 0
+}
+
+// getMemoryLinux 从 /proc/meminfo 获取 Linux 内存
+func getMemoryLinux() int {
+	cmd := exec.Command("cat", "/proc/meminfo")
+	output, err := cmd.Output()
+	if err != nil {
+		return 0
+	}
+
+	// 解析 MemTotal: 12345678 kB
+	for _, line := range strings.Split(string(output), "\n") {
+		if strings.HasPrefix(line, "MemTotal:") {
+			fields := strings.Fields(line)
+			if len(fields) >= 2 {
+				if kb, err := strconv.ParseInt(fields[1], 10, 64); err == nil {
+					return int(kb / 1024)
+				}
+			}
+		}
+	}
+	return 0
+}
+
+// getMemoryDarwin 使用 sysctl 获取 macOS 内存
+func getMemoryDarwin() int {
+	cmd := exec.Command("sysctl", "-n", "hw.memsize")
+	output, err := cmd.Output()
+	if err != nil {
+		return 0
+	}
+
+	value := strings.TrimSpace(string(output))
+	if bytes, err := strconv.ParseInt(value, 10, 64); err == nil {
+		return int(bytes / 1024 / 1024)
+	}
+	return 0
+}
+
+// WorkerConfig worker 配置
+type WorkerConfig struct {
+	GoWorkers        int // Go 端 worker 数量（探测/MIME/移动）
+	NsfwWorkers      int // NudeNet Python 推理 worker 数量
+	AnimeNsfwWorkers int // WD Tagger Python 推理 worker 数量
+}
+
+// CalculateWorkerConfig 根据系统资源自动计算 worker 配置
+// 规则：
+// - Go worker: CPU核心数 - 2（留2核给Python推理和系统）
+// - NudeNet worker: 因为模型小（~300MB），可以开更多，上限为 CPU核心数
+// - WD Tagger worker: 因为模型大（~1.5GB），受内存限制，每个worker约需2GB
+func CalculateWorkerConfig(info SystemInfo) WorkerConfig {
+	config := WorkerConfig{}
+
+	// Go worker: CPU核心数 - 2，最少为1
+	goWorkers := info.CPUCores - 2
+	if goWorkers < 1 {
+		goWorkers = 1
+	}
+	config.GoWorkers = goWorkers
+
+	// NudeNet worker: 模型小，可以开到 CPU核心数，但留一些给系统
+	nsfwWorkers := info.CPUCores
+	if nsfwWorkers > 12 {
+		nsfwWorkers = 12 // 上限12，避免过度并发
+	}
+	config.NsfwWorkers = nsfwWorkers
+
+	// WD Tagger worker: 受内存限制，每个worker约需2GB
+	// 预留4GB给系统和其他进程
+	availableMemoryGB := info.MemoryGB - 4
+	if availableMemoryGB < 0 {
+		availableMemoryGB = 0
+	}
+	animeWorkers := int(availableMemoryGB / 2)
+	if animeWorkers < 1 {
+		animeWorkers = 1
+	}
+	if animeWorkers > info.CPUCores {
+		animeWorkers = info.CPUCores // 不超过CPU核心数
+	}
+	if animeWorkers > 12 {
+		animeWorkers = 12 // 上限12
+	}
+	config.AnimeNsfwWorkers = animeWorkers
+
+	return config
+}