1.4.5.2

EdgeNode/internal/utils/fs/disk.go

@@ -0,0 +1,126 @@
+// Copyright 2023 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package fsutils
+
+import (
+	"bytes"
+	"encoding/json"
+	"github.com/iwind/TeaGo/Tea"
+	"math"
+	"os"
+	"time"
+)
+
+const diskSpeedDataFile = "disk.speed.json"
+
+type DiskSpeedCache struct {
+	Speed      Speed   `json:"speed"`
+	SpeedMB    float64 `json:"speedMB"`
+	CountTests int     `json:"countTests"` // test times
+}
+
+// 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
+	speedMB = math.Ceil(speedMB/10) * 10
+
+	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
+	}
+
+	// read old cached info
+	var cacheFile = Tea.Root + "/data/" + diskSpeedDataFile
+	var cacheInfo = &DiskSpeedCache{}
+	{
+		cacheData, cacheErr := os.ReadFile(cacheFile)
+		if cacheErr == nil {
+			var oldCacheInfo = &DiskSpeedCache{}
+			cacheErr = json.Unmarshal(cacheData, oldCacheInfo)
+			if cacheErr == nil {
+				cacheInfo = oldCacheInfo
+			}
+		}
+	}
+
+	cacheInfo.CountTests++
+
+	defer func() {
+		// write to local file
+		cacheData, jsonErr := json.Marshal(cacheInfo)
+		if jsonErr == nil {
+			_ = os.WriteFile(cacheFile, cacheData, 0666)
+		}
+	}()
+
+	isFast = speedMB > 150
+
+	if speedMB <= DiskSpeedMB {
+		return
+	}
+
+	if speedMB > 1000 {
+		DiskSpeed = SpeedExtremelyFast
+	} else if speedMB > 150 {
+		DiskSpeed = SpeedFast
+	} else if speedMB > 60 {
+		DiskSpeed = SpeedLow
+	} else {
+		DiskSpeed = SpeedExtremelySlow
+	}
+
+	DiskSpeedMB = speedMB
+
+	cacheInfo.Speed = DiskSpeed
+	cacheInfo.SpeedMB = DiskSpeedMB
+
+	return
+}