diff --git a/common/trafficcontrol/manager.go b/common/trafficcontrol/manager.go index c18a59f5..d5b6a2b5 100644 --- a/common/trafficcontrol/manager.go +++ b/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 - uploadTotal atomic.Int64 - downloadTotal atomic.Int64 + outbound adapter.OutboundManager 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() } diff --git a/common/trafficcontrol/tracker.go b/common/trafficcontrol/tracker.go index 34820f30..260eb70f 100644 --- a/common/trafficcontrol/tracker.go +++ b/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) { - upload.Add(n) - m.uploadTotal.Add(n) - }}, []N.CountFunc{func(n int64) { - download.Add(n) - m.downloadTotal.Add(n) - }}), - metadata: m.newTrackerMetadata(metadata, matchedRule, matchOutbound, upload, download), - manager: m, + ExtendedConn: bufio.NewInt64CounterConn(conn, []*atomic.Int64{upload}, []*atomic.Int64{download}), + 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) { - upload.Add(n) - m.uploadTotal.Add(n) - }}, []N.CountFunc{func(n int64) { - download.Add(n) - m.downloadTotal.Add(n) - }}), - metadata: m.newTrackerMetadata(metadata, matchedRule, matchOutbound, upload, download), - manager: m, + PacketConn: bufio.NewInt64CounterPacketConn(conn, []*atomic.Int64{upload}, nil, []*atomic.Int64{download}, nil), + 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() {