Improve traffic totals

common/trafficcontrol/manager.go

@@ -2,7 +2,6 @@ package trafficcontrol
 
 import (
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/sagernet/sing-box/adapter"
@@ -36,13 +35,13 @@ var (
 )
 
 type Manager struct {
-	outbound      adapter.OutboundManager
+	outbound adapter.OutboundManager
-	uploadTotal   atomic.Int64
-	downloadTotal atomic.Int64
 
 	connections             compatible.Map[uuid.UUID, Tracker]
 	closedConnectionsAccess sync.Mutex
 	closedConnections       list.List[TrackerMetadata]
+	closedUploadTotal       int64
+	closedDownloadTotal     int64
 
 	eventSubscriber *observable.Subscriber[ConnectionEvent]
 	eventObserver   *observable.Observer[ConnectionEvent]
@@ -98,16 +97,19 @@ func (m *Manager) join(tracker Tracker) {
 
 func (m *Manager) leave(tracker Tracker) {
 	metadata := tracker.Metadata()
+	closedAt := time.Now()
+	m.closedConnectionsAccess.Lock()
 	_, loaded := m.connections.LoadAndDelete(metadata.ID)
 	if !loaded {
+		m.closedConnectionsAccess.Unlock()
 		return
 	}
-	closedAt := time.Now()
 	metadata.ClosedAt = closedAt
 	metadataCopy := *metadata
-	m.closedConnectionsAccess.Lock()
 	if m.closedConnections.Len() >= closedConnectionsLimit {
-		m.closedConnections.PopFront()
+		evicted := m.closedConnections.PopFront()
+		m.closedUploadTotal += evicted.Upload.Load()
+		m.closedDownloadTotal += evicted.Download.Load()
 	}
 	m.closedConnections.PushBack(metadataCopy)
 	m.closedConnectionsAccess.Unlock()
@@ -120,7 +122,21 @@ func (m *Manager) leave(tracker Tracker) {
 }
 
 func (m *Manager) Total() (uplinkTotal int64, downlinkTotal int64) {
-	return m.uploadTotal.Load(), m.downloadTotal.Load()
+	m.closedConnectionsAccess.Lock()
+	defer m.closedConnectionsAccess.Unlock()
+	uplinkTotal = m.closedUploadTotal
+	downlinkTotal = m.closedDownloadTotal
+	for element := m.closedConnections.Front(); element != nil; element = element.Next() {
+		uplinkTotal += element.Value.Upload.Load()
+		downlinkTotal += element.Value.Download.Load()
+	}
+	m.connections.Range(func(_ uuid.UUID, tracker Tracker) bool {
+		metadata := tracker.Metadata()
+		uplinkTotal += metadata.Upload.Load()
+		downlinkTotal += metadata.Download.Load()
+		return true
+	})
+	return
 }
 
 func (m *Manager) ConnectionsLen() int {
@@ -168,5 +184,9 @@ func (m *Manager) CloseAllConnections() {
 func (m *Manager) Clear() {
 	m.closedConnectionsAccess.Lock()
 	defer m.closedConnectionsAccess.Unlock()
+	for element := m.closedConnections.Front(); element != nil; element = element.Next() {
+		m.closedUploadTotal += element.Value.Upload.Load()
+		m.closedDownloadTotal += element.Value.Download.Load()
+	}
 	m.closedConnections.Init()
 }

common/trafficcontrol/tracker.go

@@ -37,15 +37,9 @@ func (m *Manager) RoutedConnection(ctx context.Context, conn net.Conn, metadata
 	upload := new(atomic.Int64)
 	download := new(atomic.Int64)
 	tracker := &connTracker{
-		ExtendedConn: bufio.NewCounterConn(conn, []N.CountFunc{func(n int64) {
+		ExtendedConn: bufio.NewInt64CounterConn(conn, []*atomic.Int64{upload}, []*atomic.Int64{download}),
-			upload.Add(n)
+		metadata:     m.newTrackerMetadata(metadata, matchedRule, matchOutbound, upload, download),
-			m.uploadTotal.Add(n)
+		manager:      m,
-		}}, []N.CountFunc{func(n int64) {
-			download.Add(n)
-			m.downloadTotal.Add(n)
-		}}),
-		metadata: m.newTrackerMetadata(metadata, matchedRule, matchOutbound, upload, download),
-		manager:  m,
 	}
 	m.join(tracker)
 	return tracker
@@ -55,15 +49,9 @@ func (m *Manager) RoutedPacketConnection(ctx context.Context, conn N.PacketConn,
 	upload := new(atomic.Int64)
 	download := new(atomic.Int64)
 	tracker := &packetConnTracker{
-		PacketConn: bufio.NewCounterPacketConn(conn, []N.CountFunc{func(n int64) {
+		PacketConn: bufio.NewInt64CounterPacketConn(conn, []*atomic.Int64{upload}, nil, []*atomic.Int64{download}, nil),
-			upload.Add(n)
+		metadata:   m.newTrackerMetadata(metadata, matchedRule, matchOutbound, upload, download),
-			m.uploadTotal.Add(n)
+		manager:    m,
-		}}, []N.CountFunc{func(n int64) {
-			download.Add(n)
-			m.downloadTotal.Add(n)
-		}}),
-		metadata: m.newTrackerMetadata(metadata, matchedRule, matchOutbound, upload, download),
-		manager:  m,
 	}
 	m.join(tracker)
 	return tracker
@@ -165,12 +153,10 @@ func (t *flowTracker) AttachFlow(handle tun.FlowHandle) {
 
 func (t *flowTracker) CountForward(n int) {
 	t.metadata.Upload.Add(int64(n))
-	t.manager.uploadTotal.Add(int64(n))
 }
 
 func (t *flowTracker) CountReverse(n int) {
 	t.metadata.Download.Add(int64(n))
-	t.manager.downloadTotal.Add(int64(n))
 }
 
 func (t *flowTracker) FlowEstablished() {