From 6abc0fce517a16b3a8d4ecf93f37ca6845213336 Mon Sep 17 00:00:00 2001
From: naiba <hi@nai.ba>
Date: Mon, 12 Aug 2024 10:06:55 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E6=B5=81=E9=87=8F=E7=BB=9F=E8=AE=A1?=
 =?UTF-8?q?=E5=BC=82=E5=B8=B8?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 cmd/dashboard/controller/member_api.go | 2 ++
 model/server.go                        | 7 +++++--
 pkg/utils/utils.go                     | 3 +++
 service/rpc/nezha.go                   | 2 ++
 service/singleton/server.go            | 1 +
 5 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/cmd/dashboard/controller/member_api.go b/cmd/dashboard/controller/member_api.go
index 3f6ef5e..e614516 100644
--- a/cmd/dashboard/controller/member_api.go
+++ b/cmd/dashboard/controller/member_api.go
@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/gin-gonic/gin"
@@ -384,6 +385,7 @@ func (ma *memberAPI) addOrEditServer(c *gin.Context) {
 	} else {
 		s.Host = &model.Host{}
 		s.State = &model.HostState{}
+		s.TaskCloseLock = new(sync.Mutex)
 		singleton.ServerLock.Lock()
 		singleton.SecretToID[s.Secret] = s.ID
 		singleton.ServerList[s.ID] = &s
diff --git a/model/server.go b/model/server.go
index d121bf2..043bada 100644
--- a/model/server.go
+++ b/model/server.go
@@ -3,6 +3,7 @@ package model
 import (
 	"fmt"
 	"html/template"
+	"sync"
 	"time"
 
 	"github.com/naiba/nezha/pkg/utils"
@@ -27,8 +28,9 @@ type Server struct {
 	State      *HostState `gorm:"-"`
 	LastActive time.Time  `gorm:"-"`
 
-	TaskClose  chan error                        `gorm:"-" json:"-"`
-	TaskStream pb.NezhaService_RequestTaskServer `gorm:"-" json:"-"`
+	TaskClose     chan error                        `gorm:"-" json:"-"`
+	TaskCloseLock *sync.Mutex                       `gorm:"-" json:"-"`
+	TaskStream    pb.NezhaService_RequestTaskServer `gorm:"-" json:"-"`
 
 	PrevTransferInSnapshot  int64 `gorm:"-" json:"-"` // 上次数据点时的入站使用量
 	PrevTransferOutSnapshot int64 `gorm:"-" json:"-"` // 上次数据点时的出站使用量
@@ -39,6 +41,7 @@ func (s *Server) CopyFromRunningServer(old *Server) {
 	s.State = old.State
 	s.LastActive = old.LastActive
 	s.TaskClose = old.TaskClose
+	s.TaskCloseLock = old.TaskCloseLock
 	s.TaskStream = old.TaskStream
 	s.PrevTransferInSnapshot = old.PrevTransferInSnapshot
 	s.PrevTransferOutSnapshot = old.PrevTransferOutSnapshot
diff --git a/pkg/utils/utils.go b/pkg/utils/utils.go
index fbb2bfe..6b0ed0f 100644
--- a/pkg/utils/utils.go
+++ b/pkg/utils/utils.go
@@ -81,5 +81,8 @@ func Uint64SubInt64(a uint64, b int64) uint64 {
 	if b < 0 {
 		return a + uint64(-b)
 	}
+	if a < uint64(b) {
+		return 0
+	}
 	return a - uint64(b)
 }
diff --git a/service/rpc/nezha.go b/service/rpc/nezha.go
index 72376a6..df5643a 100644
--- a/service/rpc/nezha.go
+++ b/service/rpc/nezha.go
@@ -90,12 +90,14 @@ func (s *NezhaHandler) RequestTask(h *pb.Host, stream pb.NezhaService_RequestTas
 	}
 	closeCh := make(chan error)
 	singleton.ServerLock.RLock()
+	singleton.ServerList[clientID].TaskCloseLock.Lock()
 	// 修复不断的请求 task 但是没有 return 导致内存泄漏
 	if singleton.ServerList[clientID].TaskClose != nil {
 		close(singleton.ServerList[clientID].TaskClose)
 	}
 	singleton.ServerList[clientID].TaskStream = stream
 	singleton.ServerList[clientID].TaskClose = closeCh
+	singleton.ServerList[clientID].TaskCloseLock.Unlock()
 	singleton.ServerLock.RUnlock()
 	return <-closeCh
 }
diff --git a/service/singleton/server.go b/service/singleton/server.go
index b1209fd..21b58a3 100644
--- a/service/singleton/server.go
+++ b/service/singleton/server.go
@@ -34,6 +34,7 @@ func loadServers() {
 		innerS := s
 		innerS.Host = &model.Host{}
 		innerS.State = &model.HostState{}
+		innerS.TaskCloseLock = new(sync.Mutex)
 		ServerList[innerS.ID] = &innerS
 		SecretToID[innerS.Secret] = innerS.ID
 		ServerTagToIDList[innerS.Tag] = append(ServerTagToIDList[innerS.Tag], innerS.ID)