From 847c80a1fe2cd1c0e7ee71c41daab6fa5fa3d70a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Mon, 14 Apr 2025 23:40:54 +0800
Subject: [PATCH] Fix

---
 service/ssmapi/api.go     |   8 +--
 service/ssmapi/traffic.go | 101 ++++++++++++++++++--------------------
 2 files changed, 51 insertions(+), 58 deletions(-)

diff --git a/service/ssmapi/api.go b/service/ssmapi/api.go
index b9b753a4..6c742c13 100644
--- a/service/ssmapi/api.go
+++ b/service/ssmapi/api.go
@@ -159,15 +159,11 @@ func (s *APIServer) getStats(writer http.ResponseWriter, request *http.Request)
 	requireClear := chi.URLParam(request, "clear") == "true"
 
 	users := s.user.List()
-	s.traffic.ReadUsers(users)
+	s.traffic.ReadUsers(users, requireClear)
 	for i := range users {
 		users[i].Password = ""
 	}
-	uplinkBytes, downlinkBytes, uplinkPackets, downlinkPackets, tcpSessions, udpSessions := s.traffic.ReadGlobal()
-
-	if requireClear {
-		s.traffic.Clear()
-	}
+	uplinkBytes, downlinkBytes, uplinkPackets, downlinkPackets, tcpSessions, udpSessions := s.traffic.ReadGlobal(requireClear)
 
 	render.JSON(writer, request, render.M{
 		"uplinkBytes":     uplinkBytes,
diff --git a/service/ssmapi/traffic.go b/service/ssmapi/traffic.go
index 7f3f103e..11bd14fa 100644
--- a/service/ssmapi/traffic.go
+++ b/service/ssmapi/traffic.go
@@ -136,80 +136,77 @@ func (s *TrafficManager) TrackPacketConnection(conn N.PacketConn, metadata adapt
 func (s *TrafficManager) ReadUser(user *UserObject) {
 	s.userAccess.Lock()
 	defer s.userAccess.Unlock()
-	s.readUser(user)
+	s.readUser(user, false)
 }
 
-func (s *TrafficManager) readUser(user *UserObject) {
+func (s *TrafficManager) readUser(user *UserObject, swap bool) {
 	if counter, loaded := s.userUplink[user.UserName]; loaded {
-		user.UplinkBytes = counter.Load()
+		if swap {
+			user.UplinkBytes = counter.Swap(0)
+		} else {
+			user.UplinkBytes = counter.Load()
+		}
 	}
 	if counter, loaded := s.userDownlink[user.UserName]; loaded {
-		user.DownlinkBytes = counter.Load()
+		if swap {
+			user.DownlinkBytes = counter.Swap(0)
+		} else {
+			user.DownlinkBytes = counter.Load()
+		}
 	}
 	if counter, loaded := s.userUplinkPackets[user.UserName]; loaded {
-		user.UplinkPackets = counter.Load()
+		if swap {
+			user.UplinkPackets = counter.Swap(0)
+		} else {
+			user.UplinkPackets = counter.Load()
+		}
 	}
 	if counter, loaded := s.userDownlinkPackets[user.UserName]; loaded {
-		user.DownlinkPackets = counter.Load()
+		if swap {
+			user.DownlinkPackets = counter.Swap(0)
+		} else {
+			user.DownlinkPackets = counter.Load()
+		}
 	}
 	if counter, loaded := s.userTCPSessions[user.UserName]; loaded {
-		user.TCPSessions = counter.Load()
+		if swap {
+			user.TCPSessions = counter.Swap(0)
+		} else {
+			user.TCPSessions = counter.Load()
+		}
 	}
 	if counter, loaded := s.userUDPSessions[user.UserName]; loaded {
-		user.UDPSessions = counter.Load()
+		if swap {
+			user.UDPSessions = counter.Swap(0)
+		} else {
+			user.UDPSessions = counter.Load()
+		}
 	}
 }
 
-func (s *TrafficManager) ReadUsers(users []*UserObject) {
+func (s *TrafficManager) ReadUsers(users []*UserObject, swap bool) {
 	s.userAccess.Lock()
 	defer s.userAccess.Unlock()
 	for _, user := range users {
-		s.readUser(user)
+		s.readUser(user, swap)
 	}
 	return
 }
 
-func (s *TrafficManager) ReadGlobal() (
-	uplinkBytes int64,
-	downlinkBytes int64,
-	uplinkPackets int64,
-	downlinkPackets int64,
-	tcpSessions int64,
-	udpSessions int64,
-) {
-	return s.globalUplink.Load(),
-		s.globalDownlink.Load(),
-		s.globalUplinkPackets.Load(),
-		s.globalDownlinkPackets.Load(),
-		s.globalTCPSessions.Load(),
-		s.globalUDPSessions.Load()
-}
-
-func (s *TrafficManager) Clear() {
-	s.globalUplink.Store(0)
-	s.globalDownlink.Store(0)
-	s.globalUplinkPackets.Store(0)
-	s.globalDownlinkPackets.Store(0)
-	s.globalTCPSessions.Store(0)
-	s.globalUDPSessions.Store(0)
-	s.userAccess.Lock()
-	defer s.userAccess.Unlock()
-	for _, counter := range s.userUplink {
-		counter.Store(0)
-	}
-	for _, counter := range s.userDownlink {
-		counter.Store(0)
-	}
-	for _, counter := range s.userUplinkPackets {
-		counter.Store(0)
-	}
-	for _, counter := range s.userDownlinkPackets {
-		counter.Store(0)
-	}
-	for _, counter := range s.userTCPSessions {
-		counter.Store(0)
-	}
-	for _, counter := range s.userUDPSessions {
-		counter.Store(0)
+func (s *TrafficManager) ReadGlobal(swap bool) (uplinkBytes int64, downlinkBytes int64, uplinkPackets int64, downlinkPackets int64, tcpSessions int64, udpSessions int64) {
+	if swap {
+		return s.globalUplink.Swap(0),
+			s.globalDownlink.Swap(0),
+			s.globalUplinkPackets.Swap(0),
+			s.globalDownlinkPackets.Swap(0),
+			s.globalTCPSessions.Swap(0),
+			s.globalUDPSessions.Swap(0)
+	} else {
+		return s.globalUplink.Load(),
+			s.globalDownlink.Load(),
+			s.globalUplinkPackets.Load(),
+			s.globalDownlinkPackets.Load(),
+			s.globalTCPSessions.Load(),
+			s.globalUDPSessions.Load()
 	}
 }