diff --git a/internal/web/service/client_bulk.go b/internal/web/service/client_bulk.go index 582f8a931..25fc0b1cb 100644 --- a/internal/web/service/client_bulk.go +++ b/internal/web/service/client_bulk.go @@ -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{} diff --git a/internal/web/service/client_bulk_reset_query_count_test.go b/internal/web/service/client_bulk_reset_query_count_test.go new file mode 100644 index 000000000..c9f3707c5 --- /dev/null +++ b/internal/web/service/client_bulk_reset_query_count_test.go @@ -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) + } +} diff --git a/internal/web/service/client_lookup.go b/internal/web/service/client_lookup.go index 5cda43942..127d71e0a 100644 --- a/internal/web/service/client_lookup.go +++ b/internal/web/service/client_lookup.go @@ -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 +} diff --git a/internal/web/service/client_traffic.go b/internal/web/service/client_traffic.go index d33f0088b..7c99c03f0 100644 --- a/internal/web/service/client_traffic.go +++ b/internal/web/service/client_traffic.go @@ -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 { diff --git a/internal/web/service/inbound_util.go b/internal/web/service/inbound_util.go index 4c343ab9c..641e23f06 100644 --- a/internal/web/service/inbound_util.go +++ b/internal/web/service/inbound_util.go @@ -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 {