From c6e8052b5fb70d40dbdfa3415f9a413062130fdc Mon Sep 17 00:00:00 2001
From: GoEdgeLab <goedgecloud@gmail.com>
Date: Mon, 27 Feb 2023 10:48:16 +0800
Subject: [PATCH] =?UTF-8?q?=E4=B8=8A=E4=BC=A0=E5=B8=A6=E5=AE=BD=E6=95=B0?=
 =?UTF-8?q?=E6=8D=AE=E6=97=B6=E5=90=8C=E6=97=B6=E4=B8=8A=E4=BC=A0=E6=B5=81?=
 =?UTF-8?q?=E9=87=8F=E6=95=B0=E6=8D=AE?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/nodes/client_conn.go            |  5 ++---
 internal/nodes/listener_udp.go           |  2 +-
 internal/stats/bandwidth_stat_manager.go | 21 +++++++++++++--------
 3 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/internal/nodes/client_conn.go b/internal/nodes/client_conn.go
index cc338a4..369dd8b 100644
--- a/internal/nodes/client_conn.go
+++ b/internal/nodes/client_conn.go
@@ -180,7 +180,6 @@ func (this *ClientConn) Write(b []byte) (n int, err error) {
 	var before = time.Now()
 	n, err = this.rawConn.Write(b)
 	if n > 0 {
-
 		// 统计当前服务带宽
 		if this.serverId > 0 {
 			if !this.isLO || Tea.IsTesting() { // 环路不统计带宽，避免缓存预热等行为产生带宽
@@ -188,9 +187,9 @@ func (this *ClientConn) Write(b []byte) (n int, err error) {
 
 				var cost = time.Since(before).Seconds()
 				if cost > 1 {
-					stats.SharedBandwidthStatManager.Add(this.userId, this.serverId, int64(float64(n)/cost))
+					stats.SharedBandwidthStatManager.Add(this.userId, this.serverId, int64(float64(n)/cost), int64(n))
 				} else {
-					stats.SharedBandwidthStatManager.Add(this.userId, this.serverId, int64(n))
+					stats.SharedBandwidthStatManager.Add(this.userId, this.serverId, int64(n), int64(n))
 				}
 			}
 		}
diff --git a/internal/nodes/listener_udp.go b/internal/nodes/listener_udp.go
index 12e7c54..065f9b2 100644
--- a/internal/nodes/listener_udp.go
+++ b/internal/nodes/listener_udp.go
@@ -404,7 +404,7 @@ func NewUDPConn(server *serverconfigs.ServerConfig, addr net.Addr, proxyListener
 					stats.SharedTrafficStatManager.Add(server.Id, "", int64(n), 0, 0, 0, 0, 0, server.ShouldCheckTrafficLimit(), server.PlanId())
 
 					// 带宽
-					stats.SharedBandwidthStatManager.Add(server.UserId, server.Id, int64(n))
+					stats.SharedBandwidthStatManager.Add(server.UserId, server.Id, int64(n), int64(n))
 				}
 			}
 			if err != nil {
diff --git a/internal/stats/bandwidth_stat_manager.go b/internal/stats/bandwidth_stat_manager.go
index 266e232..c28d3d2 100644
--- a/internal/stats/bandwidth_stat_manager.go
+++ b/internal/stats/bandwidth_stat_manager.go
@@ -37,6 +37,7 @@ type BandwidthStat struct {
 	CurrentBytes     int64
 	CurrentTimestamp int64
 	MaxBytes         int64
+	TotalBytes       int64
 }
 
 // BandwidthStatManager 服务带宽统计
@@ -107,6 +108,7 @@ func (this *BandwidthStatManager) Loop() error {
 				Day:          stat.Day,
 				TimeAt:       stat.TimeAt,
 				Bytes:        stat.MaxBytes / bandwidthTimestampDelim,
+				TotalBytes:   stat.TotalBytes,
 				NodeRegionId: regionId,
 			})
 			delete(this.m, key)
@@ -132,8 +134,8 @@ func (this *BandwidthStatManager) Loop() error {
 }
 
 // Add 添加带宽数据
-func (this *BandwidthStatManager) Add(userId int64, serverId int64, bytes int64) {
-	if serverId <= 0 || bytes == 0 {
+func (this *BandwidthStatManager) Add(userId int64, serverId int64, peekBytes int64, totalBytes int64) {
+	if serverId <= 0 || (peekBytes == 0 && totalBytes == 0) {
 		return
 	}
 
@@ -146,8 +148,8 @@ func (this *BandwidthStatManager) Add(userId int64, serverId int64, bytes int64)
 	// 增加TCP Header尺寸，这里默认MTU为1500，且默认为IPv4
 	const mtu = 1500
 	const tcpHeaderSize = 20
-	if bytes > mtu {
-		bytes += bytes * tcpHeaderSize / mtu
+	if peekBytes > mtu {
+		peekBytes += peekBytes * tcpHeaderSize / mtu
 	}
 
 	this.locker.Lock()
@@ -156,22 +158,25 @@ func (this *BandwidthStatManager) Add(userId int64, serverId int64, bytes int64)
 		// 此刻如果发生用户ID（userId）的变化也忽略，等N分钟后有新记录后再换
 
 		if stat.CurrentTimestamp == timestamp {
-			stat.CurrentBytes += bytes
+			stat.CurrentBytes += peekBytes
 		} else {
-			stat.CurrentBytes = bytes
+			stat.CurrentBytes = peekBytes
 			stat.CurrentTimestamp = timestamp
 		}
 		if stat.CurrentBytes > stat.MaxBytes {
 			stat.MaxBytes = stat.CurrentBytes
 		}
+
+		stat.TotalBytes += totalBytes
 	} else {
 		this.m[key] = &BandwidthStat{
 			Day:              day,
 			TimeAt:           timeAt,
 			UserId:           userId,
 			ServerId:         serverId,
-			CurrentBytes:     bytes,
-			MaxBytes:         bytes,
+			CurrentBytes:     peekBytes,
+			MaxBytes:         peekBytes,
+			TotalBytes:       totalBytes,
 			CurrentTimestamp: timestamp,
 		}
 	}