hi-server/pkg/speedlimit/calculator.go

package speedlimit

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"time"

	"github.com/redis/go-redis/v9"
	"gorm.io/gorm"
)

// TrafficLimitRule represents a dynamic speed throttling rule.
type TrafficLimitRule struct {
	StatType     string `json:"stat_type"`
	StatValue    int64  `json:"stat_value"`
	TrafficUsage int64  `json:"traffic_usage"`
	SpeedLimit   int64  `json:"speed_limit"`
}

// ThrottleResult contains the computed speed limit status for a user subscription.
type ThrottleResult struct {
	BaseSpeed      int64   `json:"base_speed"`      // Plan base speed limit (Mbps, 0=unlimited)
	EffectiveSpeed int64   `json:"effective_speed"` // Current effective speed limit (Mbps)
	IsThrottled    bool    `json:"is_throttled"`    // Whether the user is currently throttled
	ThrottleRule   string  `json:"throttle_rule"`   // Description of the matched rule (empty if not throttled)
	UsedTrafficGB  float64 `json:"used_traffic_gb"` // Traffic used in the matched rule's window (GB)
	ThrottleStart  int64   `json:"throttle_start"`  // Window start Unix timestamp (seconds), 0 if not throttled
	ThrottleEnd    int64   `json:"throttle_end"`    // Window end Unix timestamp (seconds), 0 if not throttled
}

// CalculateWithCache computes the effective speed limit with a short Redis cache.
// It is intended for hot read paths such as node user-list pulls where many nodes
// can ask for the same subscription limits in a short period.
func CalculateWithCache(ctx context.Context, cache *redis.Client, db *gorm.DB, userId, subscribeId, baseSpeedLimit int64, trafficLimitJSON string, ttl time.Duration) *ThrottleResult {
	if cache == nil || ttl <= 0 || trafficLimitJSON == "" {
		return Calculate(ctx, db, userId, subscribeId, baseSpeedLimit, trafficLimitJSON)
	}

	key := cacheKey(userId, subscribeId, baseSpeedLimit, trafficLimitJSON)
	if cached, err := cache.Get(ctx, key).Result(); err == nil && cached != "" {
		var result ThrottleResult
		if err := json.Unmarshal([]byte(cached), &result); err == nil {
			return &result
		}
	}

	result := Calculate(ctx, db, userId, subscribeId, baseSpeedLimit, trafficLimitJSON)
	if payload, err := json.Marshal(result); err == nil {
		_ = cache.Set(ctx, key, string(payload), ttl).Err()
	}
	return result
}

// ClearCache removes a cached speed-limit calculation for a user subscription.
func ClearCache(ctx context.Context, cache *redis.Client, userId, subscribeId, baseSpeedLimit int64, trafficLimitJSON string) error {
	if cache == nil || trafficLimitJSON == "" {
		return nil
	}
	return cache.Del(ctx, cacheKey(userId, subscribeId, baseSpeedLimit, trafficLimitJSON)).Err()
}

// Calculate computes the effective speed limit for a user subscription,
// considering traffic-based throttling rules.
func Calculate(ctx context.Context, db *gorm.DB, userId, subscribeId, baseSpeedLimit int64, trafficLimitJSON string) *ThrottleResult {
	result := &ThrottleResult{
		BaseSpeed:      baseSpeedLimit,
		EffectiveSpeed: baseSpeedLimit,
	}

	if trafficLimitJSON == "" {
		return result
	}

	var rules []TrafficLimitRule
	if err := json.Unmarshal([]byte(trafficLimitJSON), &rules); err != nil {
		return result
	}

	if len(rules) == 0 {
		return result
	}

	now := time.Now()
	for _, rule := range rules {
		var startTime time.Time

		switch rule.StatType {
		case "hour":
			if rule.StatValue <= 0 {
				continue
			}
			startTime = now.Add(-time.Duration(rule.StatValue) * time.Hour)
		case "day":
			if rule.StatValue <= 0 {
				continue
			}
			startTime = now.AddDate(0, 0, -int(rule.StatValue))
		default:
			continue
		}

		var usedTraffic struct {
			Upload   int64
			Download int64
		}
		err := db.WithContext(ctx).
			Table("traffic_log").
			Select("COALESCE(SUM(upload), 0) as upload, COALESCE(SUM(download), 0) as download").
			Where("user_id = ? AND subscribe_id = ? AND timestamp >= ? AND timestamp < ?",
				userId, subscribeId, startTime, now).
			Scan(&usedTraffic).Error

		if err != nil {
			continue
		}

		usedGB := float64(usedTraffic.Upload+usedTraffic.Download) / (1024 * 1024 * 1024)

		if usedGB >= float64(rule.TrafficUsage) {
			if rule.SpeedLimit > 0 {
				if result.EffectiveSpeed == 0 || rule.SpeedLimit < result.EffectiveSpeed {
					result.EffectiveSpeed = rule.SpeedLimit
					result.IsThrottled = true
					result.UsedTrafficGB = usedGB
					result.ThrottleStart = startTime.Unix()
					result.ThrottleEnd = now.Unix()

					statLabel := "小时"
					if rule.StatType == "day" {
						statLabel = "天"
					}
					result.ThrottleRule = fmt.Sprintf("%d%s内超%dGB，限速%dMbps",
						rule.StatValue, statLabel, rule.TrafficUsage, rule.SpeedLimit)
				}
			}
		}
	}

	return result
}

func cacheKey(userId, subscribeId, baseSpeedLimit int64, trafficLimitJSON string) string {
	sum := sha256.Sum256([]byte(trafficLimitJSON))
	return fmt.Sprintf("speedlimit:%d:%d:%d:%s", userId, subscribeId, baseSpeedLimit, hex.EncodeToString(sum[:8]))
}