From 29a99985c8e1e2cf0a6b3d1e5d2ba0a46baf67e3 Mon Sep 17 00:00:00 2001
From: wyx2685 <wyx2685@gmail.com>
Date: Sun, 28 Jul 2024 13:54:47 +0900
Subject: [PATCH] =?UTF-8?q?Fix:=E4=BC=98=E5=8C=96=E6=B5=81=E9=87=8F?=
 =?UTF-8?q?=E7=BB=9F=E8=AE=A1=E5=B9=B6=E5=8F=91=E6=83=85=E5=86=B5=E5=8F=AF?=
 =?UTF-8?q?=E8=83=BD=E5=AD=98=E5=9C=A8=E7=9A=84=E5=86=B2=E7=AA=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 common/counter/traffic.go | 65 +++++++++++++++------------------------
 1 file changed, 24 insertions(+), 41 deletions(-)

diff --git a/common/counter/traffic.go b/common/counter/traffic.go
index 2af0ae4..5299415 100644
--- a/common/counter/traffic.go
+++ b/common/counter/traffic.go
@@ -6,8 +6,7 @@ import (
 )
 
 type TrafficCounter struct {
-	counters map[string]*TrafficStorage
-	lock     sync.RWMutex
+	counters sync.Map
 }
 
 type TrafficStorage struct {
@@ -16,60 +15,52 @@ type TrafficStorage struct {
 }
 
 func NewTrafficCounter() *TrafficCounter {
-	return &TrafficCounter{
-		counters: map[string]*TrafficStorage{},
-	}
+	return &TrafficCounter{}
 }
 
 func (c *TrafficCounter) GetCounter(id string) *TrafficStorage {
-	c.lock.RLock()
-	cts, ok := c.counters[id]
-	c.lock.RUnlock()
-	if !ok {
-		cts = &TrafficStorage{}
-		c.counters[id] = cts
+	if cts, ok := c.counters.Load(id); ok {
+		return cts.(*TrafficStorage)
 	}
-	return cts
+	newStorage := &TrafficStorage{}
+	if cts, loaded := c.counters.LoadOrStore(id, newStorage); loaded {
+		return cts.(*TrafficStorage)
+	}
+	return newStorage
 }
 
 func (c *TrafficCounter) GetUpCount(id string) int64 {
-	c.lock.RLock()
-	cts, ok := c.counters[id]
-	c.lock.RUnlock()
-	if ok {
-		return cts.UpCounter.Load()
+	if cts, ok := c.counters.Load(id); ok {
+		return cts.(*TrafficStorage).UpCounter.Load()
 	}
 	return 0
 }
 
 func (c *TrafficCounter) GetDownCount(id string) int64 {
-	c.lock.RLock()
-	cts, ok := c.counters[id]
-	c.lock.RUnlock()
-	if ok {
-		return cts.DownCounter.Load()
+	if cts, ok := c.counters.Load(id); ok {
+		return cts.(*TrafficStorage).DownCounter.Load()
 	}
 	return 0
 }
 
 func (c *TrafficCounter) Len() int {
-	c.lock.RLock()
-	defer c.lock.RUnlock()
-	return len(c.counters)
+	length := 0
+	c.counters.Range(func(_, _ interface{}) bool {
+		length++
+		return true
+	})
+	return length
 }
 
 func (c *TrafficCounter) Reset(id string) {
-	c.lock.RLock()
-	cts := c.GetCounter(id)
-	c.lock.RUnlock()
-	cts.UpCounter.Store(0)
-	cts.DownCounter.Store(0)
+	if cts, ok := c.counters.Load(id); ok {
+		cts.(*TrafficStorage).UpCounter.Store(0)
+		cts.(*TrafficStorage).DownCounter.Store(0)
+	}
 }
 
 func (c *TrafficCounter) Delete(id string) {
-	c.lock.Lock()
-	delete(c.counters, id)
-	c.lock.Unlock()
+	c.counters.Delete(id)
 }
 
 func (c *TrafficCounter) Rx(id string, n int) {
@@ -81,11 +72,3 @@ func (c *TrafficCounter) Tx(id string, n int) {
 	cts := c.GetCounter(id)
 	cts.UpCounter.Add(int64(n))
 }
-
-func (c *TrafficCounter) IncConn(auth string) {
-	return
-}
-
-func (c *TrafficCounter) DecConn(auth string) {
-	return
-}