fix: prune stale agent traffic counters

go-gost/x/service/global_traffic_manager.go

@@ -5,15 +5,17 @@ import (
 	"fmt"
 	"sync"
 	"time"
+
+	"github.com/go-gost/x/registry"
 )
 
 // GlobalTrafficManager 全局流量管理器（所有服务共享）
 type GlobalTrafficManager struct {
-	mu            sync.RWMutex
+	mu             sync.RWMutex
 	serviceTraffic map[string]*ServiceTraffic // key: 服务名, value: 流量数据
-	ctx           context.Context
-	cancel        context.CancelFunc
-	reportTicker  *time.Ticker
+	ctx            context.Context
+	cancel         context.CancelFunc
+	reportTicker   *time.Ticker
 }
 
 // ServiceTraffic 单个服务的流量累积
@@ -50,6 +52,9 @@ func (m *GlobalTrafficManager) AddTraffic(serviceName string, upBytes, downBytes
 	if upBytes == 0 && downBytes == 0 {
 		return
 	}
+	if !registry.ServiceRegistry().IsRegistered(serviceName) {
+		return
+	}
 
 	m.mu.Lock()
 	defer m.mu.Unlock()
@@ -70,6 +75,51 @@ func (m *GlobalTrafficManager) AddTraffic(serviceName string, upBytes, downBytes
 	traffic.mu.Unlock()
 }
 
+// RemoveServices drops cached traffic counters for services that no longer exist.
+func (m *GlobalTrafficManager) RemoveServices(serviceNames ...string) {
+	if m == nil || len(serviceNames) == 0 {
+		return
+	}
+
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	for _, name := range serviceNames {
+		if m.isTrafficEmptyLocked(name) {
+			delete(m.serviceTraffic, name)
+		}
+	}
+}
+
+// RetainServices removes traffic counters for every service not in activeNames.
+func (m *GlobalTrafficManager) RetainServices(activeNames map[string]struct{}) {
+	if m == nil {
+		return
+	}
+
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	for name := range m.serviceTraffic {
+		if _, ok := activeNames[name]; !ok {
+			if m.isTrafficEmptyLocked(name) {
+				delete(m.serviceTraffic, name)
+			}
+		}
+	}
+}
+
+func (m *GlobalTrafficManager) isTrafficEmptyLocked(name string) bool {
+	traffic, ok := m.serviceTraffic[name]
+	if !ok {
+		return true
+	}
+
+	traffic.mu.Lock()
+	defer traffic.mu.Unlock()
+	return traffic.UpBytes == 0 && traffic.DownBytes == 0
+}
+
 // startReporting 启动定时上报协程（每5秒执行一次）
 func (m *GlobalTrafficManager) startReporting() {
 
@@ -105,6 +155,7 @@ func (m *GlobalTrafficManager) collectAndReport() {
 			traffic.DownBytes = 0
 		}
 		traffic.mu.Unlock()
+		isStale := !registry.ServiceRegistry().IsRegistered(name)
 
 		if up > 0 || down > 0 {
 			reportItems = append(reportItems, TrafficReportItem{
@@ -113,6 +164,9 @@ func (m *GlobalTrafficManager) collectAndReport() {
 				D: down,
 			})
 		}
+		if isStale {
+			delete(m.serviceTraffic, name)
+		}
 	}
 
 	m.mu.Unlock()
@@ -162,4 +216,3 @@ func (m *GlobalTrafficManager) GetServiceTraffic(serviceName string) (upBytes, d
 	}
 	return
 }
-

go-gost/x/service/global_traffic_manager_test.go

@@ -0,0 +1,95 @@
+package service
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"strings"
+	"testing"
+	"time"
+)
+
+func TestGlobalTrafficManagerRemoveServicesDropsCachedEntries(t *testing.T) {
+	m := &GlobalTrafficManager{serviceTraffic: map[string]*ServiceTraffic{
+		"svc-a": {ServiceName: "svc-a"},
+		"svc-b": {ServiceName: "svc-b"},
+	}}
+
+	m.RemoveServices("svc-a")
+
+	if _, ok := m.serviceTraffic["svc-a"]; ok {
+		t.Fatalf("expected svc-a traffic entry to be removed")
+	}
+	if _, ok := m.serviceTraffic["svc-b"]; !ok {
+		t.Fatalf("expected svc-b traffic entry to remain")
+	}
+}
+
+func TestGlobalTrafficManagerRetainServicesDropsStaleEntries(t *testing.T) {
+	m := &GlobalTrafficManager{serviceTraffic: map[string]*ServiceTraffic{
+		"svc-a": {ServiceName: "svc-a"},
+		"svc-b": {ServiceName: "svc-b"},
+	}}
+
+	m.RetainServices(map[string]struct{}{"svc-b": {}})
+
+	if _, ok := m.serviceTraffic["svc-a"]; ok {
+		t.Fatalf("expected stale svc-a traffic entry to be removed")
+	}
+	if _, ok := m.serviceTraffic["svc-b"]; !ok {
+		t.Fatalf("expected active svc-b traffic entry to remain")
+	}
+}
+
+func TestGlobalTrafficManagerAddTrafficIgnoresUnregisteredService(t *testing.T) {
+	m := &GlobalTrafficManager{serviceTraffic: make(map[string]*ServiceTraffic)}
+
+	m.AddTraffic("deleted-service", 10, 20)
+
+	if _, ok := m.serviceTraffic["deleted-service"]; ok {
+		t.Fatalf("expected unregistered service traffic to be ignored")
+	}
+}
+
+func TestGlobalTrafficManagerCollectAndReportDropsStaleEntriesAfterReporting(t *testing.T) {
+	origReportDo := reportDo
+	origReportURL := httpReportURL
+	origAESCrypto := httpAESCrypto
+	defer func() {
+		reportDo = origReportDo
+		httpReportURL = origReportURL
+		httpAESCrypto = origAESCrypto
+	}()
+
+	httpReportURL = "http://panel.example.com/flow/upload?secret=abc"
+	httpAESCrypto = nil
+
+	var requestBody string
+	reportDo = func(_ context.Context, req *http.Request, _ time.Duration) (*http.Response, error) {
+		body, err := io.ReadAll(req.Body)
+		if err != nil {
+			t.Fatalf("read request body: %v", err)
+		}
+		requestBody = string(body)
+		return &http.Response{
+			StatusCode: http.StatusOK,
+			Body:       io.NopCloser(strings.NewReader("ok")),
+		}, nil
+	}
+
+	m := &GlobalTrafficManager{
+		serviceTraffic: map[string]*ServiceTraffic{
+			"stale-service": {ServiceName: "stale-service", UpBytes: 10, DownBytes: 20},
+		},
+		ctx: context.Background(),
+	}
+
+	m.collectAndReport()
+
+	if !strings.Contains(requestBody, "stale-service") {
+		t.Fatalf("expected pending stale traffic to be reported first, body=%s", requestBody)
+	}
+	if _, ok := m.serviceTraffic["stale-service"]; ok {
+		t.Fatalf("expected stale traffic entry to be removed after report collection")
+	}
+}