package handler import ( "context" "encoding/json" "fmt" "net/http" "os/exec" "regexp" "strconv" "strings" "time" "github.com/gin-gonic/gin" "flowstack/internal/crypto" "flowstack/internal/database" "flowstack/internal/model" "flowstack/internal/sshexec" ) const ( ddFileSize = ".fs-disktest" ddSizeMB = 256 ddTimeout = 60 * time.Second ) type cmdRunner func(cmd string) (string, error) func localRunner(cmd string) (string, error) { ctx, cancel := context.WithTimeout(context.Background(), ddTimeout) defer cancel() out, err := exec.CommandContext(ctx, "sh", "-c", cmd).CombinedOutput() return string(out), err } // parseDDMBps 兼容两种 dd 输出: // Linux: "268435456 bytes (268 MB, 256 MiB) copied, 1.2 s, 224 MB/s" // macOS: "268435456 bytes transferred in 1.234 secs (217543210 bytes/sec)" var ( linuxSpeedRe = regexp.MustCompile(`([\d.]+)\s*([kMGT]?)B/s`) macSpeedRe = regexp.MustCompile(`\(([\d,]+)\s*bytes/sec\)`) ) func parseDDMBps(out string) float64 { if ms := linuxSpeedRe.FindAllStringSubmatch(out, -1); len(ms) > 0 { m := ms[len(ms)-1] v, err := strconv.ParseFloat(m[1], 64) if err != nil { return 0 } switch m[2] { case "k", "K": v /= 1024 case "G": v *= 1024 case "T": v *= 1024 * 1024 } return v } if m := macSpeedRe.FindStringSubmatch(out); m != nil { v, err := strconv.ParseFloat(strings.ReplaceAll(m[1], ",", ""), 64) if err == nil { return v / 1e6 } } return 0 } // TestBuildNodeDisk POST /api/build-nodes/:id/disk-test —— dd 实测磁盘读写 func TestBuildNodeDisk(c *gin.Context) { var n model.BuildNode if err := database.DB.First(&n, c.Param("id")).Error; err != nil { c.JSON(http.StatusNotFound, gin.H{"error": "build node not found"}) return } // 本地节点直接 exec,否则 SSH(需要凭证) isLocal := n.IP == "127.0.0.1" || n.IP == "localhost" || n.IP == "::1" var run cmdRunner mode := "ssh" if isLocal { run = localRunner mode = "local" } else { if n.CredentialID == nil { c.JSON(http.StatusBadRequest, gin.H{"error": "远程节点需先在编辑中关联 SSH 凭证"}) return } var cred model.Credential if err := database.DB.First(&cred, *n.CredentialID).Error; err != nil || cred.Content == "" { c.JSON(http.StatusBadRequest, gin.H{"error": "SSH 凭证不存在或未设置内容"}) return } plain, err := crypto.DecryptField(cred.Content) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": "凭证解密失败"}) return } var cc struct { PrivateKey string `json:"private_key"` Password string `json:"password"` SSHUser string `json:"ssh_user"` SSHPort int `json:"ssh_port"` } if err := json.Unmarshal([]byte(plain), &cc); err != nil { c.JSON(http.StatusBadRequest, gin.H{"error": "凭证内容格式错误"}) return } cfg := sshexec.Config{Host: n.IP, Port: cc.SSHPort, User: cc.SSHUser, Password: cc.Password, PrivateKey: cc.PrivateKey} run = func(cmd string) (string, error) { return sshexec.Run(cfg, cmd, ddTimeout) } } ws := n.WorkspacePath if ws == "" { ws = "/data/flowstack-workspace" } file := ws + "/" + ddFileSize start := time.Now() if _, err := run(fmt.Sprintf("mkdir -p '%s'", ws)); err != nil { c.JSON(http.StatusOK, gin.H{"success": false, "message": "创建工作空间失败: " + err.Error()}) return } defer func() { _, _ = run(fmt.Sprintf("rm -f '%s'", file)) }() // 写测试:direct 优先(绕过缓存),失败降级普通模式 writeCmd := fmt.Sprintf("dd if=/dev/zero of='%s' bs=1m count=%d oflag=direct 2>&1", file, ddSizeMB) writeOut, err := run(writeCmd) direct := err == nil if err != nil || parseDDMBps(writeOut) == 0 { writeOut, err = run(fmt.Sprintf("dd if=/dev/zero of='%s' bs=1m count=%d 2>&1", file, ddSizeMB)) direct = false } if err != nil { c.JSON(http.StatusOK, gin.H{"success": false, "message": "写入测试失败: " + firstLine(writeOut, err)}) return } writeMBps := parseDDMBps(writeOut) // 读测试 readCmd := fmt.Sprintf("dd if='%s' of=/dev/null bs=1m 2>&1", file) if direct { readCmd = fmt.Sprintf("dd if='%s' of=/dev/null bs=1m iflag=direct 2>&1", file) } readOut, err := run(readCmd) if err != nil && direct { readOut, err = run(fmt.Sprintf("dd if='%s' of=/dev/null bs=1m 2>&1", file)) } if err != nil { c.JSON(http.StatusOK, gin.H{"success": false, "message": "读取测试失败: " + firstLine(readOut, err)}) return } readMBps := parseDDMBps(readOut) if direct { mode += "/direct" } else { mode += "/cache" } // 回写节点 now := time.Now() database.DB.Model(&n).Updates(map[string]any{ "disk_write_mbps": writeMBps, "disk_read_mbps": readMBps, "disk_test_mode": mode, "disk_tested_at": &now, }) c.JSON(http.StatusOK, gin.H{ "success": true, "writeMBps": writeMBps, "readMBps": readMBps, "mode": mode, "elapsedMs": time.Since(start).Milliseconds(), "message": fmt.Sprintf("写入 %.0f MB/s · 读取 %.0f MB/s(%s)", writeMBps, readMBps, map[bool]string{true: "direct 直测", false: "含缓存"}[direct]), }) }