fix: prune stale agent traffic counters

This commit is contained in:
sagitchu
2026-04-30 09:54:02 +08:00
parent f05e9480ee
commit b4192c9e94
6 changed files with 178 additions and 5 deletions
+7
View File
@@ -234,6 +234,13 @@ func (p *program) reloadConfig() error {
if err := loader.Load(cfg); err != nil {
return err
}
activeServices := make(map[string]struct{}, len(cfg.Services))
for _, svc := range cfg.Services {
if svc != nil {
activeServices[svc.Name] = struct{}{}
}
}
xservice.GetGlobalTrafficManager().RetainServices(activeServices)
if err := p.run(cfg); err != nil {
return err
+8
View File
@@ -8,6 +8,7 @@ import (
"github.com/go-gost/x/config/loader"
"github.com/go-gost/x/config/parsing/parser"
"github.com/go-gost/x/registry"
xservice "github.com/go-gost/x/service"
)
// swagger:parameters reloadConfigRequest
@@ -42,6 +43,13 @@ func reloadConfig(ctx *gin.Context) {
writeError(ctx, NewError(http.StatusBadRequest, ErrCodeInvalid, err.Error()))
return
}
activeServices := make(map[string]struct{}, len(cfg.Services))
for _, svc := range cfg.Services {
if svc != nil {
activeServices[svc.Name] = struct{}{}
}
}
xservice.GetGlobalTrafficManager().RetainServices(activeServices)
for _, svc := range registry.ServiceRegistry().GetAll() {
svc := svc
+8
View File
@@ -14,6 +14,7 @@ import (
parser "github.com/go-gost/x/config/parsing/service"
kill "github.com/go-gost/x/internal/util/port"
"github.com/go-gost/x/registry"
xservice "github.com/go-gost/x/service"
)
// swagger:parameters createServiceRequest
@@ -409,6 +410,7 @@ func deleteService(ctx *gin.Context) {
}
return nil
})
xservice.GetGlobalTrafficManager().RemoveServices(name)
ctx.JSON(http.StatusOK, Response{
Msg: "OK",
@@ -484,6 +486,12 @@ func deleteServices(ctx *gin.Context) {
return nil
})
names := make([]string, 0, len(servicesToDelete))
for _, std := range servicesToDelete {
names = append(names, std.name)
}
xservice.GetGlobalTrafficManager().RemoveServices(names...)
ctx.JSON(http.StatusOK, Response{
Msg: "OK",
})
+54 -1
View File
@@ -5,6 +5,8 @@ import (
"fmt"
"sync"
"time"
"github.com/go-gost/x/registry"
)
// GlobalTrafficManager 全局流量管理器(所有服务共享)
@@ -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
}
@@ -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")
}
}
+2
View File
@@ -11,6 +11,7 @@ import (
parser "github.com/go-gost/x/config/parsing/service"
kill "github.com/go-gost/x/internal/util/port"
"github.com/go-gost/x/registry"
xservice "github.com/go-gost/x/service"
)
func createServices(req createServicesRequest) error {
@@ -209,6 +210,7 @@ func deleteServices(req deleteServicesRequest) error {
}
return nil
})
xservice.GetGlobalTrafficManager().RemoveServices(namesToRemove...)
return nil
}