fix(traffic): stop adding capped counters in SQL

Per-client counters are stored up to TrafficMax, which a node can reach
by reporting huge usage, and the group totals, group baselines and the
per-client reset lookup added two such counters inside SQL, overflowing
it: every group listing and reset failed, and no client on that reset
cycle was reset. The counters are now read separately and added in Go,
saturating at TrafficMax.
This commit is contained in:
MHSanaei
2026-10-31 14:41:12 +01:00
parent d69b924baf
commit dcdc56c64b
4 changed files with 115 additions and 55 deletions
+4 -3
View File
@@ -123,18 +123,19 @@ type ClientResetCycle struct {
TrafficResetDay int
Enable bool
Total int64
Used int64
Up int64
Down int64
}
// Depleted reports a client the quota switched off. A reset restores that one;
// a client disabled below its quota was switched off by hand and stays off.
func (c ClientResetCycle) Depleted() bool {
return c.Total > 0 && c.Used >= c.Total
return c.Total > 0 && clampedTrafficAdd(c.Up, c.Down) >= c.Total
}
func clientResetCycles(tx *gorm.DB) *gorm.DB {
return tx.Table("clients c").
Select("c.email, c.traffic_reset_day, c.enable, COALESCE(ct.total, 0) AS total, COALESCE(ct.up, 0) + COALESCE(ct.down, 0) AS used").
Select("c.email, c.traffic_reset_day, c.enable, COALESCE(ct.total, 0) AS total, COALESCE(ct.up, 0) AS up, COALESCE(ct.down, 0) AS down").
Joins("LEFT JOIN client_traffics ct ON ct.email = c.email")
}
@@ -227,3 +227,21 @@ func TestGroupResetStillZeroesAfterBaselineAdjustments(t *testing.T) {
t.Fatalf("group reset did not zero adjusted baselines: got %+v", g)
}
}
// Counters are stored up to TrafficMax (a node can report that much); adding two in SQL
// overflowed and failed every group listing, reset and baseline shift for the panel.
func TestGroupTrafficSurvivesCountersAtTheCap(t *testing.T) {
initTrafficTestDB(t)
svc := &ClientService{}
seedGroupedClient(t, "huge-a", "vip", database.TrafficMax, database.TrafficMax)
seedGroupedClient(t, "huge-b", "vip", database.TrafficMax, database.TrafficMax)
if g := groupByName(t, svc, "vip"); g.Up != database.TrafficMax || g.ClientCount != 2 {
t.Fatalf("group = %+v, want its counters held at TrafficMax", g)
}
if err := svc.ResetGroupTraffic("vip"); err != nil {
t.Fatalf("ResetGroupTraffic: %v", err)
}
if err := adjustGroupBaselinesForRemovedTraffic(database.GetDB(), []string{"huge-a", "huge-b"}); err != nil {
t.Fatalf("shift baselines: %v", err)
}
}
+79 -52
View File
@@ -2,6 +2,7 @@ package service
import (
"encoding/json"
"errors"
"sort"
"strings"
@@ -24,13 +25,8 @@ func (s *ClientService) ListGroups() ([]GroupSummary, error) {
db := database.GetDB()
// email is unique in both clients and client_traffics, so the LEFT JOIN
// never double-counts a client's traffic.
var derived []GroupSummary
if err := db.Table("clients AS c").
Select("c.group_name AS name, COUNT(*) AS client_count, COALESCE(SUM(ct.up + ct.down), 0) AS traffic_used, COALESCE(SUM(ct.up), 0) AS up, COALESCE(SUM(ct.down), 0) AS down").
Joins("LEFT JOIN client_traffics ct ON ct.email = c.email").
Where("c.group_name <> ''").
Group("c.group_name").
Scan(&derived).Error; err != nil {
derived, err := sumGroupTraffic(db, "c.group_name <> ''")
if err != nil {
return nil, err
}
var stored []model.ClientGroup
@@ -50,14 +46,14 @@ func (s *ClientService) ListGroups() ([]GroupSummary, error) {
baseUp[g.Name] = g.ResetUp
baseDown[g.Name] = g.ResetDown
}
for _, g := range derived {
merged[g.Name] = groupAgg{count: g.ClientCount, up: g.Up, down: g.Down}
for name, g := range derived {
merged[name] = groupAgg{count: g.clients, up: g.up, down: g.down}
}
out := make([]GroupSummary, 0, len(merged))
for name, agg := range merged {
up := max(agg.up-baseUp[name], 0)
down := max(agg.down-baseDown[name], 0)
out = append(out, GroupSummary{Name: name, ClientCount: agg.count, TrafficUsed: up + down, Up: up, Down: down})
up := max(clampedTrafficAdd(agg.up, -baseUp[name]), 0)
down := max(clampedTrafficAdd(agg.down, -baseDown[name]), 0)
out = append(out, GroupSummary{Name: name, ClientCount: agg.count, TrafficUsed: clampedTrafficAdd(up, down), Up: up, Down: down})
}
sort.Slice(out, func(i, j int) bool {
return strings.ToLower(out[i].Name) < strings.ToLower(out[j].Name)
@@ -81,46 +77,35 @@ func shiftGroupBaselines(tx *gorm.DB, emails []string, sign int64) error {
if len(emails) == 0 {
return nil
}
type groupDelta struct {
Name string
Up int64
Down int64
}
totals := make(map[string]*groupDelta)
totals := make(map[string]*groupCounters)
for _, batch := range chunkStrings(emails, sqlInChunk) {
var part []groupDelta
if err := tx.Table("clients AS c").
Select("c.group_name AS name, COALESCE(SUM(ct.up), 0) AS up, COALESCE(SUM(ct.down), 0) AS down").
Joins("JOIN client_traffics ct ON ct.email = c.email").
Where("c.group_name <> '' AND c.email IN ?", batch).
Group("c.group_name").
Scan(&part).Error; err != nil {
part, err := sumGroupTraffic(tx, "c.group_name <> '' AND c.email IN ?", batch)
if err != nil {
return err
}
for i := range part {
if agg, ok := totals[part[i].Name]; ok {
agg.Up += part[i].Up
agg.Down += part[i].Down
for name, p := range part {
if agg, ok := totals[name]; ok {
agg.up = clampedTrafficAdd(agg.up, p.up)
agg.down = clampedTrafficAdd(agg.down, p.down)
} else {
totals[part[i].Name] = &part[i]
totals[name] = p
}
}
}
for name, d := range totals {
if d.Up == 0 && d.Down == 0 {
if d.up == 0 && d.down == 0 {
continue
}
res := tx.Model(&model.ClientGroup{}).Where("name = ?", name).Updates(map[string]any{
"reset_up": gorm.Expr("reset_up + ?", sign*d.Up),
"reset_down": gorm.Expr("reset_down + ?", sign*d.Down),
})
if res.Error != nil {
return res.Error
var group model.ClientGroup
err := tx.Where("name = ?", name).First(&group).Error
if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) {
return err
}
if res.RowsAffected == 0 {
if err := tx.Create(&model.ClientGroup{Name: name, ResetUp: sign * d.Up, ResetDown: sign * d.Down}).Error; err != nil {
return err
}
group.Name = name
group.ResetUp = clampedTrafficAdd(group.ResetUp, sign*d.up)
group.ResetDown = clampedTrafficAdd(group.ResetDown, sign*d.down)
if err := tx.Save(&group).Error; err != nil {
return err
}
}
return nil
@@ -151,26 +136,23 @@ func (s *ClientService) ResetGroupTraffic(name string) error {
return common.NewError("group name is required")
}
db := database.GetDB()
var agg struct {
Up int64
Down int64
}
if err := db.Table("clients AS c").
Select("COALESCE(SUM(ct.up), 0) AS up, COALESCE(SUM(ct.down), 0) AS down").
Joins("LEFT JOIN client_traffics ct ON ct.email = c.email").
Where("c.group_name = ?", name).
Scan(&agg).Error; err != nil {
totals, err := sumGroupTraffic(db, "c.group_name = ?", name)
if err != nil {
return err
}
agg := groupCounters{}
if t := totals[name]; t != nil {
agg = *t
}
var count int64
if err := db.Model(&model.ClientGroup{}).Where("name = ?", name).Count(&count).Error; err != nil {
return err
}
if count == 0 {
return db.Create(&model.ClientGroup{Name: name, ResetUp: agg.Up, ResetDown: agg.Down}).Error
return db.Create(&model.ClientGroup{Name: name, ResetUp: agg.up, ResetDown: agg.down}).Error
}
return db.Model(&model.ClientGroup{}).Where("name = ?", name).
Updates(map[string]any{"reset_up": agg.Up, "reset_down": agg.Down}).Error
Updates(map[string]any{"reset_up": agg.up, "reset_down": agg.down}).Error
}
func (s *ClientService) CreateGroup(name string) error {
@@ -450,3 +432,48 @@ func (s *ClientService) replaceGroupValue(oldName, newName string) (int, error)
}
return len(records), nil
}
// groupCounters is a group's summed traffic, held within ±TrafficMax.
type groupCounters struct {
clients int
up, down int64
}
// clampedTrafficAdd adds two counters, saturating at ±TrafficMax: stored counters reach
// the cap, and SQL SUM over two of them overflowed and failed the whole query.
func clampedTrafficAdd(a, b int64) int64 {
switch {
case b > 0 && a > database.TrafficMax-b:
return database.TrafficMax
case b < 0 && a < -database.TrafficMax-b:
return -database.TrafficMax
}
return a + b
}
// sumGroupTraffic reads one row per grouped client matching where and adds them up in Go.
func sumGroupTraffic(tx *gorm.DB, where string, args ...any) (map[string]*groupCounters, error) {
var rows []struct {
Name string
Up, Down int64
}
if err := tx.Table("clients AS c").
Select("c.group_name AS name, COALESCE(ct.up, 0) AS up, COALESCE(ct.down, 0) AS down").
Joins("LEFT JOIN client_traffics ct ON ct.email = c.email").
Where(where, args...).
Scan(&rows).Error; err != nil {
return nil, err
}
totals := make(map[string]*groupCounters)
for _, r := range rows {
agg := totals[r.Name]
if agg == nil {
agg = &groupCounters{}
totals[r.Name] = agg
}
agg.clients++
agg.up = clampedTrafficAdd(agg.up, r.Up)
agg.down = clampedTrafficAdd(agg.down, r.Down)
}
return totals, nil
}
@@ -152,3 +152,17 @@ func TestClientTrafficResetValidation(t *testing.T) {
})
}
}
// A node can report counters at TrafficMax; adding two in SQL overflowed and failed the
// lookup, so no client on that cycle was reset at all.
func TestGetClientsByTrafficResetSurvivesCountersAtTheCap(t *testing.T) {
initTrafficTestDB(t)
if err := database.GetDB().Create(&model.ClientRecord{Email: "huge", Enable: false, TrafficReset: "daily"}).Error; err != nil {
t.Fatalf("seed client: %v", err)
}
seedClientRow(t, "huge", 1, database.TrafficMax, database.TrafficMax, 100)
cycles, err := (&ClientService{}).GetClientsByTrafficReset("daily")
if err != nil || len(cycles) != 1 || !cycles[0].Depleted() {
t.Fatalf("GetClientsByTrafficReset = %+v, %v; want the client, depleted", cycles, err)
}
}