mirror of
https://github.com/Sagit-chu/flvx.git
synced 2026-09-28 07:36:38 +08:00
refactor: batch flow upload processing (#474)
* test: cover flow upload batch semantics * refactor: batch flow upload persistence * refactor: batch flow upload processing * test: harden flow upload batch regression coverage
This commit is contained in:
@@ -0,0 +1,183 @@
|
|||||||
|
package handler
|
||||||
|
|
||||||
|
import (
|
||||||
|
"log"
|
||||||
|
"sort"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"go-backend/internal/store/model"
|
||||||
|
"go-backend/internal/store/repo"
|
||||||
|
)
|
||||||
|
|
||||||
|
type flowPolicyTarget struct {
|
||||||
|
UserID int64
|
||||||
|
UserTunnelID int64
|
||||||
|
}
|
||||||
|
|
||||||
|
type flowUploadBatch struct {
|
||||||
|
flowDeltas []repo.FlowUploadCounterDelta
|
||||||
|
quotaUsage map[int64]int64
|
||||||
|
policyTargets []flowPolicyTarget
|
||||||
|
forwardTraffic map[int64]tunnelTrafficDelta
|
||||||
|
orphanServices map[string]struct{}
|
||||||
|
peerShareForwardItems map[string]flowItem
|
||||||
|
peerShareRuntimeItems map[int64]flowItem
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *Handler) buildFlowUploadBatch(items []flowItem, metas map[int64]repo.FlowUploadForwardMeta) flowUploadBatch {
|
||||||
|
batch := flowUploadBatch{
|
||||||
|
quotaUsage: make(map[int64]int64),
|
||||||
|
forwardTraffic: make(map[int64]tunnelTrafficDelta),
|
||||||
|
orphanServices: make(map[string]struct{}),
|
||||||
|
peerShareForwardItems: make(map[string]flowItem),
|
||||||
|
peerShareRuntimeItems: make(map[int64]flowItem),
|
||||||
|
}
|
||||||
|
policySeen := map[flowPolicyTarget]struct{}{}
|
||||||
|
flowSeen := map[int64]int{}
|
||||||
|
|
||||||
|
for _, item := range items {
|
||||||
|
serviceName := strings.TrimSpace(item.N)
|
||||||
|
if serviceName == "" || serviceName == "web_api" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if runtimeID, ok := parsePeerShareRuntimeServiceID(serviceName); ok {
|
||||||
|
merged := batch.peerShareRuntimeItems[runtimeID]
|
||||||
|
merged.N = serviceName
|
||||||
|
merged.U += item.U
|
||||||
|
merged.D += item.D
|
||||||
|
batch.peerShareRuntimeItems[runtimeID] = merged
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
forwardID, userID, userTunnelID, ok := parseFlowServiceIDs(serviceName)
|
||||||
|
if !ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
normalized := normalizeForwardRuntimeServiceName(serviceName)
|
||||||
|
merged := batch.peerShareForwardItems[normalized]
|
||||||
|
merged.N = normalized
|
||||||
|
merged.U += item.U
|
||||||
|
merged.D += item.D
|
||||||
|
batch.peerShareForwardItems[normalized] = merged
|
||||||
|
|
||||||
|
meta, exists := metas[forwardID]
|
||||||
|
if !exists {
|
||||||
|
batch.orphanServices[serviceName] = struct{}{}
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
raw := batch.forwardTraffic[forwardID]
|
||||||
|
raw.bytesIn += item.D
|
||||||
|
raw.bytesOut += item.U
|
||||||
|
batch.forwardTraffic[forwardID] = raw
|
||||||
|
|
||||||
|
scaledIn := int64(float64(item.D)*meta.TrafficRatio) * meta.TunnelFlow
|
||||||
|
scaledOut := int64(float64(item.U)*meta.TrafficRatio) * meta.TunnelFlow
|
||||||
|
if idx, ok := flowSeen[forwardID]; ok {
|
||||||
|
batch.flowDeltas[idx].InFlow += scaledIn
|
||||||
|
batch.flowDeltas[idx].OutFlow += scaledOut
|
||||||
|
} else {
|
||||||
|
flowSeen[forwardID] = len(batch.flowDeltas)
|
||||||
|
batch.flowDeltas = append(batch.flowDeltas, repo.FlowUploadCounterDelta{
|
||||||
|
ForwardID: forwardID,
|
||||||
|
UserID: userID,
|
||||||
|
UserTunnelID: userTunnelID,
|
||||||
|
InFlow: scaledIn,
|
||||||
|
OutFlow: scaledOut,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
batch.quotaUsage[userID] += scaledIn + scaledOut
|
||||||
|
|
||||||
|
target := flowPolicyTarget{UserID: userID, UserTunnelID: userTunnelID}
|
||||||
|
if _, seen := policySeen[target]; !seen {
|
||||||
|
policySeen[target] = struct{}{}
|
||||||
|
batch.policyTargets = append(batch.policyTargets, target)
|
||||||
|
}
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
|
sort.Slice(batch.policyTargets, func(i, j int) bool {
|
||||||
|
if batch.policyTargets[i].UserID == batch.policyTargets[j].UserID {
|
||||||
|
return batch.policyTargets[i].UserTunnelID < batch.policyTargets[j].UserTunnelID
|
||||||
|
}
|
||||||
|
return batch.policyTargets[i].UserID < batch.policyTargets[j].UserID
|
||||||
|
})
|
||||||
|
|
||||||
|
return batch
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *Handler) applyFlowUploadBatch(nodeID int64, batch flowUploadBatch, now time.Time) {
|
||||||
|
if h == nil || h.repo == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
h.applyFlowDeltasWithFallback(nodeID, batch.flowDeltas)
|
||||||
|
for userID, quota := range h.applyQuotaUsageWithFallback(nodeID, batch.quotaUsage, now) {
|
||||||
|
h.enforceUserQuotaIfNeeded(userID, quota)
|
||||||
|
}
|
||||||
|
for _, target := range batch.policyTargets {
|
||||||
|
if target.UserID <= 0 || target.UserTunnelID <= 0 {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
h.enforceFlowPolicies(target.UserID, target.UserTunnelID)
|
||||||
|
}
|
||||||
|
for serviceName := range batch.orphanServices {
|
||||||
|
h.sendDeleteOrphanedForwardService(nodeID, serviceName)
|
||||||
|
}
|
||||||
|
for serviceName, item := range batch.peerShareForwardItems {
|
||||||
|
forwardID, _, _, ok := parseFlowServiceIDs(serviceName)
|
||||||
|
if ok {
|
||||||
|
h.processPeerShareFlowFromForward(forwardID, nodeID, serviceName, item)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for runtimeID, item := range batch.peerShareRuntimeItems {
|
||||||
|
h.processPeerShareFlow(runtimeID, item)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *Handler) applyFlowDeltasWithFallback(nodeID int64, deltas []repo.FlowUploadCounterDelta) {
|
||||||
|
if h == nil || h.repo == nil || len(deltas) == 0 {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if err := h.repo.ApplyFlowUploadDeltasBatch(deltas); err == nil {
|
||||||
|
return
|
||||||
|
} else {
|
||||||
|
log.Printf("flow upload write failed op=flow.batch_apply node_id=%d err=%v", nodeID, err)
|
||||||
|
}
|
||||||
|
for _, delta := range deltas {
|
||||||
|
if err := h.repo.AddFlow(delta.ForwardID, delta.UserID, delta.UserTunnelID, delta.InFlow, delta.OutFlow); err != nil {
|
||||||
|
log.Printf("flow upload write failed op=flow.single_apply node_id=%d forward_id=%d user_id=%d user_tunnel_id=%d err=%v", nodeID, delta.ForwardID, delta.UserID, delta.UserTunnelID, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *Handler) applyQuotaUsageWithFallback(nodeID int64, usages map[int64]int64, now time.Time) map[int64]*model.UserQuotaView {
|
||||||
|
if h == nil || h.repo == nil || len(usages) == 0 {
|
||||||
|
return map[int64]*model.UserQuotaView{}
|
||||||
|
}
|
||||||
|
quotaViews, err := h.repo.AddUserQuotaUsageBatch(usages, now)
|
||||||
|
if err == nil {
|
||||||
|
return quotaViews
|
||||||
|
}
|
||||||
|
log.Printf("flow upload write failed op=quota.batch_apply node_id=%d err=%v", nodeID, err)
|
||||||
|
|
||||||
|
userIDs := make([]int64, 0, len(usages))
|
||||||
|
for userID := range usages {
|
||||||
|
if userID > 0 {
|
||||||
|
userIDs = append(userIDs, userID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
sort.Slice(userIDs, func(i, j int) bool { return userIDs[i] < userIDs[j] })
|
||||||
|
|
||||||
|
quotaViews = make(map[int64]*model.UserQuotaView, len(userIDs))
|
||||||
|
for _, userID := range userIDs {
|
||||||
|
quota, singleErr := h.repo.AddUserQuotaUsage(userID, usages[userID], now)
|
||||||
|
if singleErr != nil {
|
||||||
|
log.Printf("flow upload write failed op=quota.single_apply node_id=%d user_id=%d err=%v", nodeID, userID, singleErr)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if quota != nil {
|
||||||
|
quotaViews[userID] = quota
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return quotaViews
|
||||||
|
}
|
||||||
@@ -0,0 +1,254 @@
|
|||||||
|
package handler
|
||||||
|
|
||||||
|
import (
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"go-backend/internal/store/model"
|
||||||
|
"go-backend/internal/store/repo"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestBuildFlowUploadBatchAggregatesForwardQuotaPeerShareAndCleanupTargets(t *testing.T) {
|
||||||
|
h := &Handler{}
|
||||||
|
metas := map[int64]repo.FlowUploadForwardMeta{
|
||||||
|
20: {
|
||||||
|
ForwardID: 20,
|
||||||
|
TunnelID: 1,
|
||||||
|
TrafficRatio: 2,
|
||||||
|
TunnelFlow: 3,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
batch := h.buildFlowUploadBatch([]flowItem{
|
||||||
|
{N: "20_2_10", U: 70, D: 50},
|
||||||
|
{N: "20_2_10_tcp", U: 40, D: 30},
|
||||||
|
{N: "99_2_10", U: 12, D: 8},
|
||||||
|
{N: "fed_svc_17", U: 9, D: 1},
|
||||||
|
}, metas)
|
||||||
|
|
||||||
|
if len(batch.flowDeltas) != 1 {
|
||||||
|
t.Fatalf("expected 1 flow delta, got %d", len(batch.flowDeltas))
|
||||||
|
}
|
||||||
|
delta := batch.flowDeltas[0]
|
||||||
|
if delta.ForwardID != 20 || delta.UserID != 2 || delta.UserTunnelID != 10 {
|
||||||
|
t.Fatalf("unexpected flow delta identity: %#v", delta)
|
||||||
|
}
|
||||||
|
if delta.InFlow != 480 || delta.OutFlow != 660 {
|
||||||
|
t.Fatalf("expected scaled flow in=480 out=660, got in=%d out=%d", delta.InFlow, delta.OutFlow)
|
||||||
|
}
|
||||||
|
if batch.quotaUsage[2] != 1140 {
|
||||||
|
t.Fatalf("expected quota usage 1140, got %d", batch.quotaUsage[2])
|
||||||
|
}
|
||||||
|
if len(batch.policyTargets) != 1 {
|
||||||
|
t.Fatalf("expected 1 policy target, got %d", len(batch.policyTargets))
|
||||||
|
}
|
||||||
|
if batch.policyTargets[0].UserID != 2 || batch.policyTargets[0].UserTunnelID != 10 {
|
||||||
|
t.Fatalf("unexpected policy target: %#v", batch.policyTargets[0])
|
||||||
|
}
|
||||||
|
traffic := batch.forwardTraffic[20]
|
||||||
|
if traffic.bytesIn != 80 || traffic.bytesOut != 110 {
|
||||||
|
t.Fatalf("expected raw traffic in=80 out=110, got in=%d out=%d", traffic.bytesIn, traffic.bytesOut)
|
||||||
|
}
|
||||||
|
if _, ok := batch.orphanServices["99_2_10"]; !ok {
|
||||||
|
t.Fatalf("expected orphan service cleanup target for 99_2_10")
|
||||||
|
}
|
||||||
|
if item, ok := batch.peerShareForwardItems["99_2_10"]; !ok || item.U != 12 || item.D != 8 {
|
||||||
|
t.Fatalf("expected orphan forward to remain eligible for peer-share accounting, got %#v ok=%v", item, ok)
|
||||||
|
}
|
||||||
|
if item, ok := batch.peerShareForwardItems["20_2_10"]; !ok || item.U != 110 || item.D != 80 {
|
||||||
|
t.Fatalf("expected merged peer-share forward item, got %#v ok=%v", item, ok)
|
||||||
|
}
|
||||||
|
if item, ok := batch.peerShareRuntimeItems[17]; !ok || item.U != 9 || item.D != 1 {
|
||||||
|
t.Fatalf("expected merged peer-share runtime item, got %#v ok=%v", item, ok)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestApplyFlowUploadBatchContinuesPolicyAndPeerShareSideEffectsWhenQuotaBatchFails(t *testing.T) {
|
||||||
|
r, err := repo.Open(filepath.Join(t.TempDir(), "flow-upload-batch-quota-fail.db"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("open repo: %v", err)
|
||||||
|
}
|
||||||
|
defer r.Close()
|
||||||
|
|
||||||
|
now := time.Now()
|
||||||
|
nowMs := now.UnixMilli()
|
||||||
|
if err := r.DB().Create(&model.User{ID: 2, User: "flow-user", Pwd: "pwd", RoleID: 1, ExpTime: 2727251700000, Flow: 99999, Num: 99999, CreatedTime: nowMs, Status: 1}).Error; err != nil {
|
||||||
|
t.Fatalf("seed user: %v", err)
|
||||||
|
}
|
||||||
|
if err := r.DB().Create(&model.Tunnel{ID: 1, Name: "tunnel-1", TrafficRatio: 1, Type: 1, Protocol: "tls", Flow: 1, CreatedTime: nowMs, UpdatedTime: nowMs, Status: 1}).Error; err != nil {
|
||||||
|
t.Fatalf("seed tunnel: %v", err)
|
||||||
|
}
|
||||||
|
if err := r.DB().Create(&model.UserTunnel{ID: 10, UserID: 2, TunnelID: 1, Num: 99999, Flow: 0, ExpTime: 2727251700000, Status: 1}).Error; err != nil {
|
||||||
|
t.Fatalf("seed user tunnel: %v", err)
|
||||||
|
}
|
||||||
|
if err := r.DB().Create(&model.Forward{ID: 20, UserID: 2, UserName: "flow-user", Name: "forward-20", TunnelID: 1, RemoteAddr: "1.1.1.1:80", Strategy: "fifo", CreatedTime: nowMs, UpdatedTime: nowMs, Status: 1}).Error; err != nil {
|
||||||
|
t.Fatalf("seed forward: %v", err)
|
||||||
|
}
|
||||||
|
if err := r.CreatePeerShare(&repo.PeerShare{Name: "share", NodeID: 1, Token: "token", MaxBandwidth: 0, CurrentFlow: 0, PortRangeStart: 31000, PortRangeEnd: 31010, IsActive: 1, CreatedTime: nowMs, UpdatedTime: nowMs}); err != nil {
|
||||||
|
t.Fatalf("create peer share: %v", err)
|
||||||
|
}
|
||||||
|
share, err := r.GetPeerShareByToken("token")
|
||||||
|
if err != nil || share == nil {
|
||||||
|
t.Fatalf("load peer share: %v", err)
|
||||||
|
}
|
||||||
|
if err := r.DB().Exec(`
|
||||||
|
INSERT INTO peer_share_runtime(share_id, node_id, reservation_id, resource_key, binding_id, role, chain_name, service_name, protocol, strategy, port, target, applied, status, created_time, updated_time)
|
||||||
|
VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||||
|
`, share.ID, 1, "svc-r1", "svc-rk1", "", "forward", "", "20_2_10", "tcp", "fifo", 31001, "", 1, 1, nowMs, nowMs).Error; err != nil {
|
||||||
|
t.Fatalf("insert peer share runtime: %v", err)
|
||||||
|
}
|
||||||
|
if err := r.DB().Exec(`
|
||||||
|
CREATE TRIGGER fail_user_quota_insert
|
||||||
|
BEFORE INSERT ON user_quota
|
||||||
|
BEGIN
|
||||||
|
SELECT RAISE(FAIL, 'quota insert blocked for test');
|
||||||
|
END;
|
||||||
|
`).Error; err != nil {
|
||||||
|
t.Fatalf("create quota failure trigger: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
h := &Handler{repo: r}
|
||||||
|
h.applyFlowUploadBatch(1, flowUploadBatch{
|
||||||
|
flowDeltas: []repo.FlowUploadCounterDelta{{ForwardID: 20, UserID: 2, UserTunnelID: 10, InFlow: 80, OutFlow: 120}},
|
||||||
|
quotaUsage: map[int64]int64{2: 200},
|
||||||
|
policyTargets: []flowPolicyTarget{{UserID: 2, UserTunnelID: 10}},
|
||||||
|
peerShareForwardItems: map[string]flowItem{"20_2_10": {N: "20_2_10", U: 120, D: 80}},
|
||||||
|
}, now)
|
||||||
|
|
||||||
|
if got := mustQueryInt(t, r, `SELECT status FROM forward WHERE id = 20`); got != 0 {
|
||||||
|
t.Fatalf("expected flow-policy enforcement to pause forward after quota failure, got status=%d", got)
|
||||||
|
}
|
||||||
|
updatedShare, err := r.GetPeerShare(share.ID)
|
||||||
|
if err != nil || updatedShare == nil {
|
||||||
|
t.Fatalf("reload peer share: %v", err)
|
||||||
|
}
|
||||||
|
if updatedShare.CurrentFlow != 200 {
|
||||||
|
t.Fatalf("expected peer-share flow accounting to continue after quota failure, got %d", updatedShare.CurrentFlow)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestApplyFlowUploadBatchContinuesPeerShareSideEffectsWhenFlowBatchFails(t *testing.T) {
|
||||||
|
r, err := repo.Open(filepath.Join(t.TempDir(), "flow-upload-batch-flow-fail.db"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("open repo: %v", err)
|
||||||
|
}
|
||||||
|
defer r.Close()
|
||||||
|
|
||||||
|
now := time.Now()
|
||||||
|
nowMs := now.UnixMilli()
|
||||||
|
if err := r.DB().Create(&model.User{ID: 2, User: "flow-user", Pwd: "pwd", RoleID: 1, ExpTime: 2727251700000, Flow: 99999, Num: 99999, CreatedTime: nowMs, Status: 1}).Error; err != nil {
|
||||||
|
t.Fatalf("seed user: %v", err)
|
||||||
|
}
|
||||||
|
if err := r.DB().Create(&model.Tunnel{ID: 1, Name: "tunnel-1", TrafficRatio: 1, Type: 1, Protocol: "tls", Flow: 1, CreatedTime: nowMs, UpdatedTime: nowMs, Status: 1}).Error; err != nil {
|
||||||
|
t.Fatalf("seed tunnel: %v", err)
|
||||||
|
}
|
||||||
|
if err := r.DB().Create(&model.UserTunnel{ID: 10, UserID: 2, TunnelID: 1, Num: 99999, Flow: 0, ExpTime: 2727251700000, Status: 1}).Error; err != nil {
|
||||||
|
t.Fatalf("seed user tunnel: %v", err)
|
||||||
|
}
|
||||||
|
if err := r.DB().Create(&model.Forward{ID: 20, UserID: 2, UserName: "flow-user", Name: "forward-20", TunnelID: 1, RemoteAddr: "1.1.1.1:80", Strategy: "fifo", CreatedTime: nowMs, UpdatedTime: nowMs, Status: 1}).Error; err != nil {
|
||||||
|
t.Fatalf("seed forward: %v", err)
|
||||||
|
}
|
||||||
|
if err := r.DB().Create(&model.Forward{ID: 21, UserID: 2, UserName: "flow-user", Name: "forward-21", TunnelID: 1, RemoteAddr: "1.1.1.1:81", Strategy: "fifo", CreatedTime: nowMs, UpdatedTime: nowMs, Status: 1}).Error; err != nil {
|
||||||
|
t.Fatalf("seed second forward: %v", err)
|
||||||
|
}
|
||||||
|
if err := r.CreatePeerShare(&repo.PeerShare{Name: "share", NodeID: 1, Token: "token", MaxBandwidth: 0, CurrentFlow: 0, PortRangeStart: 31000, PortRangeEnd: 31010, IsActive: 1, CreatedTime: nowMs, UpdatedTime: nowMs}); err != nil {
|
||||||
|
t.Fatalf("create peer share: %v", err)
|
||||||
|
}
|
||||||
|
share, err := r.GetPeerShareByToken("token")
|
||||||
|
if err != nil || share == nil {
|
||||||
|
t.Fatalf("load peer share: %v", err)
|
||||||
|
}
|
||||||
|
if err := r.DB().Exec(`
|
||||||
|
INSERT INTO peer_share_runtime(share_id, node_id, reservation_id, resource_key, binding_id, role, chain_name, service_name, protocol, strategy, port, target, applied, status, created_time, updated_time)
|
||||||
|
VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||||
|
`, share.ID, 1, "svc-r1", "svc-rk1", "", "forward", "", "20_2_10", "tcp", "fifo", 31001, "", 1, 1, nowMs, nowMs).Error; err != nil {
|
||||||
|
t.Fatalf("insert peer share runtime: %v", err)
|
||||||
|
}
|
||||||
|
if err := r.DB().Exec(`
|
||||||
|
CREATE TRIGGER fail_forward_flow_update
|
||||||
|
BEFORE UPDATE ON forward
|
||||||
|
WHEN NEW.id = 21 AND (NEW.in_flow != OLD.in_flow OR NEW.out_flow != OLD.out_flow)
|
||||||
|
BEGIN
|
||||||
|
SELECT RAISE(FAIL, 'forward flow update blocked for test');
|
||||||
|
END;
|
||||||
|
`).Error; err != nil {
|
||||||
|
t.Fatalf("create flow failure trigger: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
h := &Handler{repo: r}
|
||||||
|
h.applyFlowUploadBatch(1, flowUploadBatch{
|
||||||
|
flowDeltas: []repo.FlowUploadCounterDelta{
|
||||||
|
{ForwardID: 20, UserID: 2, UserTunnelID: 10, InFlow: 80, OutFlow: 120},
|
||||||
|
{ForwardID: 21, UserID: 2, UserTunnelID: 10, InFlow: 30, OutFlow: 40},
|
||||||
|
},
|
||||||
|
quotaUsage: map[int64]int64{2: 200},
|
||||||
|
policyTargets: []flowPolicyTarget{{UserID: 2, UserTunnelID: 10}},
|
||||||
|
peerShareForwardItems: map[string]flowItem{"20_2_10": {N: "20_2_10", U: 120, D: 80}},
|
||||||
|
}, now)
|
||||||
|
|
||||||
|
if got := mustQueryInt(t, r, `SELECT status FROM forward WHERE id = 20`); got != 0 {
|
||||||
|
t.Fatalf("expected flow-policy enforcement to pause forward after flow batch failure, got status=%d", got)
|
||||||
|
}
|
||||||
|
updatedShare, err := r.GetPeerShare(share.ID)
|
||||||
|
if err != nil || updatedShare == nil {
|
||||||
|
t.Fatalf("reload peer share: %v", err)
|
||||||
|
}
|
||||||
|
if updatedShare.CurrentFlow != 200 {
|
||||||
|
t.Fatalf("expected peer-share flow accounting to continue after flow batch failure, got %d", updatedShare.CurrentFlow)
|
||||||
|
}
|
||||||
|
if got := mustQueryInt(t, r, `SELECT in_flow FROM forward WHERE id = 20`); got != 80 {
|
||||||
|
t.Fatalf("expected flow fallback to persist forward 20 in_flow=80, got %d", got)
|
||||||
|
}
|
||||||
|
if got := mustQueryInt(t, r, `SELECT in_flow FROM forward WHERE id = 21`); got != 0 {
|
||||||
|
t.Fatalf("expected failed forward 21 delta to remain unapplied, got %d", got)
|
||||||
|
}
|
||||||
|
if got := mustQueryInt(t, r, `SELECT in_flow FROM user WHERE id = 2`); got != 80 {
|
||||||
|
t.Fatalf("expected flow fallback to preserve successful user totals, got %d", got)
|
||||||
|
}
|
||||||
|
if got := mustQueryInt(t, r, `SELECT in_flow FROM user_tunnel WHERE id = 10`); got != 80 {
|
||||||
|
t.Fatalf("expected flow fallback to preserve successful user_tunnel totals, got %d", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestApplyFlowUploadBatchFallsBackToPerUserQuotaUpdates(t *testing.T) {
|
||||||
|
r, err := repo.Open(filepath.Join(t.TempDir(), "flow-upload-batch-quota-fallback.db"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("open repo: %v", err)
|
||||||
|
}
|
||||||
|
defer r.Close()
|
||||||
|
|
||||||
|
now := time.Now()
|
||||||
|
nowMs := now.UnixMilli()
|
||||||
|
dayKey := int64(now.Year()*10000 + int(now.Month())*100 + now.Day())
|
||||||
|
monthKey := int64(now.Year()*100 + int(now.Month()))
|
||||||
|
if err := r.DB().Exec(`INSERT INTO user(id, user, pwd, role_id, exp_time, flow, in_flow, out_flow, flow_reset_time, num, created_time, updated_time, status) VALUES(2, 'u2', 'pwd', 1, 2727251700000, 99999, 0, 0, 1, 99999, ?, ?, 1)`, nowMs, nowMs).Error; err != nil {
|
||||||
|
t.Fatalf("insert user 2: %v", err)
|
||||||
|
}
|
||||||
|
if err := r.DB().Exec(`INSERT INTO user(id, user, pwd, role_id, exp_time, flow, in_flow, out_flow, flow_reset_time, num, created_time, updated_time, status) VALUES(3, 'u3', 'pwd', 1, 2727251700000, 99999, 0, 0, 1, 99999, ?, ?, 1)`, nowMs, nowMs).Error; err != nil {
|
||||||
|
t.Fatalf("insert user 3: %v", err)
|
||||||
|
}
|
||||||
|
if err := r.DB().Exec(`INSERT INTO user_quota(user_id, daily_limit_gb, monthly_limit_gb, daily_used_bytes, monthly_used_bytes, day_key, month_key, disabled_by_quota, disabled_at, paused_forward_ids, created_time, updated_time) VALUES(2, 0, 0, 0, 0, ?, ?, 0, 0, '', ?, ?), (3, 0, 0, 0, 0, ?, ?, 0, 0, '', ?, ?)`, dayKey, monthKey, nowMs, nowMs, dayKey, monthKey, nowMs, nowMs).Error; err != nil {
|
||||||
|
t.Fatalf("insert user quotas: %v", err)
|
||||||
|
}
|
||||||
|
if err := r.DB().Exec(`
|
||||||
|
CREATE TRIGGER fail_user_3_quota_update
|
||||||
|
BEFORE UPDATE ON user_quota
|
||||||
|
WHEN NEW.user_id = 3 AND (NEW.daily_used_bytes != OLD.daily_used_bytes OR NEW.monthly_used_bytes != OLD.monthly_used_bytes)
|
||||||
|
BEGIN
|
||||||
|
SELECT RAISE(FAIL, 'quota update blocked for user 3');
|
||||||
|
END;
|
||||||
|
`).Error; err != nil {
|
||||||
|
t.Fatalf("create quota fallback trigger: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
h := &Handler{repo: r}
|
||||||
|
h.applyFlowUploadBatch(1, flowUploadBatch{quotaUsage: map[int64]int64{2: 200, 3: 300}}, now)
|
||||||
|
|
||||||
|
if got := mustQueryInt(t, r, `SELECT daily_used_bytes FROM user_quota WHERE user_id = 2`); got != 200 {
|
||||||
|
t.Fatalf("expected quota fallback to persist user 2 usage, got %d", got)
|
||||||
|
}
|
||||||
|
if got := mustQueryInt(t, r, `SELECT daily_used_bytes FROM user_quota WHERE user_id = 3`); got != 0 {
|
||||||
|
t.Fatalf("expected failed user 3 quota delta to remain unapplied, got %d", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -7,6 +7,7 @@ import (
|
|||||||
"encoding/json"
|
"encoding/json"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
|
"log"
|
||||||
"net/http"
|
"net/http"
|
||||||
"net/url"
|
"net/url"
|
||||||
"sort"
|
"sort"
|
||||||
@@ -797,11 +798,16 @@ func (h *Handler) flowUpload(w http.ResponseWriter, r *http.Request) {
|
|||||||
if err == nil && strings.TrimSpace(raw) != "" {
|
if err == nil && strings.TrimSpace(raw) != "" {
|
||||||
var items []flowItem
|
var items []flowItem
|
||||||
if json.Unmarshal([]byte(raw), &items) == nil {
|
if json.Unmarshal([]byte(raw), &items) == nil {
|
||||||
nowMs := time.Now().UnixMilli()
|
now := time.Now()
|
||||||
h.recordTunnelMetricsFromFlowItems(node.ID, items, nowMs)
|
forwardIDs := collectFlowUploadForwardIDs(items)
|
||||||
for _, item := range items {
|
metas, metaErr := h.repo.GetFlowUploadForwardMetas(forwardIDs)
|
||||||
h.processFlowItem(node.ID, item)
|
if metaErr != nil {
|
||||||
|
log.Printf("flow upload metadata lookup failed node_id=%d err=%v", node.ID, metaErr)
|
||||||
|
metas = map[int64]repo.FlowUploadForwardMeta{}
|
||||||
}
|
}
|
||||||
|
batch := h.buildFlowUploadBatch(items, metas)
|
||||||
|
h.recordTunnelMetricsFromForwardBatch(node.ID, batch.forwardTraffic, metas, now.UnixMilli())
|
||||||
|
h.applyFlowUploadBatch(node.ID, batch, now)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -6,6 +6,7 @@ import (
|
|||||||
"time"
|
"time"
|
||||||
|
|
||||||
"go-backend/internal/store/model"
|
"go-backend/internal/store/model"
|
||||||
|
"go-backend/internal/store/repo"
|
||||||
)
|
)
|
||||||
|
|
||||||
type tunnelTrafficDelta struct {
|
type tunnelTrafficDelta struct {
|
||||||
@@ -21,75 +22,42 @@ func unixMilliBucketMinute(nowMs int64) int64 {
|
|||||||
return nowMs - (nowMs % minuteMs)
|
return nowMs - (nowMs % minuteMs)
|
||||||
}
|
}
|
||||||
|
|
||||||
func (h *Handler) recordTunnelMetricsFromFlowItems(nodeID int64, items []flowItem, nowMs int64) {
|
func collectFlowUploadForwardIDs(items []flowItem) []int64 {
|
||||||
if h == nil || h.repo == nil {
|
ids := make([]int64, 0, len(items))
|
||||||
return
|
seen := make(map[int64]struct{}, len(items))
|
||||||
}
|
for _, item := range items {
|
||||||
if nodeID <= 0 || len(items) == 0 {
|
forwardID, _, _, ok := parseFlowServiceIDs(strings.TrimSpace(item.N))
|
||||||
return
|
if !ok || forwardID <= 0 {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if _, exists := seen[forwardID]; exists {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
seen[forwardID] = struct{}{}
|
||||||
|
ids = append(ids, forwardID)
|
||||||
}
|
}
|
||||||
|
return ids
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *Handler) recordTunnelMetricsFromForwardBatch(nodeID int64, forwardDeltas map[int64]tunnelTrafficDelta, metas map[int64]repo.FlowUploadForwardMeta, nowMs int64) {
|
||||||
|
if h == nil || h.repo == nil || nodeID <= 0 || len(forwardDeltas) == 0 {
|
||||||
|
return
|
||||||
|
}
|
||||||
bucketTs := unixMilliBucketMinute(nowMs)
|
bucketTs := unixMilliBucketMinute(nowMs)
|
||||||
if bucketTs <= 0 {
|
if bucketTs <= 0 {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
forwardDeltas := make(map[int64]tunnelTrafficDelta)
|
|
||||||
var skippedParse, skippedZero int
|
|
||||||
for _, item := range items {
|
|
||||||
name := strings.TrimSpace(item.N)
|
|
||||||
if name == "" || name == "web_api" {
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
forwardID, _, _, ok := parseFlowServiceIDs(name)
|
|
||||||
if !ok {
|
|
||||||
skippedParse++
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
if item.D == 0 && item.U == 0 {
|
|
||||||
skippedZero++
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
d := forwardDeltas[forwardID]
|
|
||||||
d.bytesIn += item.D
|
|
||||||
d.bytesOut += item.U
|
|
||||||
forwardDeltas[forwardID] = d
|
|
||||||
}
|
|
||||||
if len(forwardDeltas) == 0 {
|
|
||||||
if len(items) > 0 {
|
|
||||||
log.Printf("monitoring debug op=tunnel_metric.no_forward_deltas node_id=%d items=%d skipped_parse=%d skipped_zero=%d", nodeID, len(items), skippedParse, skippedZero)
|
|
||||||
}
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
forwardIDs := make([]int64, 0, len(forwardDeltas))
|
|
||||||
for id := range forwardDeltas {
|
|
||||||
forwardIDs = append(forwardIDs, id)
|
|
||||||
}
|
|
||||||
|
|
||||||
forwardTunnelMap, err := h.repo.MapForwardIDsToTunnelIDs(forwardIDs)
|
|
||||||
if err != nil {
|
|
||||||
log.Printf("monitoring write skipped op=tunnel_metric.map_forward_to_tunnel node_id=%d err=%v", nodeID, err)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
if len(forwardTunnelMap) == 0 {
|
|
||||||
log.Printf("monitoring debug op=tunnel_metric.no_tunnel_map node_id=%d forward_ids=%v", nodeID, forwardIDs)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
tunnelAgg := make(map[int64]tunnelTrafficDelta)
|
tunnelAgg := make(map[int64]tunnelTrafficDelta)
|
||||||
for forwardID, delta := range forwardDeltas {
|
for forwardID, delta := range forwardDeltas {
|
||||||
tunnelID := forwardTunnelMap[forwardID]
|
meta, ok := metas[forwardID]
|
||||||
if tunnelID <= 0 {
|
if !ok || meta.TunnelID <= 0 {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
a := tunnelAgg[tunnelID]
|
current := tunnelAgg[meta.TunnelID]
|
||||||
a.bytesIn += delta.bytesIn
|
current.bytesIn += delta.bytesIn
|
||||||
a.bytesOut += delta.bytesOut
|
current.bytesOut += delta.bytesOut
|
||||||
tunnelAgg[tunnelID] = a
|
tunnelAgg[meta.TunnelID] = current
|
||||||
}
|
|
||||||
if len(tunnelAgg) == 0 {
|
|
||||||
return
|
|
||||||
}
|
}
|
||||||
|
|
||||||
metrics := make([]*model.TunnelMetric, 0, len(tunnelAgg))
|
metrics := make([]*model.TunnelMetric, 0, len(tunnelAgg))
|
||||||
@@ -98,14 +66,11 @@ func (h *Handler) recordTunnelMetricsFromFlowItems(nodeID int64, items []flowIte
|
|||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
metrics = append(metrics, &model.TunnelMetric{
|
metrics = append(metrics, &model.TunnelMetric{
|
||||||
TunnelID: tunnelID,
|
TunnelID: tunnelID,
|
||||||
NodeID: nodeID,
|
NodeID: nodeID,
|
||||||
Timestamp: bucketTs,
|
Timestamp: bucketTs,
|
||||||
BytesIn: delta.bytesIn,
|
BytesIn: delta.bytesIn,
|
||||||
BytesOut: delta.bytesOut,
|
BytesOut: delta.bytesOut,
|
||||||
Connections: 0,
|
|
||||||
Errors: 0,
|
|
||||||
AvgLatencyMs: 0,
|
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
if len(metrics) == 0 {
|
if len(metrics) == 0 {
|
||||||
@@ -114,7 +79,7 @@ func (h *Handler) recordTunnelMetricsFromFlowItems(nodeID int64, items []flowIte
|
|||||||
|
|
||||||
if err := h.repo.UpsertTunnelMetricBuckets(metrics); err != nil {
|
if err := h.repo.UpsertTunnelMetricBuckets(metrics); err != nil {
|
||||||
log.Printf("monitoring write failed op=tunnel_metric.upsert_buckets node_id=%d bucket_ts=%d count=%d err=%v", nodeID, bucketTs, len(metrics), err)
|
log.Printf("monitoring write failed op=tunnel_metric.upsert_buckets node_id=%d bucket_ts=%d count=%d err=%v", nodeID, bucketTs, len(metrics), err)
|
||||||
} else {
|
return
|
||||||
log.Printf("monitoring ok op=tunnel_metric.upsert_buckets node_id=%d bucket_ts=%d count=%d", nodeID, bucketTs, len(metrics))
|
|
||||||
}
|
}
|
||||||
|
log.Printf("monitoring ok op=tunnel_metric.upsert_buckets node_id=%d bucket_ts=%d count=%d", nodeID, bucketTs, len(metrics))
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -61,6 +61,14 @@ type Repository struct {
|
|||||||
db *gorm.DB
|
db *gorm.DB
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type FlowUploadCounterDelta struct {
|
||||||
|
ForwardID int64
|
||||||
|
UserID int64
|
||||||
|
UserTunnelID int64
|
||||||
|
InFlow int64
|
||||||
|
OutFlow int64
|
||||||
|
}
|
||||||
|
|
||||||
func (r *Repository) DB() *gorm.DB {
|
func (r *Repository) DB() *gorm.DB {
|
||||||
if r == nil {
|
if r == nil {
|
||||||
return nil
|
return nil
|
||||||
@@ -68,6 +76,79 @@ func (r *Repository) DB() *gorm.DB {
|
|||||||
return r.db
|
return r.db
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func sortedFlowUploadTargetIDs(totals map[int64][2]int64) []int64 {
|
||||||
|
ids := make([]int64, 0, len(totals))
|
||||||
|
for id := range totals {
|
||||||
|
ids = append(ids, id)
|
||||||
|
}
|
||||||
|
sort.Slice(ids, func(i, j int) bool { return ids[i] < ids[j] })
|
||||||
|
return ids
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *Repository) ApplyFlowUploadDeltasBatch(deltas []FlowUploadCounterDelta) error {
|
||||||
|
if r == nil || r.db == nil {
|
||||||
|
return errors.New("repository not initialized")
|
||||||
|
}
|
||||||
|
if len(deltas) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
forwardTotals := make(map[int64][2]int64, len(deltas))
|
||||||
|
userTotals := make(map[int64][2]int64, len(deltas))
|
||||||
|
userTunnelTotals := make(map[int64][2]int64, len(deltas))
|
||||||
|
for _, delta := range deltas {
|
||||||
|
if delta.ForwardID > 0 {
|
||||||
|
current := forwardTotals[delta.ForwardID]
|
||||||
|
current[0] += delta.InFlow
|
||||||
|
current[1] += delta.OutFlow
|
||||||
|
forwardTotals[delta.ForwardID] = current
|
||||||
|
}
|
||||||
|
if delta.UserID > 0 {
|
||||||
|
current := userTotals[delta.UserID]
|
||||||
|
current[0] += delta.InFlow
|
||||||
|
current[1] += delta.OutFlow
|
||||||
|
userTotals[delta.UserID] = current
|
||||||
|
}
|
||||||
|
if delta.UserTunnelID > 0 {
|
||||||
|
current := userTunnelTotals[delta.UserTunnelID]
|
||||||
|
current[0] += delta.InFlow
|
||||||
|
current[1] += delta.OutFlow
|
||||||
|
userTunnelTotals[delta.UserTunnelID] = current
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return r.db.Transaction(func(tx *gorm.DB) error {
|
||||||
|
for _, forwardID := range sortedFlowUploadTargetIDs(forwardTotals) {
|
||||||
|
total := forwardTotals[forwardID]
|
||||||
|
if err := tx.Model(&model.Forward{}).Where("id = ?", forwardID).UpdateColumns(map[string]interface{}{
|
||||||
|
"in_flow": gorm.Expr("in_flow + ?", total[0]),
|
||||||
|
"out_flow": gorm.Expr("out_flow + ?", total[1]),
|
||||||
|
}).Error; err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, userID := range sortedFlowUploadTargetIDs(userTotals) {
|
||||||
|
total := userTotals[userID]
|
||||||
|
if err := tx.Model(&model.User{}).Where("id = ?", userID).UpdateColumns(map[string]interface{}{
|
||||||
|
"in_flow": gorm.Expr("in_flow + ?", total[0]),
|
||||||
|
"out_flow": gorm.Expr("out_flow + ?", total[1]),
|
||||||
|
}).Error; err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, userTunnelID := range sortedFlowUploadTargetIDs(userTunnelTotals) {
|
||||||
|
total := userTunnelTotals[userTunnelID]
|
||||||
|
if err := tx.Model(&model.UserTunnel{}).Where("id = ?", userTunnelID).UpdateColumns(map[string]interface{}{
|
||||||
|
"in_flow": gorm.Expr("in_flow + ?", total[0]),
|
||||||
|
"out_flow": gorm.Expr("out_flow + ?", total[1]),
|
||||||
|
}).Error; err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
// ─── Open / Close ────────────────────────────────────────────────────
|
// ─── Open / Close ────────────────────────────────────────────────────
|
||||||
|
|
||||||
func Open(path string) (*Repository, error) {
|
func Open(path string) (*Repository, error) {
|
||||||
|
|||||||
@@ -9,6 +9,90 @@ import (
|
|||||||
"go-backend/internal/store/model"
|
"go-backend/internal/store/model"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
type FlowUploadForwardMeta struct {
|
||||||
|
ForwardID int64
|
||||||
|
TunnelID int64
|
||||||
|
TrafficRatio float64
|
||||||
|
TunnelFlow int64
|
||||||
|
}
|
||||||
|
|
||||||
|
const flowUploadForwardMetaChunkSize = 500
|
||||||
|
|
||||||
|
func chunkFlowUploadForwardIDs(ids []int64) [][]int64 {
|
||||||
|
if len(ids) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
chunks := make([][]int64, 0, (len(ids)+flowUploadForwardMetaChunkSize-1)/flowUploadForwardMetaChunkSize)
|
||||||
|
for start := 0; start < len(ids); start += flowUploadForwardMetaChunkSize {
|
||||||
|
end := start + flowUploadForwardMetaChunkSize
|
||||||
|
if end > len(ids) {
|
||||||
|
end = len(ids)
|
||||||
|
}
|
||||||
|
chunks = append(chunks, ids[start:end])
|
||||||
|
}
|
||||||
|
return chunks
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *Repository) GetFlowUploadForwardMetas(forwardIDs []int64) (map[int64]FlowUploadForwardMeta, error) {
|
||||||
|
if r == nil || r.db == nil {
|
||||||
|
return nil, errors.New("repository not initialized")
|
||||||
|
}
|
||||||
|
if len(forwardIDs) == 0 {
|
||||||
|
return map[int64]FlowUploadForwardMeta{}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
ids := make([]int64, 0, len(forwardIDs))
|
||||||
|
seen := make(map[int64]struct{}, len(forwardIDs))
|
||||||
|
for _, id := range forwardIDs {
|
||||||
|
if id <= 0 {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if _, ok := seen[id]; ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
seen[id] = struct{}{}
|
||||||
|
ids = append(ids, id)
|
||||||
|
}
|
||||||
|
if len(ids) == 0 {
|
||||||
|
return map[int64]FlowUploadForwardMeta{}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
type row struct {
|
||||||
|
ForwardID int64 `gorm:"column:forward_id"`
|
||||||
|
TunnelID int64 `gorm:"column:tunnel_id"`
|
||||||
|
TrafficRatio float64 `gorm:"column:traffic_ratio"`
|
||||||
|
TunnelFlow int64 `gorm:"column:tunnel_flow"`
|
||||||
|
}
|
||||||
|
|
||||||
|
out := make(map[int64]FlowUploadForwardMeta, len(ids))
|
||||||
|
for _, chunk := range chunkFlowUploadForwardIDs(ids) {
|
||||||
|
var rows []row
|
||||||
|
err := r.db.Table("forward AS f").
|
||||||
|
Select("f.id AS forward_id, f.tunnel_id AS tunnel_id, t.traffic_ratio AS traffic_ratio, t.flow AS tunnel_flow").
|
||||||
|
Joins("LEFT JOIN tunnel t ON t.id = f.tunnel_id").
|
||||||
|
Where("f.id IN ?", chunk).
|
||||||
|
Scan(&rows).Error
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
for _, row := range rows {
|
||||||
|
if row.TunnelFlow <= 0 {
|
||||||
|
row.TunnelFlow = 1
|
||||||
|
}
|
||||||
|
if row.TrafficRatio <= 0 {
|
||||||
|
row.TrafficRatio = 1
|
||||||
|
}
|
||||||
|
out[row.ForwardID] = FlowUploadForwardMeta{
|
||||||
|
ForwardID: row.ForwardID,
|
||||||
|
TunnelID: row.TunnelID,
|
||||||
|
TrafficRatio: row.TrafficRatio,
|
||||||
|
TunnelFlow: row.TunnelFlow,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return out, nil
|
||||||
|
}
|
||||||
|
|
||||||
func (r *Repository) UpdateForwardStatus(forwardID int64, status int, now int64) error {
|
func (r *Repository) UpdateForwardStatus(forwardID int64, status int, now int64) error {
|
||||||
if r == nil || r.db == nil {
|
if r == nil || r.db == nil {
|
||||||
return errors.New("repository not initialized")
|
return errors.New("repository not initialized")
|
||||||
|
|||||||
@@ -0,0 +1,142 @@
|
|||||||
|
package repo
|
||||||
|
|
||||||
|
import (
|
||||||
|
"path/filepath"
|
||||||
|
"reflect"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestChunkFlowUploadForwardIDs(t *testing.T) {
|
||||||
|
ids := make([]int64, 0, 1001)
|
||||||
|
for i := int64(1); i <= 1001; i++ {
|
||||||
|
ids = append(ids, i)
|
||||||
|
}
|
||||||
|
|
||||||
|
chunks := chunkFlowUploadForwardIDs(ids)
|
||||||
|
if len(chunks) != 3 {
|
||||||
|
t.Fatalf("expected 3 chunks, got %d", len(chunks))
|
||||||
|
}
|
||||||
|
if len(chunks[0]) != 500 || len(chunks[1]) != 500 || len(chunks[2]) != 1 {
|
||||||
|
t.Fatalf("unexpected chunk sizes: %d, %d, %d", len(chunks[0]), len(chunks[1]), len(chunks[2]))
|
||||||
|
}
|
||||||
|
if chunks[0][0] != 1 || chunks[1][0] != 501 || chunks[2][0] != 1001 {
|
||||||
|
t.Fatalf("unexpected chunk boundaries: %#v %#v %#v", chunks[0][:1], chunks[1][:1], chunks[2][:1])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSortedFlowUploadTargetIDs(t *testing.T) {
|
||||||
|
totals := map[int64][2]int64{
|
||||||
|
9: {1, 1},
|
||||||
|
2: {1, 1},
|
||||||
|
7: {1, 1},
|
||||||
|
}
|
||||||
|
|
||||||
|
got := sortedFlowUploadTargetIDs(totals)
|
||||||
|
want := []int64{2, 7, 9}
|
||||||
|
if !reflect.DeepEqual(got, want) {
|
||||||
|
t.Fatalf("expected sorted ids %v, got %v", want, got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGetFlowUploadForwardMetasAndApplyFlowUploadDeltasBatch(t *testing.T) {
|
||||||
|
r, err := Open(filepath.Join(t.TempDir(), "flow-batch.db"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("open repo: %v", err)
|
||||||
|
}
|
||||||
|
defer r.Close()
|
||||||
|
|
||||||
|
now := time.Now().UnixMilli()
|
||||||
|
if err := r.DB().Exec(`INSERT INTO user(id, user, pwd, role_id, exp_time, flow, in_flow, out_flow, flow_reset_time, num, created_time, updated_time, status) VALUES(2, 'u2', 'pwd', 1, 2727251700000, 99999, 0, 0, 1, 99999, ?, ?, 1)`, now, now).Error; err != nil {
|
||||||
|
t.Fatalf("insert user: %v", err)
|
||||||
|
}
|
||||||
|
if err := r.DB().Exec(`INSERT INTO tunnel(id, name, traffic_ratio, type, protocol, flow, created_time, updated_time, status, in_ip, inx) VALUES(1, 't1', 2.0, 1, 'tls', 3, ?, ?, 1, NULL, 0)`, now, now).Error; err != nil {
|
||||||
|
t.Fatalf("insert tunnel: %v", err)
|
||||||
|
}
|
||||||
|
if err := r.DB().Exec(`INSERT INTO user_tunnel(id, user_id, tunnel_id, speed_id, num, flow, in_flow, out_flow, flow_reset_time, exp_time, status) VALUES(10, 2, 1, NULL, 99999, 99999, 0, 0, 1, 2727251700000, 1)`).Error; err != nil {
|
||||||
|
t.Fatalf("insert user_tunnel: %v", err)
|
||||||
|
}
|
||||||
|
if err := r.DB().Exec(`INSERT INTO forward(id, user_id, user_name, name, tunnel_id, remote_addr, strategy, in_flow, out_flow, created_time, updated_time, status, inx) VALUES(20, 2, 'u2', 'f20', 1, '1.1.1.1:80', 'fifo', 0, 0, ?, ?, 1, 0)`, now, now).Error; err != nil {
|
||||||
|
t.Fatalf("insert forward: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
metas, err := r.GetFlowUploadForwardMetas([]int64{20, 99})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("get metas: %v", err)
|
||||||
|
}
|
||||||
|
if metas[20].TunnelID != 1 || metas[20].TrafficRatio != 2 || metas[20].TunnelFlow != 3 {
|
||||||
|
t.Fatalf("unexpected meta for forward 20: %#v", metas[20])
|
||||||
|
}
|
||||||
|
if _, ok := metas[99]; ok {
|
||||||
|
t.Fatalf("did not expect meta for missing forward 99")
|
||||||
|
}
|
||||||
|
|
||||||
|
err = r.ApplyFlowUploadDeltasBatch([]FlowUploadCounterDelta{{ForwardID: 20, UserID: 2, UserTunnelID: 10, InFlow: 480, OutFlow: 660}})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("apply flow batch: %v", err)
|
||||||
|
}
|
||||||
|
if got := mustFlowBatchCount(t, r, `SELECT in_flow FROM forward WHERE id = 20`); got != 480 {
|
||||||
|
t.Fatalf("expected forward in_flow=480, got %d", got)
|
||||||
|
}
|
||||||
|
if got := mustFlowBatchCount(t, r, `SELECT out_flow FROM user WHERE id = 2`); got != 660 {
|
||||||
|
t.Fatalf("expected user out_flow=660, got %d", got)
|
||||||
|
}
|
||||||
|
if got := mustFlowBatchCount(t, r, `SELECT in_flow FROM user_tunnel WHERE id = 10`); got != 480 {
|
||||||
|
t.Fatalf("expected user_tunnel in_flow=480, got %d", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGetFlowUploadForwardMetasKeepsForwardsWhenTunnelRowMissing(t *testing.T) {
|
||||||
|
r, err := Open(filepath.Join(t.TempDir(), "flow-batch-missing-tunnel.db"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("open repo: %v", err)
|
||||||
|
}
|
||||||
|
defer r.Close()
|
||||||
|
|
||||||
|
now := time.Now().UnixMilli()
|
||||||
|
if err := r.DB().Exec(`INSERT INTO forward(id, user_id, user_name, name, tunnel_id, remote_addr, strategy, in_flow, out_flow, created_time, updated_time, status, inx) VALUES(25, 2, 'u2', 'f25', 99, '1.1.1.1:80', 'fifo', 0, 0, ?, ?, 1, 0)`, now, now).Error; err != nil {
|
||||||
|
t.Fatalf("insert forward: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
metas, err := r.GetFlowUploadForwardMetas([]int64{25})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("get metas: %v", err)
|
||||||
|
}
|
||||||
|
meta, ok := metas[25]
|
||||||
|
if !ok {
|
||||||
|
t.Fatalf("expected metadata for forward with missing tunnel row")
|
||||||
|
}
|
||||||
|
if meta.ForwardID != 25 || meta.TunnelID != 99 || meta.TrafficRatio != 1 || meta.TunnelFlow != 1 {
|
||||||
|
t.Fatalf("unexpected fallback meta: %#v", meta)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAddUserQuotaUsageBatchReturnsNormalizedViews(t *testing.T) {
|
||||||
|
r, err := Open(filepath.Join(t.TempDir(), "quota-batch.db"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("open repo: %v", err)
|
||||||
|
}
|
||||||
|
defer r.Close()
|
||||||
|
|
||||||
|
now := time.Now()
|
||||||
|
nowMs := now.UnixMilli()
|
||||||
|
if err := r.DB().Exec(`INSERT INTO user(id, user, pwd, role_id, exp_time, flow, in_flow, out_flow, flow_reset_time, num, created_time, updated_time, status) VALUES(2, 'u2', 'pwd', 1, 2727251700000, 99999, 0, 0, 1, 99999, ?, ?, 1)`, nowMs, nowMs).Error; err != nil {
|
||||||
|
t.Fatalf("insert user: %v", err)
|
||||||
|
}
|
||||||
|
views, err := r.AddUserQuotaUsageBatch(map[int64]int64{2: 1140}, now)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("batch quota update: %v", err)
|
||||||
|
}
|
||||||
|
if views[2] == nil || views[2].DailyUsedBytes != 1140 || views[2].MonthlyUsedBytes != 1140 {
|
||||||
|
t.Fatalf("unexpected quota view: %#v", views[2])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func mustFlowBatchCount(t *testing.T, r *Repository, query string, args ...interface{}) int64 {
|
||||||
|
t.Helper()
|
||||||
|
var value int64
|
||||||
|
if err := r.DB().Raw(query, args...).Row().Scan(&value); err != nil {
|
||||||
|
t.Fatalf("query %q failed: %v", query, err)
|
||||||
|
}
|
||||||
|
return value
|
||||||
|
}
|
||||||
@@ -3,6 +3,7 @@ package repo
|
|||||||
import (
|
import (
|
||||||
"errors"
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"sort"
|
||||||
"strconv"
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
@@ -255,6 +256,54 @@ func (r *Repository) AddUserQuotaUsage(userID int64, usedBytes int64, now time.T
|
|||||||
return normalizeUserQuotaView(result, now), nil
|
return normalizeUserQuotaView(result, now), nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (r *Repository) AddUserQuotaUsageBatch(usages map[int64]int64, now time.Time) (map[int64]*model.UserQuotaView, error) {
|
||||||
|
if r == nil || r.db == nil {
|
||||||
|
return nil, errors.New("repository not initialized")
|
||||||
|
}
|
||||||
|
if len(usages) == 0 {
|
||||||
|
return map[int64]*model.UserQuotaView{}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
result := make(map[int64]*model.UserQuotaView, len(usages))
|
||||||
|
err := r.db.Transaction(func(tx *gorm.DB) error {
|
||||||
|
userIDs := make([]int64, 0, len(usages))
|
||||||
|
for userID := range usages {
|
||||||
|
if userID > 0 {
|
||||||
|
userIDs = append(userIDs, userID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
sort.Slice(userIDs, func(i, j int) bool { return userIDs[i] < userIDs[j] })
|
||||||
|
|
||||||
|
for _, userID := range userIDs {
|
||||||
|
q, err := r.loadOrCreateUserQuotaTx(tx, userID, now)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
applyUserQuotaWindowRoll(q, now)
|
||||||
|
if usages[userID] > 0 {
|
||||||
|
q.DailyUsedBytes += usages[userID]
|
||||||
|
q.MonthlyUsedBytes += usages[userID]
|
||||||
|
}
|
||||||
|
q.UpdatedTime = now.UnixMilli()
|
||||||
|
if err := tx.Model(&model.UserQuota{}).Where("user_id = ?", userID).Updates(map[string]interface{}{
|
||||||
|
"daily_used_bytes": q.DailyUsedBytes,
|
||||||
|
"monthly_used_bytes": q.MonthlyUsedBytes,
|
||||||
|
"day_key": q.DayKey,
|
||||||
|
"month_key": q.MonthKey,
|
||||||
|
"updated_time": q.UpdatedTime,
|
||||||
|
}).Error; err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
result[userID] = normalizeUserQuotaView(cloneUserQuotaView(*q), now)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return result, nil
|
||||||
|
}
|
||||||
|
|
||||||
func (r *Repository) MarkUserQuotaDisabled(userID int64, pausedForwardIDs []int64, now int64) error {
|
func (r *Repository) MarkUserQuotaDisabled(userID int64, pausedForwardIDs []int64, now int64) error {
|
||||||
if r == nil || r.db == nil {
|
if r == nil || r.db == nil {
|
||||||
return errors.New("repository not initialized")
|
return errors.New("repository not initialized")
|
||||||
|
|||||||
@@ -0,0 +1,107 @@
|
|||||||
|
package contract_test
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"encoding/json"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"go-backend/internal/store/model"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestFlowUploadAggregatesRepeatedItemsAndDisablesQuotaImmediately(t *testing.T) {
|
||||||
|
secret := "monitoring-jwt-secret"
|
||||||
|
router, repo := setupContractRouter(t, secret)
|
||||||
|
now := time.Now()
|
||||||
|
nowMs := now.UnixMilli()
|
||||||
|
dayKey := int64(now.Year()*10000 + int(now.Month())*100 + now.Day())
|
||||||
|
monthKey := int64(now.Year()*100 + int(now.Month()))
|
||||||
|
const bytesPerGB = int64(1024 * 1024 * 1024)
|
||||||
|
|
||||||
|
node := &model.Node{Name: "node-1", Secret: "node-secret", ServerIP: "127.0.0.1", Port: "10000-10010", TCPListenAddr: "[::]", UDPListenAddr: "[::]", CreatedTime: nowMs, Status: 1}
|
||||||
|
if err := repo.DB().Create(node).Error; err != nil {
|
||||||
|
t.Fatalf("seed node: %v", err)
|
||||||
|
}
|
||||||
|
if err := repo.DB().Exec(`INSERT INTO user(id, user, pwd, role_id, exp_time, flow, in_flow, out_flow, flow_reset_time, num, created_time, updated_time, status) VALUES(2, 'flow_user', 'pwd', 1, 2727251700000, 99999, 0, 0, 1, 99999, ?, ?, 1)`, nowMs, nowMs).Error; err != nil {
|
||||||
|
t.Fatalf("insert user: %v", err)
|
||||||
|
}
|
||||||
|
tunnel := &model.Tunnel{Name: "tunnel-1", TrafficRatio: 1.0, Type: 1, Protocol: "tls", Flow: 1, CreatedTime: nowMs, UpdatedTime: nowMs, Status: 1}
|
||||||
|
if err := repo.DB().Create(tunnel).Error; err != nil {
|
||||||
|
t.Fatalf("seed tunnel: %v", err)
|
||||||
|
}
|
||||||
|
if err := repo.DB().Exec(`INSERT INTO user_tunnel(id, user_id, tunnel_id, speed_id, num, flow, in_flow, out_flow, flow_reset_time, exp_time, status) VALUES(10, 2, ?, NULL, 99999, 99999, 0, 0, 1, 2727251700000, 1)`, tunnel.ID).Error; err != nil {
|
||||||
|
t.Fatalf("insert user_tunnel: %v", err)
|
||||||
|
}
|
||||||
|
forward := &model.Forward{ID: 20, UserID: 2, UserName: "flow_user", Name: "forward-20", TunnelID: tunnel.ID, RemoteAddr: "1.1.1.1:80", Strategy: "fifo", CreatedTime: nowMs, UpdatedTime: nowMs, Status: 1}
|
||||||
|
if err := repo.DB().Create(forward).Error; err != nil {
|
||||||
|
t.Fatalf("seed forward: %v", err)
|
||||||
|
}
|
||||||
|
if err := repo.DB().Exec(`INSERT INTO user_quota(user_id, daily_limit_gb, monthly_limit_gb, daily_used_bytes, monthly_used_bytes, day_key, month_key, disabled_by_quota, disabled_at, paused_forward_ids, created_time, updated_time) VALUES(2, 1, 0, ?, ?, ?, ?, 0, 0, '', ?, ?)`, bytesPerGB-100, bytesPerGB-100, dayKey, monthKey, nowMs, nowMs).Error; err != nil {
|
||||||
|
t.Fatalf("insert user_quota: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
body, err := json.Marshal([]map[string]interface{}{
|
||||||
|
{"n": "20_2_10", "u": 70, "d": 50},
|
||||||
|
{"n": "20_2_10_tcp", "u": 40, "d": 30},
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("marshal body: %v", err)
|
||||||
|
}
|
||||||
|
req := httptest.NewRequest(http.MethodPost, "/flow/upload?secret="+node.Secret, bytes.NewReader(body))
|
||||||
|
res := httptest.NewRecorder()
|
||||||
|
router.ServeHTTP(res, req)
|
||||||
|
|
||||||
|
if res.Code != http.StatusOK {
|
||||||
|
t.Fatalf("expected status 200, got %d", res.Code)
|
||||||
|
}
|
||||||
|
if got := mustQueryInt(t, repo, `SELECT status FROM forward WHERE id = 20`); got != 0 {
|
||||||
|
t.Fatalf("expected forward paused immediately, got status=%d", got)
|
||||||
|
}
|
||||||
|
if got := mustQueryInt(t, repo, `SELECT disabled_by_quota FROM user_quota WHERE user_id = 2`); got != 1 {
|
||||||
|
t.Fatalf("expected quota disabled flag=1, got %d", got)
|
||||||
|
}
|
||||||
|
if got := mustQueryInt(t, repo, `SELECT in_flow FROM forward WHERE id = 20`); got != 80 {
|
||||||
|
t.Fatalf("expected forward in_flow=80, got %d", got)
|
||||||
|
}
|
||||||
|
if got := mustQueryInt(t, repo, `SELECT out_flow FROM forward WHERE id = 20`); got != 110 {
|
||||||
|
t.Fatalf("expected forward out_flow=110, got %d", got)
|
||||||
|
}
|
||||||
|
metrics, err := repo.GetTunnelMetrics(tunnel.ID, 0, nowMs+60_000)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("get tunnel metrics: %v", err)
|
||||||
|
}
|
||||||
|
if len(metrics) != 1 || metrics[0].BytesIn != 80 || metrics[0].BytesOut != 110 {
|
||||||
|
t.Fatalf("expected one aggregated metric row, got %#v", metrics)
|
||||||
|
}
|
||||||
|
|
||||||
|
body, err = json.Marshal([]map[string]interface{}{
|
||||||
|
{"n": "20_2_10", "u": 10, "d": 20},
|
||||||
|
{"n": "20_2_10", "u": 10, "d": 20},
|
||||||
|
{"n": "20_2_10_tcp", "u": 10, "d": 20},
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("marshal body: %v", err)
|
||||||
|
}
|
||||||
|
req = httptest.NewRequest(http.MethodPost, "/flow/upload?secret="+node.Secret, bytes.NewReader(body))
|
||||||
|
res = httptest.NewRecorder()
|
||||||
|
router.ServeHTTP(res, req)
|
||||||
|
|
||||||
|
if res.Code != http.StatusOK {
|
||||||
|
t.Fatalf("expected second request status 200, got %d", res.Code)
|
||||||
|
}
|
||||||
|
if got := mustQueryInt(t, repo, `SELECT in_flow FROM forward WHERE id = 20`); got != 140 {
|
||||||
|
t.Fatalf("expected forward in_flow=140 after second request, got %d", got)
|
||||||
|
}
|
||||||
|
if got := mustQueryInt(t, repo, `SELECT out_flow FROM forward WHERE id = 20`); got != 140 {
|
||||||
|
t.Fatalf("expected forward out_flow=140 after second request, got %d", got)
|
||||||
|
}
|
||||||
|
metrics, err = repo.GetTunnelMetrics(tunnel.ID, 0, nowMs+60_000)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("get tunnel metrics after second request: %v", err)
|
||||||
|
}
|
||||||
|
if len(metrics) != 1 || metrics[0].BytesIn != 140 || metrics[0].BytesOut != 140 {
|
||||||
|
t.Fatalf("expected one aggregated metric row after second request, got %#v", metrics)
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user