mirror of
https://github.com/Sagit-chu/flvx.git
synced 2026-10-12 04:06:37 +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 {
|
if err := loader.Load(cfg); err != nil {
|
||||||
return err
|
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 {
|
if err := p.run(cfg); err != nil {
|
||||||
return err
|
return err
|
||||||
|
|||||||
@@ -8,6 +8,7 @@ import (
|
|||||||
"github.com/go-gost/x/config/loader"
|
"github.com/go-gost/x/config/loader"
|
||||||
"github.com/go-gost/x/config/parsing/parser"
|
"github.com/go-gost/x/config/parsing/parser"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
|
xservice "github.com/go-gost/x/service"
|
||||||
)
|
)
|
||||||
|
|
||||||
// swagger:parameters reloadConfigRequest
|
// swagger:parameters reloadConfigRequest
|
||||||
@@ -42,6 +43,13 @@ func reloadConfig(ctx *gin.Context) {
|
|||||||
writeError(ctx, NewError(http.StatusBadRequest, ErrCodeInvalid, err.Error()))
|
writeError(ctx, NewError(http.StatusBadRequest, ErrCodeInvalid, err.Error()))
|
||||||
return
|
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() {
|
for _, svc := range registry.ServiceRegistry().GetAll() {
|
||||||
svc := svc
|
svc := svc
|
||||||
|
|||||||
@@ -14,6 +14,7 @@ import (
|
|||||||
parser "github.com/go-gost/x/config/parsing/service"
|
parser "github.com/go-gost/x/config/parsing/service"
|
||||||
kill "github.com/go-gost/x/internal/util/port"
|
kill "github.com/go-gost/x/internal/util/port"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
|
xservice "github.com/go-gost/x/service"
|
||||||
)
|
)
|
||||||
|
|
||||||
// swagger:parameters createServiceRequest
|
// swagger:parameters createServiceRequest
|
||||||
@@ -409,6 +410,7 @@ func deleteService(ctx *gin.Context) {
|
|||||||
}
|
}
|
||||||
return nil
|
return nil
|
||||||
})
|
})
|
||||||
|
xservice.GetGlobalTrafficManager().RemoveServices(name)
|
||||||
|
|
||||||
ctx.JSON(http.StatusOK, Response{
|
ctx.JSON(http.StatusOK, Response{
|
||||||
Msg: "OK",
|
Msg: "OK",
|
||||||
@@ -484,6 +486,12 @@ func deleteServices(ctx *gin.Context) {
|
|||||||
return nil
|
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{
|
ctx.JSON(http.StatusOK, Response{
|
||||||
Msg: "OK",
|
Msg: "OK",
|
||||||
})
|
})
|
||||||
|
|||||||
@@ -5,15 +5,17 @@ import (
|
|||||||
"fmt"
|
"fmt"
|
||||||
"sync"
|
"sync"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"github.com/go-gost/x/registry"
|
||||||
)
|
)
|
||||||
|
|
||||||
// GlobalTrafficManager 全局流量管理器(所有服务共享)
|
// GlobalTrafficManager 全局流量管理器(所有服务共享)
|
||||||
type GlobalTrafficManager struct {
|
type GlobalTrafficManager struct {
|
||||||
mu sync.RWMutex
|
mu sync.RWMutex
|
||||||
serviceTraffic map[string]*ServiceTraffic // key: 服务名, value: 流量数据
|
serviceTraffic map[string]*ServiceTraffic // key: 服务名, value: 流量数据
|
||||||
ctx context.Context
|
ctx context.Context
|
||||||
cancel context.CancelFunc
|
cancel context.CancelFunc
|
||||||
reportTicker *time.Ticker
|
reportTicker *time.Ticker
|
||||||
}
|
}
|
||||||
|
|
||||||
// ServiceTraffic 单个服务的流量累积
|
// ServiceTraffic 单个服务的流量累积
|
||||||
@@ -50,6 +52,9 @@ func (m *GlobalTrafficManager) AddTraffic(serviceName string, upBytes, downBytes
|
|||||||
if upBytes == 0 && downBytes == 0 {
|
if upBytes == 0 && downBytes == 0 {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
if !registry.ServiceRegistry().IsRegistered(serviceName) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
m.mu.Lock()
|
m.mu.Lock()
|
||||||
defer m.mu.Unlock()
|
defer m.mu.Unlock()
|
||||||
@@ -70,6 +75,51 @@ func (m *GlobalTrafficManager) AddTraffic(serviceName string, upBytes, downBytes
|
|||||||
traffic.mu.Unlock()
|
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秒执行一次)
|
// startReporting 启动定时上报协程(每5秒执行一次)
|
||||||
func (m *GlobalTrafficManager) startReporting() {
|
func (m *GlobalTrafficManager) startReporting() {
|
||||||
|
|
||||||
@@ -105,6 +155,7 @@ func (m *GlobalTrafficManager) collectAndReport() {
|
|||||||
traffic.DownBytes = 0
|
traffic.DownBytes = 0
|
||||||
}
|
}
|
||||||
traffic.mu.Unlock()
|
traffic.mu.Unlock()
|
||||||
|
isStale := !registry.ServiceRegistry().IsRegistered(name)
|
||||||
|
|
||||||
if up > 0 || down > 0 {
|
if up > 0 || down > 0 {
|
||||||
reportItems = append(reportItems, TrafficReportItem{
|
reportItems = append(reportItems, TrafficReportItem{
|
||||||
@@ -113,6 +164,9 @@ func (m *GlobalTrafficManager) collectAndReport() {
|
|||||||
D: down,
|
D: down,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
if isStale {
|
||||||
|
delete(m.serviceTraffic, name)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
m.mu.Unlock()
|
m.mu.Unlock()
|
||||||
@@ -162,4 +216,3 @@ func (m *GlobalTrafficManager) GetServiceTraffic(serviceName string) (upBytes, d
|
|||||||
}
|
}
|
||||||
return
|
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"
|
parser "github.com/go-gost/x/config/parsing/service"
|
||||||
kill "github.com/go-gost/x/internal/util/port"
|
kill "github.com/go-gost/x/internal/util/port"
|
||||||
"github.com/go-gost/x/registry"
|
"github.com/go-gost/x/registry"
|
||||||
|
xservice "github.com/go-gost/x/service"
|
||||||
)
|
)
|
||||||
|
|
||||||
func createServices(req createServicesRequest) error {
|
func createServices(req createServicesRequest) error {
|
||||||
@@ -209,6 +210,7 @@ func deleteServices(req deleteServicesRequest) error {
|
|||||||
}
|
}
|
||||||
return nil
|
return nil
|
||||||
})
|
})
|
||||||
|
xservice.GetGlobalTrafficManager().RemoveServices(namesToRemove...)
|
||||||
|
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user