mirror of
https://github.com/MHSanaei/3x-ui.git
synced 2026-11-02 11:15:03 +03:00
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:
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user