写文件增加负载保护

internal/utils/fs/disk.go

@@ -4,11 +4,20 @@ 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"`
+}
+
 // CheckDiskWritingSpeed test disk writing speed
 func CheckDiskWritingSpeed() (speedMB float64, err error) {
 	var tempDir = os.TempDir()
@@ -66,8 +75,13 @@ func CheckDiskIsFast() (speedMB float64, isFast bool, err error) {
 	if err != nil {
 		return
 	}
+
 	isFast = speedMB > 150
 
+	if speedMB <= DiskSpeedMB {
+		return
+	}
+
 	if speedMB > 1000 {
 		DiskSpeed = SpeedExtremelyFast
 	} else if speedMB > 150 {
@@ -79,8 +93,15 @@ func CheckDiskIsFast() (speedMB float64, isFast bool, err error) {
 	}
 	calculateDiskMaxWrites()
 
-	if speedMB > DiskSpeedMB {
 	DiskSpeedMB = speedMB
+
+	// write to local file
+	cacheData, jsonErr := json.Marshal(&DiskSpeedCache{
+		Speed:   DiskSpeed,
+		SpeedMB: DiskSpeedMB,
+	})
+	if jsonErr == nil {
+		_ = os.WriteFile(Tea.Root+"/data/"+diskSpeedDataFile, cacheData, 0666)
 	}
 
 	return

internal/utils/fs/status.go

@@ -3,8 +3,11 @@
 package fsutils
 
 import (
+	"encoding/json"
 	teaconst "github.com/TeaOSLab/EdgeNode/internal/const"
 	"github.com/iwind/TeaGo/Tea"
+	"github.com/shirou/gopsutil/v3/load"
+	"os"
 	"sync/atomic"
 	"time"
 )
@@ -38,6 +41,14 @@ var (
 	DiskSpeedMB   float64
 )
 
+var IsInHighLoad = false
+var IsInExtremelyHighLoad = false
+
+const (
+	highLoad1Threshold          = 20
+	extremelyHighLoad1Threshold = 40
+)
+
 func init() {
 	if !teaconst.IsMain {
 		return
@@ -45,6 +56,18 @@ func init() {
 
 	// test disk
 	go func() {
+		// load last result from local disk
+		cacheData, cacheErr := os.ReadFile(Tea.Root + "/data/" + diskSpeedDataFile)
+		if cacheErr == nil {
+			var cache = &DiskSpeedCache{}
+			err := json.Unmarshal(cacheData, cache)
+			if err == nil && cache.SpeedMB > 0 {
+				DiskSpeedMB = cache.SpeedMB
+				DiskSpeed = cache.Speed
+				calculateDiskMaxWrites()
+			}
+		}
+
 		// initial check
 		_, _, _ = CheckDiskIsFast()
 
@@ -61,6 +84,16 @@ func init() {
 			}
 		}
 	}()
+
+	// check high load
+	go func() {
+		var ticker = time.NewTicker(5 * time.Second)
+		for range ticker.C {
+			stat, _ := load.Avg()
+			IsInExtremelyHighLoad = stat != nil && stat.Load1 > extremelyHighLoad1Threshold
+			IsInHighLoad = stat != nil && stat.Load1 > highLoad1Threshold && !DiskIsFast()
+		}
+	}()
 }
 
 func DiskIsFast() bool {
@@ -78,6 +111,10 @@ func DiskIsExtremelyFast() bool {
 var countWrites int32 = 0
 
 func WriteReady() bool {
+	if IsInExtremelyHighLoad {
+		return false
+	}
+
 	return atomic.LoadInt32(&countWrites) < DiskMaxWrites
 }