修复(#12): 优化用户列表限速计算性能 (#14)

Co-authored-by: multica-agent <github@multica.ai>
2026-06-07 23:29:15 -07:00
parent b98d718f3c
commit 24caf58987
2 changed files with 517 additions and 17 deletions
@@ -1,11 +1,24 @@
 package server

 import (
+	"context"
 	"fmt"
+	"net/http"
+	"regexp"
 	"slices"
 	"testing"
+	"time"

+	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/gin-gonic/gin"
 	"github.com/perfect-panel/server/internal/model/node"
+	"github.com/perfect-panel/server/internal/model/subscribe"
+	"github.com/perfect-panel/server/internal/model/user"
+	"github.com/perfect-panel/server/internal/svc"
+	"github.com/perfect-panel/server/pkg/logger"
+	"github.com/perfect-panel/server/pkg/speedlimit"
+	"gorm.io/driver/mysql"
+	"gorm.io/gorm"
 )

 // TestNormalizeServerUserListProtocol 验证客户端 hysteria2 兼容字段被映射回
@@ -76,3 +89,252 @@ func TestServerUserListCacheKey_ShapeMatchesDelPath(t *testing.T) {
 		t.Fatalf("write-path key %q not present in Del-path enumeration %v", writeShape, delKeys)
 	}
 }
+
+func TestServerUserSpeedLimitInputs(t *testing.T) {
+	planTrafficLimit := `[{"stat_type":"day","stat_value":1,"traffic_usage":10,"speed_limit":5}]`
+	userTrafficLimit := `[{"stat_type":"day","stat_value":1,"traffic_usage":1,"speed_limit":3}]`
+
+	cases := []struct {
+		name             string
+		sub              *subscribe.Subscribe
+		userSub          *user.Subscribe
+		wantSpeedLimit   int64
+		wantTrafficLimit string
+	}{
+		{
+			name: "未设置用户覆盖时返回套餐基础限速",
+			sub: &subscribe.Subscribe{
+				SpeedLimit:   100,
+				TrafficLimit: planTrafficLimit,
+			},
+			userSub:          &user.Subscribe{},
+			wantSpeedLimit:   100,
+			wantTrafficLimit: planTrafficLimit,
+		},
+		{
+			name: "用户 speed_limit 覆盖套餐基础限速",
+			sub: &subscribe.Subscribe{
+				SpeedLimit:   100,
+				TrafficLimit: planTrafficLimit,
+			},
+			userSub: &user.Subscribe{
+				SpeedLimit: 30,
+			},
+			wantSpeedLimit:   30,
+			wantTrafficLimit: planTrafficLimit,
+		},
+		{
+			name: "用户 traffic_limit 覆盖套餐规则",
+			sub: &subscribe.Subscribe{
+				SpeedLimit:   100,
+				TrafficLimit: planTrafficLimit,
+			},
+			userSub: &user.Subscribe{
+				TrafficLimit: &userTrafficLimit,
+			},
+			wantSpeedLimit:   100,
+			wantTrafficLimit: userTrafficLimit,
+		},
+		{
+			name: "用户 traffic_limit 为空时保留套餐规则",
+			sub: &subscribe.Subscribe{
+				SpeedLimit:   100,
+				TrafficLimit: planTrafficLimit,
+			},
+			userSub: &user.Subscribe{
+				TrafficLimit: ptrString(""),
+			},
+			wantSpeedLimit:   100,
+			wantTrafficLimit: planTrafficLimit,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			gotSpeedLimit, gotTrafficLimit := serverUserSpeedLimitInputs(tc.sub, tc.userSub)
+			if gotSpeedLimit != tc.wantSpeedLimit {
+				t.Fatalf("speedLimit = %d, want %d", gotSpeedLimit, tc.wantSpeedLimit)
+			}
+			if gotTrafficLimit != tc.wantTrafficLimit {
+				t.Fatalf("trafficLimit = %q, want %q", gotTrafficLimit, tc.wantTrafficLimit)
+			}
+		})
+	}
+}
+
+func TestCalculateServerUserSpeedLimits_BatchTrafficRules(t *testing.T) {
+	db, mock, cleanup := newServerUserListTestDB(t)
+	defer cleanup()
+
+	ctx, _ := gin.CreateTestContext(nil)
+	ctx.Request = httptestNewRequest()
+	logic := &GetServerUserListLogic{
+		Logger: logger.WithContext(ctx.Request.Context()),
+		ctx:    ctx,
+		svcCtx: &svc.ServiceContext{DB: db},
+	}
+
+	trafficLimit := `[{"stat_type":"day","stat_value":1,"traffic_usage":10,"speed_limit":5}]`
+	candidates := map[int64]serverUserSpeedLimitCandidate{
+		101: {userID: 1001, userSubscribeID: 101, baseSpeed: 0, trafficLimit: trafficLimit},
+		102: {userID: 1002, userSubscribeID: 102, baseSpeed: 0, trafficLimit: trafficLimit},
+		103: {userID: 1003, userSubscribeID: 103, baseSpeed: 20, trafficLimit: trafficLimit},
+	}
+
+	mock.ExpectQuery("FROM `traffic_log`").
+		WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg()).
+		WillReturnRows(sqlmock.NewRows([]string{"user_id", "user_subscribe_id", "upload", "download"}).
+			AddRow(1001, 101, gb(4), gb(5)).
+			AddRow(1002, 102, gb(4), gb(6)).
+			AddRow(1003, 103, gb(12), int64(0)))
+
+	got := logic.calculateServerUserSpeedLimits(candidates)
+	assertSpeedLimit(t, got, 101, 0)
+	assertSpeedLimit(t, got, 102, 5)
+	assertSpeedLimit(t, got, 103, 5)
+	assertServerUserListExpectations(t, mock)
+}
+
+func TestCalculateServerUserSpeedLimits_UserTrafficLimitOverride(t *testing.T) {
+	db, mock, cleanup := newServerUserListTestDB(t)
+	defer cleanup()
+
+	ctx, _ := gin.CreateTestContext(nil)
+	ctx.Request = httptestNewRequest()
+	logic := &GetServerUserListLogic{
+		Logger: logger.WithContext(ctx.Request.Context()),
+		ctx:    ctx,
+		svcCtx: &svc.ServiceContext{DB: db},
+	}
+
+	candidates := map[int64]serverUserSpeedLimitCandidate{
+		201: {
+			userID:          2001,
+			userSubscribeID: 201,
+			baseSpeed:       50,
+			trafficLimit:    `[{"stat_type":"day","stat_value":1,"traffic_usage":2,"speed_limit":3}]`,
+		},
+		202: {
+			userID:          2002,
+			userSubscribeID: 202,
+			baseSpeed:       50,
+			trafficLimit:    "",
+		},
+	}
+
+	mock.ExpectQuery("FROM `traffic_log`").
+		WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg()).
+		WillReturnRows(sqlmock.NewRows([]string{"user_id", "user_subscribe_id", "upload", "download"}).
+			AddRow(2001, 201, gb(2), int64(0)))
+
+	got := logic.calculateServerUserSpeedLimits(candidates)
+	assertSpeedLimit(t, got, 201, 3)
+	assertSpeedLimit(t, got, 202, 50)
+	assertServerUserListExpectations(t, mock)
+}
+
+func TestServerUserTrafficRuleWindow(t *testing.T) {
+	now := time.Date(2026, 6, 8, 12, 0, 0, 0, time.UTC)
+
+	cases := []struct {
+		name string
+		rule speedlimit.TrafficLimitRule
+		want time.Time
+		ok   bool
+	}{
+		{
+			name: "hour window",
+			rule: speedlimit.TrafficLimitRule{StatType: "hour", StatValue: 2},
+			want: now.Add(-2 * time.Hour),
+			ok:   true,
+		},
+		{
+			name: "day window",
+			rule: speedlimit.TrafficLimitRule{StatType: "day", StatValue: 1},
+			want: now.AddDate(0, 0, -1),
+			ok:   true,
+		},
+		{
+			name: "invalid stat type",
+			rule: speedlimit.TrafficLimitRule{StatType: "week", StatValue: 1},
+			ok:   false,
+		},
+		{
+			name: "invalid stat value",
+			rule: speedlimit.TrafficLimitRule{StatType: "day", StatValue: 0},
+			ok:   false,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			got, ok := serverUserTrafficRuleWindow(tc.rule, now)
+			if ok != tc.ok {
+				t.Fatalf("ok = %v, want %v", ok, tc.ok)
+			}
+			if !ok {
+				return
+			}
+			if !got.start.Equal(tc.want) || !got.end.Equal(now) {
+				t.Fatalf("window = [%s, %s), want [%s, %s)", got.start, got.end, tc.want, now)
+			}
+		})
+	}
+}
+
+func newServerUserListTestDB(t *testing.T) (*gorm.DB, sqlmock.Sqlmock, func()) {
+	t.Helper()
+
+	sqlDB, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherFunc(func(expectedSQL, actualSQL string) error {
+		if matched, _ := regexp.MatchString(expectedSQL, actualSQL); matched {
+			return nil
+		}
+		t.Fatalf("unexpected SQL\nexpected pattern: %s\nactual: %s", expectedSQL, actualSQL)
+		return nil
+	})))
+	if err != nil {
+		t.Fatalf("create sqlmock: %v", err)
+	}
+
+	db, err := gorm.Open(mysql.New(mysql.Config{
+		Conn:                      sqlDB,
+		SkipInitializeWithVersion: true,
+	}), &gorm.Config{})
+	if err != nil {
+		_ = sqlDB.Close()
+		t.Fatalf("open gorm db: %v", err)
+	}
+
+	return db, mock, func() {
+		_ = sqlDB.Close()
+	}
+}
+
+func assertServerUserListExpectations(t *testing.T, mock sqlmock.Sqlmock) {
+	t.Helper()
+
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Fatalf("unmet sql expectations: %v", err)
+	}
+}
+
+func assertSpeedLimit(t *testing.T, got map[int64]int64, userSubscribeID int64, want int64) {
+	t.Helper()
+
+	if got[userSubscribeID] != want {
+		t.Fatalf("speed limit for user_subscribe_id=%d = %d, want %d", userSubscribeID, got[userSubscribeID], want)
+	}
+}
+
+func gb(value int64) int64 {
+	return value * 1024 * 1024 * 1024
+}
+
+func ptrString(value string) *string {
+	return &value
+}
+
+func httptestNewRequest() *http.Request {
+	req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, "/v1/server/user", nil)
+	return req
+}