test(scale): cover traffic poll, ws payloads, ip-limit job, sub and xray config at 500k

The paths that run continuously in production had no scale coverage: the 5s traffic poll (AddTraffic with its auto-renew and depleted scans), the websocket snapshot the job broadcasts while a browser is connected, the 10s ip-limit job (hasLimitIp LIKE scan + per-email settings parse), a subscription fetch inside a huge inbound, and the full Xray config build. New benchmarks reuse the XUI_SCALE_TEST / XUI_DB_TYPE gating and stay log-only. Sizes default to 10k/100k; XUI_SCALE_SIZES=500000 raises the ladder without editing code. seedScaleDataset writes inbounds, clients, client_inbounds and client_traffics directly in one transaction instead of SyncInbound, so a 500k seed takes seconds. XUI_SCALE_DB_PATH persists the seeded SQLite file for manual smoke runs against a live panel.
2026-08-19 17:40:59 +00:00 · 2026-07-02 16:12:46 +02:00
parent 28f7690224
commit 4fc301682f
6 changed files with 885 additions and 3 deletions
@@ -1,22 +1,28 @@
 package service

 import (
+	"fmt"
 	"os"
 	"path/filepath"
+	"strconv"
 	"strings"
 	"testing"
+	"time"

 	"github.com/mhsanaei/3x-ui/v3/internal/config"
 	"github.com/mhsanaei/3x-ui/v3/internal/database"
+	"github.com/mhsanaei/3x-ui/v3/internal/database/model"
 	xuilogger "github.com/mhsanaei/3x-ui/v3/internal/logger"
+	"github.com/mhsanaei/3x-ui/v3/internal/xray"

 	"github.com/op/go-logging"
 	"gorm.io/gorm"
 )

 // setupScaleDB initializes the DB for a scale benchmark on either Postgres
-// (XUI_DB_TYPE=postgres + XUI_DB_DSN) or SQLite (XUI_SCALE_TEST=1, temp file),
-// and registers cleanup. Skips the test when neither backend is configured.
+// (XUI_DB_TYPE=postgres + XUI_DB_DSN) or SQLite (XUI_SCALE_TEST=1, temp file;
+// XUI_SCALE_DB_PATH persists the DB for manual smoke runs), and registers
+// cleanup. Skips the test when neither backend is configured.
 func setupScaleDB(t *testing.T) {
 	t.Helper()
 	xuilogger.InitLogger(logging.ERROR)
@@ -31,7 +37,10 @@ func setupScaleDB(t *testing.T) {

 	switch strings.ToLower(strings.TrimSpace(os.Getenv("XUI_SCALE_TEST"))) {
 	case "1", "true", "yes":
-		dbPath := filepath.Join(t.TempDir(), "scale.db")
+		dbPath := strings.TrimSpace(os.Getenv("XUI_SCALE_DB_PATH"))
+		if dbPath == "" {
+			dbPath = filepath.Join(t.TempDir(), "scale.db")
+		}
 		if err := database.InitDB(dbPath); err != nil {
 			t.Fatalf("InitDB(sqlite): %v", err)
 		}
@@ -42,6 +51,144 @@ func setupScaleDB(t *testing.T) {
 	t.Skip("set XUI_SCALE_TEST=1 (sqlite) or XUI_DB_TYPE=postgres + XUI_DB_DSN (postgres) to run the scale benchmark")
 }

+// scaleSizes returns the default size ladder unless XUI_SCALE_SIZES overrides
+// it with a comma-separated list (e.g. "500000" or "10000,100000,500000").
+func scaleSizes(t *testing.T, def ...int) []int {
+	t.Helper()
+	raw := strings.TrimSpace(os.Getenv("XUI_SCALE_SIZES"))
+	if raw == "" {
+		return def
+	}
+	var out []int
+	for _, part := range strings.Split(raw, ",") {
+		part = strings.TrimSpace(part)
+		if part == "" {
+			continue
+		}
+		n, err := strconv.Atoi(part)
+		if err != nil || n <= 0 {
+			t.Fatalf("XUI_SCALE_SIZES: invalid size %q", part)
+		}
+		out = append(out, n)
+	}
+	if len(out) == 0 {
+		return def
+	}
+	return out
+}
+
+type scaleDataset struct {
+	inboundIds []int
+	tags       []string
+	emails     []string
+	perInbound [][]model.Client
+}
+
+// seedScaleDataset seeds n healthy clients (future expiry, unfilled quota)
+// spread across numInbounds inbounds, writing inbounds, clients,
+// client_inbounds and client_traffics directly in one transaction — orders of
+// magnitude faster than SyncInbound and one fsync instead of thousands.
+func seedScaleDataset(t *testing.T, n, numInbounds int) scaleDataset {
+	t.Helper()
+	db := database.GetDB()
+	resetScaleTables(t, db, "inbounds", "clients", "client_inbounds", "client_traffics")
+
+	clients := makeScaleClients(n)
+	exp := time.Now().AddDate(1, 0, 0).UnixMilli()
+	for i := range clients {
+		clients[i].ExpiryTime = exp
+		clients[i].TotalGB = 100 << 30
+	}
+
+	ds := scaleDataset{emails: emailsOf(clients)}
+	start := time.Now()
+	tx := db.Begin()
+	if tx.Error != nil {
+		t.Fatalf("begin seed tx: %v", tx.Error)
+	}
+	committed := false
+	defer func() {
+		if !committed {
+			tx.Rollback()
+		}
+	}()
+
+	per := n / numInbounds
+	for i := range numInbounds {
+		lo, hi := i*per, (i+1)*per
+		if i == numInbounds-1 {
+			hi = n
+		}
+		chunk := clients[lo:hi]
+		ib := &model.Inbound{
+			UserId:   1,
+			Tag:      fmt.Sprintf("scale-%d-%d", n, i),
+			Enable:   true,
+			Port:     41000 + i,
+			Protocol: model.VLESS,
+			Settings: clientsSettings(t, chunk),
+		}
+		if err := tx.Create(ib).Error; err != nil {
+			t.Fatalf("seed inbound %d: %v", i, err)
+		}
+
+		records := make([]*model.ClientRecord, len(chunk))
+		for j := range chunk {
+			records[j] = chunk[j].ToRecord()
+		}
+		if err := tx.CreateInBatches(records, 500).Error; err != nil {
+			t.Fatalf("seed clients %d: %v", i, err)
+		}
+
+		links := make([]model.ClientInbound, len(records))
+		for j := range records {
+			links[j] = model.ClientInbound{ClientId: records[j].Id, InboundId: ib.Id}
+		}
+		if err := tx.CreateInBatches(links, 1000).Error; err != nil {
+			t.Fatalf("seed client_inbounds %d: %v", i, err)
+		}
+
+		traffics := make([]xray.ClientTraffic, len(chunk))
+		for j := range chunk {
+			traffics[j] = xray.ClientTraffic{
+				InboundId:  ib.Id,
+				Email:      chunk[j].Email,
+				Enable:     true,
+				Total:      chunk[j].TotalGB,
+				ExpiryTime: chunk[j].ExpiryTime,
+			}
+		}
+		if err := tx.CreateInBatches(traffics, 1000).Error; err != nil {
+			t.Fatalf("seed client_traffics %d: %v", i, err)
+		}
+
+		ds.inboundIds = append(ds.inboundIds, ib.Id)
+		ds.tags = append(ds.tags, ib.Tag)
+		ds.perInbound = append(ds.perInbound, chunk)
+	}
+
+	if err := tx.Commit().Error; err != nil {
+		t.Fatalf("commit seed tx: %v", err)
+	}
+	committed = true
+	db.Exec("ANALYZE")
+	t.Logf("seeded N=%d across %d inbound(s) in %v", n, numInbounds, time.Since(start).Round(time.Millisecond))
+	return ds
+}
+
+// sampleEmails picks k evenly spaced emails so active clients span the id range.
+func sampleEmails(emails []string, k int) []string {
+	if k >= len(emails) {
+		return emails
+	}
+	out := make([]string, 0, k)
+	step := len(emails) / k
+	for i := 0; i < len(emails) && len(out) < k; i += step {
+		out = append(out, emails[i])
+	}
+	return out
+}
+
 // resetScaleTables empties the given tables between sub-sizes. Postgres uses a
 // single TRUNCATE ... CASCADE; SQLite deletes per table and clears the
 // autoincrement counters so ids restart like RESTART IDENTITY.