From 6f875001f25f204362239070fed5be4c55cf7ce3 Mon Sep 17 00:00:00 2001
From: bjdgyc <bjdgyc@163.com>
Date: Sun, 9 Oct 2022 15:16:22 +0800
Subject: [PATCH] =?UTF-8?q?=E6=9B=BF=E6=8D=A2=20go.uber.org/atomic?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 server/sessdata/online.go       |  9 ++++-----
 server/sessdata/session.go      | 29 +++++++++++++++--------------
 server/sessdata/session_test.go |  4 ++--
 server/sessdata/statsinfo.go    |  5 ++---
 4 files changed, 23 insertions(+), 24 deletions(-)

diff --git a/server/sessdata/online.go b/server/sessdata/online.go
index ca97aa7..e8dff7a 100644
--- a/server/sessdata/online.go
+++ b/server/sessdata/online.go
@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"net"
 	"sort"
-	"sync/atomic"
 	"time"
 
 	"github.com/bjdgyc/anylink/pkg/utils"
@@ -57,10 +56,10 @@ func OnlineSess() []Online {
 				TunName:          v.CSess.IfName,
 				Mtu:              v.CSess.Mtu,
 				Client:           v.CSess.Client,
-				BandwidthUp:      utils.HumanByte(atomic.LoadUint32(&v.CSess.BandwidthUpPeriod)) + "/s",
-				BandwidthDown:    utils.HumanByte(atomic.LoadUint32(&v.CSess.BandwidthDownPeriod)) + "/s",
-				BandwidthUpAll:   utils.HumanByte(atomic.LoadUint64(&v.CSess.BandwidthUpAll)),
-				BandwidthDownAll: utils.HumanByte(atomic.LoadUint64(&v.CSess.BandwidthDownAll)),
+				BandwidthUp:      utils.HumanByte(v.CSess.BandwidthUpPeriod.Load()) + "/s",
+				BandwidthDown:    utils.HumanByte(v.CSess.BandwidthDownPeriod.Load()) + "/s",
+				BandwidthUpAll:   utils.HumanByte(v.CSess.BandwidthUpAll.Load()),
+				BandwidthDownAll: utils.HumanByte(v.CSess.BandwidthDownAll.Load()),
 				LastLogin:        v.LastLogin,
 			}
 			datas = append(datas, val)
diff --git a/server/sessdata/session.go b/server/sessdata/session.go
index e773782..70f3e7e 100644
--- a/server/sessdata/session.go
+++ b/server/sessdata/session.go
@@ -14,6 +14,7 @@ import (
 	"github.com/bjdgyc/anylink/dbdata"
 	"github.com/bjdgyc/anylink/pkg/utils"
 	"github.com/ivpusic/grpool"
+	atomic2 "go.uber.org/atomic"
 )
 
 var (
@@ -39,12 +40,12 @@ type ConnSession struct {
 	CstpDpd             int
 	Group               *dbdata.Group
 	Limit               *LimitRater
-	BandwidthUp         uint32 // 使用上行带宽 Byte
-	BandwidthDown       uint32 // 使用下行带宽 Byte
-	BandwidthUpPeriod   uint32 // 前一周期的总量
-	BandwidthDownPeriod uint32
-	BandwidthUpAll      uint64 // 使用上行带宽总量
-	BandwidthDownAll    uint64 // 使用下行带宽总量
+	BandwidthUp         atomic2.Uint32 // 使用上行带宽 Byte
+	BandwidthDown       atomic2.Uint32 // 使用下行带宽 Byte
+	BandwidthUpPeriod   atomic2.Uint32 // 前一周期的总量
+	BandwidthDownPeriod atomic2.Uint32
+	BandwidthUpAll      atomic2.Uint64 // 使用上行带宽总量
+	BandwidthDownAll    atomic2.Uint64 // 使用下行带宽总量
 	closeOnce           sync.Once
 	CloseChan           chan struct{}
 	PayloadIn           chan *Payload
@@ -286,14 +287,14 @@ func (cs *ConnSession) ratePeriod() {
 		}
 
 		// 实时流量清零
-		rtUp := atomic.SwapUint32(&cs.BandwidthUp, 0)
-		rtDown := atomic.SwapUint32(&cs.BandwidthDown, 0)
+		rtUp := cs.BandwidthUp.Swap(0)
+		rtDown := cs.BandwidthDown.Swap(0)
 		// 设置上一周期每秒的流量
-		atomic.SwapUint32(&cs.BandwidthUpPeriod, rtUp/BandwidthPeriodSec)
-		atomic.SwapUint32(&cs.BandwidthDownPeriod, rtDown/BandwidthPeriodSec)
+		cs.BandwidthUpPeriod.Swap(rtUp / BandwidthPeriodSec)
+		cs.BandwidthDownPeriod.Swap(rtDown / BandwidthPeriodSec)
 		// 累加所有流量
-		atomic.AddUint64(&cs.BandwidthUpAll, uint64(rtUp))
-		atomic.AddUint64(&cs.BandwidthDownAll, uint64(rtDown))
+		cs.BandwidthUpAll.Add(uint64(rtUp))
+		cs.BandwidthDownAll.Add(uint64(rtDown))
 	}
 }
 
@@ -323,11 +324,11 @@ func (cs *ConnSession) SetIfName(name string) {
 
 func (cs *ConnSession) RateLimit(byt int, isUp bool) error {
 	if isUp {
-		atomic.AddUint32(&cs.BandwidthUp, uint32(byt))
+		cs.BandwidthUp.Add(uint32(byt))
 		return nil
 	}
 	// 只对下行速率限制
-	atomic.AddUint32(&cs.BandwidthDown, uint32(byt))
+	cs.BandwidthDown.Add(uint32(byt))
 	if cs.Limit == nil {
 		return nil
 	}
diff --git a/server/sessdata/session_test.go b/server/sessdata/session_test.go
index 26b7f3e..c9219b2 100644
--- a/server/sessdata/session_test.go
+++ b/server/sessdata/session_test.go
@@ -30,9 +30,9 @@ func TestConnSession(t *testing.T) {
 
 	err := cSess.RateLimit(100, true)
 	ast.Nil(err)
-	ast.Equal(cSess.BandwidthUp, uint32(100))
+	ast.Equal(cSess.BandwidthUp.Load(), uint32(100))
 	err = cSess.RateLimit(200, false)
 	ast.Nil(err)
-	ast.Equal(cSess.BandwidthDown, uint32(200))
+	ast.Equal(cSess.BandwidthDown.Load(), uint32(200))
 	cSess.Close()
 }
diff --git a/server/sessdata/statsinfo.go b/server/sessdata/statsinfo.go
index e07a0f2..79c3b9f 100644
--- a/server/sessdata/statsinfo.go
+++ b/server/sessdata/statsinfo.go
@@ -2,7 +2,6 @@ package sessdata
 
 import (
 	"encoding/json"
-	"sync/atomic"
 	"time"
 
 	"github.com/bjdgyc/anylink/dbdata"
@@ -34,8 +33,8 @@ func saveStatsInfo() {
 					onlineNum += 1
 					numGroups[v.CSess.Group.Id] += 1
 					// 网络吞吐
-					userUp := atomic.LoadUint32(&v.CSess.BandwidthUpPeriod)
-					userDown := atomic.LoadUint32(&v.CSess.BandwidthDownPeriod)
+					userUp := v.CSess.BandwidthUpPeriod.Load()
+					userDown := v.CSess.BandwidthDownPeriod.Load()
 					if userUp > 0 {
 						upGroups[v.CSess.Group.Id] += userUp
 					}