增加硬盘速度检测命令：edge-node disk speed

cmd/edge-node/main.go

@@ -7,6 +7,7 @@ import (
 	"github.com/TeaOSLab/EdgeNode/internal/apps"
 	teaconst "github.com/TeaOSLab/EdgeNode/internal/const"
 	"github.com/TeaOSLab/EdgeNode/internal/nodes"
+	fsutils "github.com/TeaOSLab/EdgeNode/internal/utils/fs"
 	_ "github.com/iwind/TeaGo/bootstrap"
 	"github.com/iwind/TeaGo/logs"
 	"github.com/iwind/TeaGo/maps"
@@ -365,6 +366,29 @@ func main() {
 		}
 		fmt.Println(string(statsJSON))
 	})
+	app.On("disk", func() {
+		var args = os.Args[2:]
+		if len(args) > 0 {
+			switch args[0] {
+			case "speed":
+				speedMB, isFast, err := fsutils.CheckDiskIsFast()
+				if err != nil {
+					fmt.Println("[ERROR]" + err.Error())
+				} else {
+					fmt.Printf("Speed: %.2fMB/s\n", speedMB)
+					if isFast {
+						fmt.Println("IsFast: true")
+					} else {
+						fmt.Println("IsFast: false")
+					}
+				}
+			default:
+				fmt.Println("Usage: edge-node disk [speed]")
+			}
+		} else {
+			fmt.Println("Usage: edge-node disk [speed]")
+		}
+	})
 	app.Run(func() {
 		var node = nodes.NewNode()
 		node.Start()

internal/nodes/node.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"github.com/TeaOSLab/EdgeCommon/pkg/configutils"
 	iplib "github.com/TeaOSLab/EdgeCommon/pkg/iplibrary"
 	"github.com/TeaOSLab/EdgeCommon/pkg/nodeconfigs"
@@ -26,6 +27,7 @@ import (
 	"github.com/TeaOSLab/EdgeNode/internal/utils"
 	_ "github.com/TeaOSLab/EdgeNode/internal/utils/agents" // 引入Agent管理器
 	_ "github.com/TeaOSLab/EdgeNode/internal/utils/clock"  // 触发时钟更新
+	fsutils "github.com/TeaOSLab/EdgeNode/internal/utils/fs"
 	"github.com/TeaOSLab/EdgeNode/internal/utils/jsonutils"
 	"github.com/TeaOSLab/EdgeNode/internal/waf"
 	"github.com/andybalholm/brotli"
@@ -141,7 +143,7 @@ func (this *Node) Start() {
 	// 调整系统参数
 	this.checkSystem()
 
-	// 检查硬盘类型
+	// 检查硬盘
 	this.checkDisk()
 
 	// 启动事件
@@ -1109,23 +1111,16 @@ func (this *Node) checkSystem() {
 
 // 检查硬盘
 func (this *Node) checkDisk() {
-	if runtime.GOOS != "linux" {
+	speedMB, isFast, err := fsutils.CheckDiskIsFast()
+	if err != nil {
+		remotelogs.Error("NODE", "check disk speed failed: "+err.Error())
 		return
 	}
-	for n := 'a'; n <= 'z'; n++ {
-		for _, path := range []string{
-			"/sys/block/vd" + string(n) + "/queue/rotational",
-			"/sys/block/sd" + string(n) + "/queue/rotational",
-		} {
-			data, err := os.ReadFile(path)
-			if err != nil {
-				continue
-			}
-			if string(data) == "0" {
-				teaconst.DiskIsFast = true
-			}
-			return
-		}
+	teaconst.DiskIsFast = isFast
+	if isFast {
+		remotelogs.Println("NODE", "disk is fast, writing test speed: "+fmt.Sprintf("%.2fMB/s", speedMB))
+	} else {
+		remotelogs.Println("NODE", "disk is slow, writing test speed: "+fmt.Sprintf("%.2fMB/s", speedMB))
 	}
 }
 

internal/utils/fs/disk.go

@@ -0,0 +1,69 @@
+// Copyright 2023 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package fsutils
+
+import (
+	"bytes"
+	"os"
+	"time"
+)
+
+// CheckDiskWritingSpeed test disk writing speed
+func CheckDiskWritingSpeed() (speedMB float64, err error) {
+	var tempDir = os.TempDir()
+	if len(tempDir) == 0 {
+		tempDir = "/tmp"
+	}
+
+	const filename = "edge-disk-writing-test.data"
+	var path = tempDir + "/" + filename
+	_ = os.Remove(path) // always try to delete the file
+
+	fp, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY, 0666)
+	if err != nil {
+		return 0, err
+	}
+
+	var isClosed bool
+	defer func() {
+		if !isClosed {
+			_ = fp.Close()
+		}
+
+		_ = os.Remove(path)
+	}()
+
+	var data = bytes.Repeat([]byte{'A'}, 16<<20)
+	var before = time.Now()
+	_, err = fp.Write(data)
+	if err != nil {
+		return 0, err
+	}
+
+	err = fp.Sync()
+	if err != nil {
+		return 0, err
+	}
+
+	err = fp.Close()
+	if err != nil {
+		return 0, err
+	}
+
+	var costSeconds = time.Since(before).Seconds()
+	speedMB = float64(len(data)) / (1 << 20) / costSeconds
+
+	isClosed = true
+
+	return
+}
+
+// CheckDiskIsFast check disk is 'fast' disk to write
+func CheckDiskIsFast() (speedMB float64, isFast bool, err error) {
+	speedMB, err = CheckDiskWritingSpeed()
+	if err != nil {
+		return
+	}
+	isFast = speedMB > 120
+	return
+}

internal/utils/fs/disk_test_test.go

@@ -0,0 +1,16 @@
+// Copyright 2023 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package fsutils_test
+
+import (
+	fsutils "github.com/TeaOSLab/EdgeNode/internal/utils/fs"
+	"testing"
+)
+
+func TestCheckDiskWritingSpeed(t *testing.T) {
+	t.Log(fsutils.CheckDiskWritingSpeed())
+}
+
+func TestCheckDiskIsFast(t *testing.T) {
+	t.Log(fsutils.CheckDiskIsFast())
+}