mirror of
https://github.com/Sagit-chu/flvx.git
synced 2026-09-28 07:36:38 +08:00
fix: prune stale agent traffic counters
This commit is contained in:
@@ -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,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
|
||||
|
||||
@@ -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",
|
||||
})
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user