From c3364df6cb7dc2057e84202673d4dd092c3a8666 Mon Sep 17 00:00:00 2001
From: GoEdgeLab <goedgecloud@gmail.com>
Date: Sat, 7 Oct 2023 09:39:37 +0800
Subject: [PATCH] =?UTF-8?q?=E5=86=99=E6=96=87=E4=BB=B6=E5=A2=9E=E5=8A=A0?=
 =?UTF-8?q?=E8=B4=9F=E8=BD=BD=E4=BF=9D=E6=8A=A4?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/utils/fs/disk.go   | 25 +++++++++++++++++++++++--
 internal/utils/fs/status.go | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/internal/utils/fs/disk.go b/internal/utils/fs/disk.go
index 89a8549..c6eef8a 100644
--- a/internal/utils/fs/disk.go
+++ b/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
+	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
diff --git a/internal/utils/fs/status.go b/internal/utils/fs/status.go
index c60be56..659f447 100644
--- a/internal/utils/fs/status.go
+++ b/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
 }