fix: 流量统计异常

nezhahq · Aug 12, 2024 · 6abc0fc · 6abc0fc
1 parent b1d77a1
commit 6abc0fc
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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)