fix: 流量统计异常

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

model/notification_test.go

@@ -72,11 +72,11 @@ func execCase(t *testing.T, item testSt) {
 			UdpConnCount:   0,
 			ProcessCount:   0,
 		},
-		LastActive:            time.Time{},
-		TaskClose:             nil,
-		TaskStream:            nil,
-		PrevHourlyTransferIn:  0,
-		PrevHourlyTransferOut: 0,
+		LastActive:              time.Time{},
+		TaskClose:               nil,
+		TaskStream:              nil,
+		PrevTransferInSnapshot:  0,
+		PrevTransferOutSnapshot: 0,
 	}
 	ns := NotificationServerBundle{
 		Notification: &n,

model/rule.go

@@ -5,6 +5,8 @@ import (
 	"time"
 
 	"gorm.io/gorm"
+
+	"github.com/naiba/nezha/pkg/utils"
 )
 
 const (
@@ -103,21 +105,21 @@ func (u *Rule) Snapshot(cycleTransferStats *CycleTransferStats, server *Server,
 			src = float64(server.LastActive.Unix())
 		}
 	case "transfer_in_cycle":
-		src = float64(server.State.NetInTransfer - uint64(server.PrevHourlyTransferIn))
+		src = float64(utils.Uint64SubInt64(server.State.NetInTransfer, server.PrevTransferInSnapshot))
 		if u.CycleInterval != 0 {
 			var res NResult
 			db.Model(&Transfer{}).Select("SUM(`in`) AS n").Where("datetime(`created_at`) >= datetime(?) AND server_id = ?", u.GetTransferDurationStart().UTC(), server.ID).Scan(&res)
 			src += float64(res.N)
 		}
 	case "transfer_out_cycle":
-		src = float64(server.State.NetOutTransfer - uint64(server.PrevHourlyTransferOut))
+		src = float64(utils.Uint64SubInt64(server.State.NetOutTransfer, server.PrevTransferOutSnapshot))
 		if u.CycleInterval != 0 {
 			var res NResult
 			db.Model(&Transfer{}).Select("SUM(`out`) AS n").Where("datetime(`created_at`) >= datetime(?) AND server_id = ?", u.GetTransferDurationStart().UTC(), server.ID).Scan(&res)
 			src += float64(res.N)
 		}
 	case "transfer_all_cycle":
-		src = float64(server.State.NetOutTransfer - uint64(server.PrevHourlyTransferOut) + server.State.NetInTransfer - uint64(server.PrevHourlyTransferIn))
+		src = float64(utils.Uint64SubInt64(server.State.NetOutTransfer, server.PrevTransferOutSnapshot) + utils.Uint64SubInt64(server.State.NetInTransfer, server.PrevTransferInSnapshot))
 		if u.CycleInterval != 0 {
 			var res NResult
 			db.Model(&Transfer{}).Select("SUM(`in`+`out`) AS n").Where("datetime(`created_at`) >= datetime(?) AND server_id = ?", u.GetTransferDurationStart().UTC(), server.ID).Scan(&res)

model/server.go

@@ -3,6 +3,7 @@ package model
 import (
 	"fmt"
 	"html/template"
+	"sync"
 	"time"
 
 	"github.com/naiba/nezha/pkg/utils"
@@ -27,11 +28,12 @@ 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:"-"`
 
-	PrevHourlyTransferIn  int64 `gorm:"-" json:"-"` // 上次数据点时的入站使用量
-	PrevHourlyTransferOut int64 `gorm:"-" json:"-"` // 上次数据点时的出站使用量
+	PrevTransferInSnapshot  int64 `gorm:"-" json:"-"` // 上次数据点时的入站使用量
+	PrevTransferOutSnapshot int64 `gorm:"-" json:"-"` // 上次数据点时的出站使用量
 }
 
 func (s *Server) CopyFromRunningServer(old *Server) {
@@ -39,9 +41,10 @@ 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.PrevHourlyTransferIn = old.PrevHourlyTransferIn
-	s.PrevHourlyTransferOut = old.PrevHourlyTransferOut
+	s.PrevTransferInSnapshot = old.PrevTransferInSnapshot
+	s.PrevTransferOutSnapshot = old.PrevTransferOutSnapshot
 }
 
 func boolToString(b bool) string {

pkg/utils/utils.go

@@ -76,3 +76,13 @@ func GenerateRandomString(n int) (string, error) {
 	}
 	return string(ret), nil
 }
+
+func Uint64SubInt64(a uint64, b int64) uint64 {
+	if b < 0 {
+		return a + uint64(-b)
+	}
+	if a < uint64(b) {
+		return 0
+	}
+	return a - uint64(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
 }
@@ -112,10 +114,10 @@ func (s *NezhaHandler) ReportSystemState(c context.Context, r *pb.State) (*pb.Re
 	singleton.ServerList[clientID].LastActive = time.Now()
 	singleton.ServerList[clientID].State = &state
 
-	// 如果从未记录过，先打点，等到小时时间点时入库
-	if singleton.ServerList[clientID].PrevHourlyTransferIn == 0 || singleton.ServerList[clientID].PrevHourlyTransferOut == 0 {
-		singleton.ServerList[clientID].PrevHourlyTransferIn = int64(state.NetInTransfer)
-		singleton.ServerList[clientID].PrevHourlyTransferOut = int64(state.NetOutTransfer)
+	// 应对 dashboard 重启的情况，如果从未记录过，先打点，等到小时时间点时入库
+	if singleton.ServerList[clientID].PrevTransferInSnapshot == 0 || singleton.ServerList[clientID].PrevTransferOutSnapshot == 0 {
+		singleton.ServerList[clientID].PrevTransferInSnapshot = int64(state.NetInTransfer)
+		singleton.ServerList[clientID].PrevTransferOutSnapshot = int64(state.NetOutTransfer)
 	}
 
 	return &pb.Receipt{Proced: true}, nil
@@ -135,9 +137,8 @@ func (s *NezhaHandler) ReportSystemInfo(c context.Context, r *pb.Host) (*pb.Rece
 	// 检查并更新DDNS
 	if singleton.Conf.DDNS.Enable &&
 		singleton.ServerList[clientID].EnableDDNS &&
-		singleton.ServerList[clientID].Host != nil &&
 		host.IP != "" &&
-		singleton.ServerList[clientID].Host.IP != host.IP {
+		(singleton.ServerList[clientID].Host == nil || singleton.ServerList[clientID].Host.IP != host.IP) {
 		serverDomain := singleton.ServerList[clientID].DDNSDomain
 		if singleton.Conf.DDNS.Provider == "" {
 			provider, err = singleton.GetDDNSProviderFromProfile(singleton.ServerList[clientID].DDNSProfile)
@@ -164,10 +165,9 @@ func (s *NezhaHandler) ReportSystemInfo(c context.Context, r *pb.Host) (*pb.Rece
 	}
 
 	// 发送IP变动通知
-	if singleton.Conf.EnableIPChangeNotification &&
+	if singleton.ServerList[clientID].Host != nil && singleton.Conf.EnableIPChangeNotification &&
 		((singleton.Conf.Cover == model.ConfigCoverAll && !singleton.Conf.IgnoredIPNotificationServerIDs[clientID]) ||
 			(singleton.Conf.Cover == model.ConfigCoverIgnoreAll && singleton.Conf.IgnoredIPNotificationServerIDs[clientID])) &&
-		singleton.ServerList[clientID].Host != nil &&
 		singleton.ServerList[clientID].Host.IP != "" &&
 		host.IP != "" &&
 		singleton.ServerList[clientID].Host.IP != host.IP {
@@ -184,10 +184,14 @@ func (s *NezhaHandler) ReportSystemInfo(c context.Context, r *pb.Host) (*pb.Rece
 			nil)
 	}
 
-	// 判断是否是机器重启，如果是机器重启要录入最后记录的流量里面
-	if singleton.ServerList[clientID].Host.BootTime < host.BootTime {
-		singleton.ServerList[clientID].PrevHourlyTransferIn = singleton.ServerList[clientID].PrevHourlyTransferIn - int64(singleton.ServerList[clientID].State.NetInTransfer)
-		singleton.ServerList[clientID].PrevHourlyTransferOut = singleton.ServerList[clientID].PrevHourlyTransferOut - int64(singleton.ServerList[clientID].State.NetOutTransfer)
+	/**
+	 * 这里的 singleton 中的数据都是关机前的旧数据
+	 * 当 agent 重启时，bootTime 变大，agent 端会先上报 host 信息，然后上报 state 信息
+	 * 这是可以借助上报顺序的空档，将停机前的流量统计数据标记下来，加到下一个小时的数据点上
+	 */
+	if singleton.ServerList[clientID].Host != nil && singleton.ServerList[clientID].Host.BootTime < host.BootTime {
+		singleton.ServerList[clientID].PrevTransferInSnapshot = singleton.ServerList[clientID].PrevTransferInSnapshot - int64(singleton.ServerList[clientID].State.NetInTransfer)
+		singleton.ServerList[clientID].PrevTransferOutSnapshot = singleton.ServerList[clientID].PrevTransferOutSnapshot - int64(singleton.ServerList[clientID].State.NetOutTransfer)
 	}
 
 	// 不要冲掉国家码

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)

service/singleton/singleton.go

@@ -100,14 +100,14 @@ func RecordTransferHourlyUsage() {
 	for id, server := range ServerList {
 		tx := model.Transfer{
 			ServerID: id,
-			In:       server.State.NetInTransfer - uint64(server.PrevHourlyTransferIn),
-			Out:      server.State.NetOutTransfer - uint64(server.PrevHourlyTransferOut),
+			In:       utils.Uint64SubInt64(server.State.NetInTransfer, server.PrevTransferInSnapshot),
+			Out:      utils.Uint64SubInt64(server.State.NetOutTransfer, server.PrevTransferOutSnapshot),
 		}
 		if tx.In == 0 && tx.Out == 0 {
 			continue
 		}
-		server.PrevHourlyTransferIn = int64(server.State.NetInTransfer)
-		server.PrevHourlyTransferOut = int64(server.State.NetOutTransfer)
+		server.PrevTransferInSnapshot = int64(server.State.NetInTransfer)
+		server.PrevTransferOutSnapshot = int64(server.State.NetOutTransfer)
 		tx.CreatedAt = nowTrimSeconds
 		txs = append(txs, tx)
 	}