mirror of
https://github.com/MHSanaei/3x-ui.git
synced 2026-09-10 20:27:15 +00:00
perf(clients): batch the client record lookup in bulk operations
BulkResetTraffic resolved every address with its own GetRecordByEmail call, one SELECT per email, purely to find the disabled clients it has to re-enable. Resetting 30 clients issued 30 queries before the batched transaction even started, while BulkAdjust, BulkDelete and BulkSetEnable next to it already loaded their rows with a single chunked IN query. Those three carried a verbatim copy each of both the trim/dedupe loop and the chunked record load, so the reuse is the fix: trimmedUniqueEmails now delegates to the existing uniqueNonEmptyStrings, and clientRecordsByEmail holds the one chunked lookup all four call sites share. A DB failure during the lookup now aborts the reset instead of being swallowed per email; a missing row is still skipped, as before. The new test drives BulkResetTraffic with 3 and with 30 emails and fails unless both issue the same number of SELECTs against clients.
This commit is contained in:
@@ -344,36 +344,16 @@ func (s *ClientService) BulkAdjust(inboundSvc *InboundService, emails []string,
|
||||
|
||||
addExpiryMs := int64(addDays) * 24 * 60 * 60 * 1000
|
||||
|
||||
seen := map[string]struct{}{}
|
||||
cleanEmails := make([]string, 0, len(emails))
|
||||
for _, e := range emails {
|
||||
e = strings.TrimSpace(e)
|
||||
if e == "" {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[e]; ok {
|
||||
continue
|
||||
}
|
||||
seen[e] = struct{}{}
|
||||
cleanEmails = append(cleanEmails, e)
|
||||
}
|
||||
cleanEmails := trimmedUniqueEmails(emails)
|
||||
if len(cleanEmails) == 0 {
|
||||
return result, false, nil
|
||||
}
|
||||
|
||||
db := database.GetDB()
|
||||
|
||||
var records []model.ClientRecord
|
||||
for _, batch := range chunkStrings(cleanEmails, sqlInChunk) {
|
||||
var rows []model.ClientRecord
|
||||
if err := db.Where("email IN ?", batch).Find(&rows).Error; err != nil {
|
||||
return result, false, err
|
||||
}
|
||||
records = append(records, rows...)
|
||||
}
|
||||
recordsByEmail := make(map[string]*model.ClientRecord, len(records))
|
||||
for i := range records {
|
||||
recordsByEmail[records[i].Email] = &records[i]
|
||||
recordsByEmail, err := clientRecordsByEmail(db, cleanEmails)
|
||||
if err != nil {
|
||||
return result, false, err
|
||||
}
|
||||
|
||||
skippedReasons := map[string]string{}
|
||||
@@ -894,38 +874,22 @@ type BulkDeleteReport struct {
|
||||
func (s *ClientService) BulkDelete(inboundSvc *InboundService, emails []string, keepTraffic bool) (BulkDeleteResult, bool, error) {
|
||||
result := BulkDeleteResult{}
|
||||
|
||||
seen := map[string]struct{}{}
|
||||
cleanEmails := make([]string, 0, len(emails))
|
||||
for _, e := range emails {
|
||||
e = strings.TrimSpace(e)
|
||||
if e == "" {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[e]; ok {
|
||||
continue
|
||||
}
|
||||
seen[e] = struct{}{}
|
||||
cleanEmails = append(cleanEmails, e)
|
||||
}
|
||||
cleanEmails := trimmedUniqueEmails(emails)
|
||||
if len(cleanEmails) == 0 {
|
||||
return result, false, nil
|
||||
}
|
||||
|
||||
db := database.GetDB()
|
||||
|
||||
var records []model.ClientRecord
|
||||
for _, batch := range chunkStrings(cleanEmails, sqlInChunk) {
|
||||
var rows []model.ClientRecord
|
||||
if err := db.Where("email IN ?", batch).Find(&rows).Error; err != nil {
|
||||
return result, false, err
|
||||
}
|
||||
records = append(records, rows...)
|
||||
recordsByEmail, err := clientRecordsByEmail(db, cleanEmails)
|
||||
if err != nil {
|
||||
return result, false, err
|
||||
}
|
||||
recordsByEmail := make(map[string]*model.ClientRecord, len(records))
|
||||
tombstoneEmails := make([]string, 0, len(records))
|
||||
for i := range records {
|
||||
recordsByEmail[records[i].Email] = &records[i]
|
||||
tombstoneEmails = append(tombstoneEmails, records[i].Email)
|
||||
tombstoneEmails := make([]string, 0, len(recordsByEmail))
|
||||
for _, email := range cleanEmails {
|
||||
if recordsByEmail[email] != nil {
|
||||
tombstoneEmails = append(tombstoneEmails, email)
|
||||
}
|
||||
}
|
||||
tombstoneClientEmails(tombstoneEmails)
|
||||
|
||||
@@ -1579,36 +1543,16 @@ type BulkSetEnableReport struct {
|
||||
func (s *ClientService) BulkSetEnable(inboundSvc *InboundService, emails []string, enable bool) (BulkSetEnableResult, bool, error) {
|
||||
result := BulkSetEnableResult{}
|
||||
|
||||
seen := map[string]struct{}{}
|
||||
cleanEmails := make([]string, 0, len(emails))
|
||||
for _, e := range emails {
|
||||
e = strings.TrimSpace(e)
|
||||
if e == "" {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[e]; ok {
|
||||
continue
|
||||
}
|
||||
seen[e] = struct{}{}
|
||||
cleanEmails = append(cleanEmails, e)
|
||||
}
|
||||
cleanEmails := trimmedUniqueEmails(emails)
|
||||
if len(cleanEmails) == 0 {
|
||||
return result, false, nil
|
||||
}
|
||||
|
||||
db := database.GetDB()
|
||||
|
||||
var records []model.ClientRecord
|
||||
for _, batch := range chunkStrings(cleanEmails, sqlInChunk) {
|
||||
var rows []model.ClientRecord
|
||||
if err := db.Where("email IN ?", batch).Find(&rows).Error; err != nil {
|
||||
return result, false, err
|
||||
}
|
||||
records = append(records, rows...)
|
||||
}
|
||||
recordsByEmail := make(map[string]*model.ClientRecord, len(records))
|
||||
for i := range records {
|
||||
recordsByEmail[records[i].Email] = &records[i]
|
||||
recordsByEmail, err := clientRecordsByEmail(db, cleanEmails)
|
||||
if err != nil {
|
||||
return result, false, err
|
||||
}
|
||||
|
||||
skippedReasons := map[string]string{}
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/mhsanaei/3x-ui/v3/internal/database"
|
||||
"github.com/mhsanaei/3x-ui/v3/internal/database/model"
|
||||
)
|
||||
|
||||
// countClientTableQueries runs fn with a callback counting SELECTs against the
|
||||
// clients table, so a per-email lookup shows up as growth with the batch size.
|
||||
func countClientTableQueries(t *testing.T, name string, fn func()) int {
|
||||
t.Helper()
|
||||
db := database.GetDB()
|
||||
var n int64
|
||||
cb := "test:count_clients_query_" + name
|
||||
if err := db.Callback().Query().Before("gorm:query").Register(cb, func(tx *gorm.DB) {
|
||||
if tx.Statement != nil && tx.Statement.Table == "clients" {
|
||||
atomic.AddInt64(&n, 1)
|
||||
}
|
||||
}); err != nil {
|
||||
t.Fatalf("register query callback: %v", err)
|
||||
}
|
||||
defer func() {
|
||||
if err := db.Callback().Query().Remove(cb); err != nil {
|
||||
t.Errorf("remove query callback: %v", err)
|
||||
}
|
||||
}()
|
||||
fn()
|
||||
return int(atomic.LoadInt64(&n))
|
||||
}
|
||||
|
||||
func seedEnabledClientsForReset(t *testing.T, svc *ClientService, port int, n int, prefix string) []string {
|
||||
t.Helper()
|
||||
clients := make([]model.Client, 0, n)
|
||||
for i := range n {
|
||||
email := fmt.Sprintf("%s-%d@x", prefix, i)
|
||||
clients = append(clients, model.Client{
|
||||
Email: email,
|
||||
ID: fmt.Sprintf("%08d-1111-1111-1111-111111111111", i),
|
||||
SubID: email,
|
||||
Enable: true,
|
||||
})
|
||||
}
|
||||
ib := mkInbound(t, port, model.VLESS, clientsSettings(t, clients))
|
||||
if err := svc.SyncInbound(nil, ib.Id, clients); err != nil {
|
||||
t.Fatalf("seed linkage: %v", err)
|
||||
}
|
||||
emails := make([]string, 0, n)
|
||||
for _, c := range clients {
|
||||
mkTraffic(t, ib.Id, c.Email, 100, 200, 0, 0, true)
|
||||
emails = append(emails, c.Email)
|
||||
}
|
||||
return emails
|
||||
}
|
||||
|
||||
// TestBulkResetTraffic_DoesNotQueryPerEmail pins BulkResetTraffic's client
|
||||
// lookup to a batched read: the number of SELECTs on clients must not grow
|
||||
// with the number of emails reset.
|
||||
func TestBulkResetTraffic_DoesNotQueryPerEmail(t *testing.T) {
|
||||
setupBulkDB(t)
|
||||
svc := &ClientService{}
|
||||
inboundSvc := &InboundService{}
|
||||
|
||||
few := seedEnabledClientsForReset(t, svc, 53010, 3, "few")
|
||||
many := seedEnabledClientsForReset(t, svc, 53011, 30, "many")
|
||||
|
||||
fewCount := countClientTableQueries(t, "few", func() {
|
||||
if _, err := svc.BulkResetTraffic(inboundSvc, few); err != nil {
|
||||
t.Fatalf("BulkResetTraffic(few): %v", err)
|
||||
}
|
||||
})
|
||||
manyCount := countClientTableQueries(t, "many", func() {
|
||||
if _, err := svc.BulkResetTraffic(inboundSvc, many); err != nil {
|
||||
t.Fatalf("BulkResetTraffic(many): %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
if manyCount != fewCount {
|
||||
t.Fatalf("clients SELECTs: %d emails -> %d, %d emails -> %d; want the same batched count",
|
||||
len(few), fewCount, len(many), manyCount)
|
||||
}
|
||||
}
|
||||
@@ -270,3 +270,25 @@ func (s *ClientService) findInboundIdsByClientEmail(email string) ([]int, error)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// clientRecordsByEmail batch-loads client rows for emails, keyed by email.
|
||||
// Callers pass an already-deduplicated list; absent addresses are simply
|
||||
// missing from the map.
|
||||
func clientRecordsByEmail(tx *gorm.DB, emails []string) (map[string]*model.ClientRecord, error) {
|
||||
if tx == nil {
|
||||
tx = database.GetDB()
|
||||
}
|
||||
var records []model.ClientRecord
|
||||
for _, batch := range chunkStrings(emails, sqlInChunk) {
|
||||
var rows []model.ClientRecord
|
||||
if err := tx.Where("email IN ?", batch).Find(&rows).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
records = append(records, rows...)
|
||||
}
|
||||
byEmail := make(map[string]*model.ClientRecord, len(records))
|
||||
for i := range records {
|
||||
byEmail[records[i].Email] = &records[i]
|
||||
}
|
||||
return byEmail, nil
|
||||
}
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/mhsanaei/3x-ui/v3/internal/database"
|
||||
@@ -60,36 +59,29 @@ func (s *ClientService) BulkResetTraffic(inboundSvc *InboundService, emails []st
|
||||
if len(emails) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
seen := map[string]struct{}{}
|
||||
cleanEmails := make([]string, 0, len(emails))
|
||||
for _, e := range emails {
|
||||
e = strings.TrimSpace(e)
|
||||
if e == "" {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[e]; ok {
|
||||
continue
|
||||
}
|
||||
seen[e] = struct{}{}
|
||||
cleanEmails = append(cleanEmails, e)
|
||||
}
|
||||
cleanEmails := trimmedUniqueEmails(emails)
|
||||
if len(cleanEmails) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
recordsByEmail, err := clientRecordsByEmail(nil, cleanEmails)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
for _, e := range cleanEmails {
|
||||
rec, err := s.GetRecordByEmail(nil, e)
|
||||
if err == nil && !rec.Enable {
|
||||
updated := rec.ToClient()
|
||||
updated.Enable = true
|
||||
if _, uErr := s.Update(inboundSvc, rec.Id, *updated, rec.LimitHwid); uErr != nil {
|
||||
logger.Warning("Failed to auto-enable client during bulk traffic reset:", uErr)
|
||||
}
|
||||
rec := recordsByEmail[e]
|
||||
if rec == nil || rec.Enable {
|
||||
continue
|
||||
}
|
||||
updated := rec.ToClient()
|
||||
updated.Enable = true
|
||||
if _, uErr := s.Update(inboundSvc, rec.Id, *updated, rec.LimitHwid); uErr != nil {
|
||||
logger.Warning("Failed to auto-enable client during bulk traffic reset:", uErr)
|
||||
}
|
||||
}
|
||||
|
||||
affected := 0
|
||||
err := submitTrafficWrite(func() error {
|
||||
err = submitTrafficWrite(func() error {
|
||||
db := database.GetDB()
|
||||
return db.Transaction(func(tx *gorm.DB) error {
|
||||
if err := adjustGroupBaselinesForRemovedTraffic(tx, cleanEmails); err != nil {
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
package service
|
||||
|
||||
import "strings"
|
||||
|
||||
// sqliteMaxVars is a safe ceiling for the number of bind parameters in a
|
||||
// single SQL statement. SQLite's SQLITE_MAX_VARIABLE_NUMBER is 999 on builds
|
||||
// before 3.32 and 32766 after; staying under 999 keeps queries portable
|
||||
@@ -38,6 +40,16 @@ func uniqueNonEmptyStrings(in []string) []string {
|
||||
return out
|
||||
}
|
||||
|
||||
// trimmedUniqueEmails trims each address before deduplicating, so the bulk
|
||||
// client operations treat " a@x " and "a@x" as the same row.
|
||||
func trimmedUniqueEmails(in []string) []string {
|
||||
trimmed := make([]string, 0, len(in))
|
||||
for _, e := range in {
|
||||
trimmed = append(trimmed, strings.TrimSpace(e))
|
||||
}
|
||||
return uniqueNonEmptyStrings(trimmed)
|
||||
}
|
||||
|
||||
// uniqueInts returns a deduplicated copy of in, preserving order of first occurrence.
|
||||
func uniqueInts(in []int) []int {
|
||||
if len(in) == 0 {
|
||||
|
||||
Reference in New Issue
Block a user