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])) }