diff --git a/go-gost/program.go b/go-gost/program.go index 696c8e6..eae3d1b 100644 --- a/go-gost/program.go +++ b/go-gost/program.go @@ -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 diff --git a/go-gost/x/api/config_reload.go b/go-gost/x/api/config_reload.go index f36a56e..dbfb003 100644 --- a/go-gost/x/api/config_reload.go +++ b/go-gost/x/api/config_reload.go @@ -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 diff --git a/go-gost/x/api/config_service.go b/go-gost/x/api/config_service.go index 36cd35a..bf5898b 100644 --- a/go-gost/x/api/config_service.go +++ b/go-gost/x/api/config_service.go @@ -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", }) diff --git a/go-gost/x/service/global_traffic_manager.go b/go-gost/x/service/global_traffic_manager.go index 2280f0c..a693151 100644 --- a/go-gost/x/service/global_traffic_manager.go +++ b/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 } - diff --git a/go-gost/x/service/global_traffic_manager_test.go b/go-gost/x/service/global_traffic_manager_test.go new file mode 100644 index 0000000..9bf39fc --- /dev/null +++ b/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") + } +} diff --git a/go-gost/x/socket/service.go b/go-gost/x/socket/service.go index f2d7d24..731ea19 100644 --- a/go-gost/x/socket/service.go +++ b/go-gost/x/socket/service.go @@ -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 }