package inbound import ( "context" "fmt" "math" "strconv" "strings" "sync" "time" "github.com/gogf/gf/v2/container/gvar" "github.com/gogf/gf/v2/frame/g" "git.5iot.com/iotlcm-public/iotlcm-worker-svr/internal/model" ) const ( inboundUpstreamRateLimitConfigKey = "batch.inbound.upstreamRateLimit" defaultInboundUpstreamRateLimitMode = "redis" defaultInboundUpstreamRateLimitRedisGroup = "default" defaultInboundUpstreamRateLimitKeyPrefix = "iotlcm:worker:inbound:upstream_basic_info" defaultInboundUpstreamQPS = 20 defaultInboundUpstreamBurst = 1 defaultInboundUpstreamRequestCost = 1 defaultInboundChannelTempKeyCacheTTL = time.Minute inboundRateLimitFallbackLocal = "local" inboundRateLimitFallbackError = "error" inboundCarrierTelecom = "telecom" inboundCarrierMobile = "mobile" ) type upstreamBasicInfoRateLimiter interface { Wait(ctx context.Context, chanKey string, requestCount int) error } type redisEvaler interface { Eval(ctx context.Context, script string, numKeys int64, keys []string, args []any) (*gvar.Var, error) } type inboundUpstreamRateLimitConfig struct { Enabled *bool `json:"enabled"` Mode string `json:"mode"` RedisGroup string `json:"redisGroup"` KeyPrefix string `json:"keyPrefix"` RedisFailureFallback string `json:"redisFailureFallback"` QPS float64 `json:"qps"` Burst int `json:"burst"` RequestCost int `json:"requestCost"` Default inboundUpstreamRateLimitRule `json:"default"` Carriers map[string]inboundUpstreamRateLimitRule `json:"carriers"` Channels map[string]inboundUpstreamRateLimitRule `json:"channels"` } type inboundUpstreamRateLimitRule struct { Enabled *bool `json:"enabled"` QPS float64 `json:"qps"` Burst int `json:"burst"` RequestCost int `json:"requestCost"` } type inboundUpstreamRateLimitSetting struct { enabled bool mode string redisGroup string keyPrefix string redisFailureFallback string qps float64 burst int requestCost int carrier string tempKey string skipReason string } type inboundChannelMetadata struct { TempKey string `orm:"temp_key"` Operator int `orm:"operator"` } type inboundChannelTempKeyCacheEntry struct { metadata inboundChannelMetadata expiresAt time.Time } type configUpstreamBasicInfoRateLimiter struct { redis func(string) redisEvaler channelMetadata func(context.Context, string) (inboundChannelMetadata, error) channelTempKey func(context.Context, string) (string, error) local *localTokenBucketLimiter channelCache sync.Map redisLogTimes sync.Map skipLogTimes sync.Map } func newUpstreamBasicInfoRateLimiter() *configUpstreamBasicInfoRateLimiter { return &configUpstreamBasicInfoRateLimiter{ redis: func(name string) redisEvaler { return g.Redis(name) }, channelMetadata: queryInboundChannelMetadata, local: newLocalTokenBucketLimiter(), } } func (l *configUpstreamBasicInfoRateLimiter) Wait(ctx context.Context, chanKey string, requestCount int) error { if l == nil { return nil } cfg := loadInboundUpstreamRateLimitConfig(ctx) if !inboundUpstreamRateLimitGloballyEnabled(cfg) { return nil } metadata, err := l.lookupChannelMetadata(ctx, chanKey) if err != nil { return err } setting := resolveInboundUpstreamRateLimitFromChannelMetadata(cfg, chanKey, metadata.Operator, metadata.TempKey) if !setting.enabled || setting.qps <= 0 || setting.mode == "off" { if setting.skipReason != "" { l.logSkipped(ctx, chanKey, setting) } return nil } if requestCount <= 0 { requestCount = 1 } permits := requestCount * setting.requestCost key := inboundUpstreamRateLimitKey(setting.keyPrefix, chanKey, setting.carrier) for i := 0; i < permits; i++ { if err := l.waitOne(ctx, setting, key); err != nil { return err } } return nil } func (l *configUpstreamBasicInfoRateLimiter) waitOne(ctx context.Context, setting inboundUpstreamRateLimitSetting, key string) error { if setting.mode == "local" { if l.local == nil { return fmt.Errorf("local rate limiter is nil") } return l.local.Wait(ctx, key, setting.qps, setting.burst) } if l.redis == nil { return l.handleRedisError(ctx, setting, key, fmt.Errorf("redis client resolver is nil")) } redis := l.redis(setting.redisGroup) if redis == nil { return l.handleRedisError(ctx, setting, key, fmt.Errorf("redis client is nil")) } if err := (&redisTokenBucketLimiter{redis: redis}).Wait(ctx, key, setting.qps, setting.burst); err != nil { return l.handleRedisError(ctx, setting, key, err) } return nil } func (l *configUpstreamBasicInfoRateLimiter) handleRedisError(ctx context.Context, setting inboundUpstreamRateLimitSetting, key string, err error) error { if setting.redisFailureFallback == inboundRateLimitFallbackError { return err } if l.local == nil { return err } l.logRedisFallback(ctx, key, err) return l.local.Wait(ctx, key, setting.qps, setting.burst) } func (l *configUpstreamBasicInfoRateLimiter) logRedisFallback(ctx context.Context, key string, err error) { now := time.Now() if value, ok := l.redisLogTimes.Load(key); ok { if last, ok := value.(time.Time); ok && now.Sub(last) < 10*time.Second { return } } l.redisLogTimes.Store(key, now) g.Log().Errorf(ctx, "inbound upstream basic-info redis rate limiter failed, fallback local, key:%s, err:%+v", key, err) } func (l *configUpstreamBasicInfoRateLimiter) lookupChannelMetadata(ctx context.Context, chanKey string) (inboundChannelMetadata, error) { chanKey = strings.TrimSpace(chanKey) if chanKey == "" { return inboundChannelMetadata{}, nil } now := time.Now() if value, ok := l.channelCache.Load(chanKey); ok { if entry, ok := value.(inboundChannelTempKeyCacheEntry); ok && now.Before(entry.expiresAt) { return entry.metadata, nil } } if l.channelMetadata != nil { metadata, err := l.channelMetadata(ctx, chanKey) if err != nil { return inboundChannelMetadata{}, fmt.Errorf("query channel metadata failed, chanKey:%s: %w", chanKey, err) } metadata.TempKey = strings.TrimSpace(metadata.TempKey) l.channelCache.Store(chanKey, inboundChannelTempKeyCacheEntry{ metadata: metadata, expiresAt: now.Add(defaultInboundChannelTempKeyCacheTTL), }) return metadata, nil } if l.channelTempKey == nil { return inboundChannelMetadata{}, fmt.Errorf("channel metadata resolver is nil") } tempKey, err := l.channelTempKey(ctx, chanKey) if err != nil { return inboundChannelMetadata{}, fmt.Errorf("query channel template failed, chanKey:%s: %w", chanKey, err) } metadata := inboundChannelMetadata{TempKey: strings.TrimSpace(tempKey)} l.channelCache.Store(chanKey, inboundChannelTempKeyCacheEntry{ metadata: metadata, expiresAt: now.Add(defaultInboundChannelTempKeyCacheTTL), }) return metadata, nil } func (l *configUpstreamBasicInfoRateLimiter) logSkipped(ctx context.Context, chanKey string, setting inboundUpstreamRateLimitSetting) { key := strings.TrimSpace(chanKey) + ":" + setting.skipReason now := time.Now() if value, ok := l.skipLogTimes.Load(key); ok { if last, ok := value.(time.Time); ok && now.Sub(last) < 10*time.Second { return } } l.skipLogTimes.Store(key, now) g.Log().Warningf(ctx, "inbound upstream basic-info rate limit skipped, chanKey:%s, tempKey:%s, reason:%s", chanKey, setting.tempKey, setting.skipReason) } func queryInboundChannelMetadata(ctx context.Context, chanKey string) (inboundChannelMetadata, error) { var row inboundChannelMetadata if err := g.DB().Model("t_sim_channels").Ctx(ctx). Fields("temp_key,operator"). Where("chan_key", strings.TrimSpace(chanKey)). Limit(1). Scan(&row); err != nil { return inboundChannelMetadata{}, err } row.TempKey = strings.TrimSpace(row.TempKey) return row, nil } func queryInboundChannelTempKey(ctx context.Context, chanKey string) (string, error) { metadata, err := queryInboundChannelMetadata(ctx, chanKey) if err != nil { return "", err } return metadata.TempKey, nil } func resolveInboundUpstreamRateLimit(ctx context.Context, chanKey string) inboundUpstreamRateLimitSetting { return resolveInboundUpstreamRateLimitWithTempKey(ctx, chanKey, "") } func resolveInboundUpstreamRateLimitWithTempKey(ctx context.Context, chanKey, tempKey string) inboundUpstreamRateLimitSetting { cfg := loadInboundUpstreamRateLimitConfig(ctx) return resolveInboundUpstreamRateLimitFromConfig(cfg, chanKey, tempKey) } func resolveInboundUpstreamRateLimitWithChannelMetadata(ctx context.Context, chanKey string, operator int, tempKey string) inboundUpstreamRateLimitSetting { cfg := loadInboundUpstreamRateLimitConfig(ctx) return resolveInboundUpstreamRateLimitFromChannelMetadata(cfg, chanKey, operator, tempKey) } func loadInboundUpstreamRateLimitConfig(ctx context.Context) inboundUpstreamRateLimitConfig { cfg := defaultInboundUpstreamRateLimitConfig() value, err := g.Cfg().Get(ctx, inboundUpstreamRateLimitConfigKey) if err != nil { g.Log().Warningf(ctx, "load inbound upstream rate limit config failed: %+v", err) return normalizeInboundUpstreamRateLimitConfig(cfg) } if value != nil && !value.IsNil() { if err = value.Scan(&cfg); err != nil { g.Log().Warningf(ctx, "parse inbound upstream rate limit config failed: %+v", err) cfg = defaultInboundUpstreamRateLimitConfig() } } return normalizeInboundUpstreamRateLimitConfig(cfg) } func defaultInboundUpstreamRateLimitConfig() inboundUpstreamRateLimitConfig { return inboundUpstreamRateLimitConfig{ Mode: defaultInboundUpstreamRateLimitMode, RedisGroup: defaultInboundUpstreamRateLimitRedisGroup, KeyPrefix: defaultInboundUpstreamRateLimitKeyPrefix, RedisFailureFallback: inboundRateLimitFallbackError, Default: inboundUpstreamRateLimitRule{ Enabled: inboundBool(false), }, Carriers: map[string]inboundUpstreamRateLimitRule{ inboundCarrierTelecom: { Enabled: inboundBool(true), QPS: defaultInboundUpstreamQPS, Burst: defaultInboundUpstreamBurst, RequestCost: defaultInboundUpstreamRequestCost, }, inboundCarrierMobile: { Enabled: inboundBool(false), }, }, } } func normalizeInboundUpstreamRateLimitConfig(cfg inboundUpstreamRateLimitConfig) inboundUpstreamRateLimitConfig { cfg.Mode = normalizeInboundRateLimitMode(cfg.Mode) cfg.RedisGroup = firstNonEmpty(strings.TrimSpace(cfg.RedisGroup), defaultInboundUpstreamRateLimitRedisGroup) cfg.KeyPrefix = firstNonEmpty(strings.TrimSpace(cfg.KeyPrefix), defaultInboundUpstreamRateLimitKeyPrefix) cfg.RedisFailureFallback = normalizeInboundRateLimitFallback(cfg.RedisFailureFallback) if cfg.Default.Enabled == nil { cfg.Default.Enabled = inboundBool(false) } if len(cfg.Carriers) == 0 { cfg.Carriers = make(map[string]inboundUpstreamRateLimitRule, 2) } else { normalized := make(map[string]inboundUpstreamRateLimitRule, len(cfg.Carriers)) for key, rule := range cfg.Carriers { normalized[normalizeInboundCarrier(key)] = rule } cfg.Carriers = normalized } telecomRule, ok := cfg.Carriers[inboundCarrierTelecom] if !ok { telecomRule = inboundUpstreamRateLimitRule{} } if telecomRule.Enabled == nil { telecomRule.Enabled = inboundBool(true) } if telecomRule.QPS <= 0 { telecomRule.QPS = defaultInboundUpstreamQPS } if telecomRule.Burst <= 0 { telecomRule.Burst = defaultInboundUpstreamBurst } if telecomRule.RequestCost <= 0 { telecomRule.RequestCost = defaultInboundUpstreamRequestCost } cfg.Carriers[inboundCarrierTelecom] = telecomRule mobileRule, ok := cfg.Carriers[inboundCarrierMobile] if !ok { mobileRule = inboundUpstreamRateLimitRule{} } if mobileRule.Enabled == nil { mobileRule.Enabled = inboundBool(false) } cfg.Carriers[inboundCarrierMobile] = mobileRule return cfg } func resolveInboundUpstreamRateLimitFromConfig(cfg inboundUpstreamRateLimitConfig, chanKey, tempKey string) inboundUpstreamRateLimitSetting { return resolveInboundUpstreamRateLimitFromChannelMetadata(cfg, chanKey, 0, tempKey) } func resolveInboundUpstreamRateLimitFromChannelMetadata(cfg inboundUpstreamRateLimitConfig, chanKey string, operator int, tempKey string) inboundUpstreamRateLimitSetting { cfg = normalizeInboundUpstreamRateLimitConfig(cfg) setting := inboundUpstreamRateLimitSetting{ enabled: true, mode: cfg.Mode, redisGroup: cfg.RedisGroup, keyPrefix: cfg.KeyPrefix, redisFailureFallback: cfg.RedisFailureFallback, qps: defaultInboundUpstreamQPS, burst: defaultInboundUpstreamBurst, requestCost: defaultInboundUpstreamRequestCost, tempKey: strings.TrimSpace(tempKey), } if cfg.Enabled != nil { setting.enabled = *cfg.Enabled } applyInboundRateLimitRule(&setting, inboundUpstreamRateLimitRule{ QPS: cfg.QPS, Burst: cfg.Burst, RequestCost: cfg.RequestCost, }) applyInboundRateLimitRule(&setting, cfg.Default) carrier, known := inboundCarrierForOperator(operator) if model.IsTelecomTemplate(tempKey) { carrier, known = inboundCarrierTelecom, true } if known { setting.carrier = carrier if rule, ok := inboundCarrierRateLimitRule(cfg.Carriers, carrier); ok { applyInboundRateLimitRule(&setting, rule) } } else if operator == 0 { setting.skipReason = "missing_channel_operator" } else { setting.skipReason = "unsupported_channel_operator" } if rule, ok := inboundChannelRateLimitRule(cfg.Channels, chanKey); ok { applyInboundRateLimitRule(&setting, rule) } return normalizeInboundRateLimitSetting(setting) } func inboundUpstreamRateLimitGloballyEnabled(cfg inboundUpstreamRateLimitConfig) bool { cfg = normalizeInboundUpstreamRateLimitConfig(cfg) if cfg.Mode == "off" { return false } if cfg.Enabled == nil { return true } return *cfg.Enabled } func applyInboundRateLimitRule(setting *inboundUpstreamRateLimitSetting, rule inboundUpstreamRateLimitRule) { if setting == nil { return } if rule.Enabled != nil { setting.enabled = *rule.Enabled } if rule.QPS > 0 { setting.qps = rule.QPS } if rule.Burst > 0 { setting.burst = rule.Burst } if rule.RequestCost > 0 { setting.requestCost = rule.RequestCost } } func inboundChannelRateLimitRule(configs map[string]inboundUpstreamRateLimitRule, chanKey string) (inboundUpstreamRateLimitRule, bool) { chanKey = strings.TrimSpace(chanKey) if chanKey == "" || len(configs) == 0 { return inboundUpstreamRateLimitRule{}, false } if cfg, ok := configs[chanKey]; ok { return cfg, true } lower := strings.ToLower(chanKey) for key, cfg := range configs { if strings.ToLower(strings.TrimSpace(key)) == lower { return cfg, true } } return inboundUpstreamRateLimitRule{}, false } func inboundCarrierRateLimitRule(configs map[string]inboundUpstreamRateLimitRule, carrier string) (inboundUpstreamRateLimitRule, bool) { carrier = normalizeInboundCarrier(carrier) if carrier == "" || len(configs) == 0 { return inboundUpstreamRateLimitRule{}, false } if cfg, ok := configs[carrier]; ok { return cfg, true } for key, cfg := range configs { if normalizeInboundCarrier(key) == carrier { return cfg, true } } return inboundUpstreamRateLimitRule{}, false } func inboundCarrierForOperator(operator int) (string, bool) { switch operator { case 1: return inboundCarrierMobile, true case 2: return inboundCarrierTelecom, true default: return "", false } } func normalizeInboundCarrier(value string) string { switch strings.ToLower(strings.TrimSpace(value)) { case "", "unknown": return "" case inboundCarrierTelecom, "dx", "ctcc", "china_telecom", "chinatelecom": return inboundCarrierTelecom case inboundCarrierMobile, "yd", "cmcc", "china_mobile", "chinamobile": return inboundCarrierMobile default: return strings.ToLower(strings.TrimSpace(value)) } } func normalizeInboundRateLimitSetting(setting inboundUpstreamRateLimitSetting) inboundUpstreamRateLimitSetting { if setting.qps <= 0 || math.IsNaN(setting.qps) || math.IsInf(setting.qps, 0) { setting.qps = defaultInboundUpstreamQPS } if setting.burst <= 0 { setting.burst = defaultInboundUpstreamBurst } if setting.requestCost <= 0 { setting.requestCost = defaultInboundUpstreamRequestCost } return setting } func normalizeInboundRateLimitMode(mode string) string { switch strings.ToLower(strings.TrimSpace(mode)) { case "", "redis", "distributed": return defaultInboundUpstreamRateLimitMode case "local", "memory": return "local" case "off", "disable", "disabled", "none": return "off" default: return defaultInboundUpstreamRateLimitMode } } func normalizeInboundRateLimitFallback(value string) string { switch strings.ToLower(strings.TrimSpace(value)) { case "", inboundRateLimitFallbackError: return inboundRateLimitFallbackError case inboundRateLimitFallbackLocal: return inboundRateLimitFallbackLocal default: return inboundRateLimitFallbackError } } func inboundUpstreamRateLimitKey(prefix, chanKey, carrier string) string { prefix = strings.TrimRight(strings.TrimSpace(prefix), ":") if prefix == "" { prefix = defaultInboundUpstreamRateLimitKeyPrefix } chanKey = strings.TrimLeft(strings.TrimSpace(chanKey), ":") carrier = strings.TrimLeft(normalizeInboundCarrier(carrier), ":") if carrier != "" { return prefix + ":carrier:" + carrier + ":chan:" + chanKey } return prefix + ":chan:" + chanKey } func inboundBool(value bool) *bool { return &value } type localTokenBucketLimiter struct { mu sync.Mutex buckets map[string]*tokenBucket } func newLocalTokenBucketLimiter() *localTokenBucketLimiter { return &localTokenBucketLimiter{buckets: make(map[string]*tokenBucket)} } func (l *localTokenBucketLimiter) Wait(ctx context.Context, key string, qps float64, burst int) error { if qps <= 0 { return nil } bucket := l.bucket(key) delay := bucket.reserve(time.Now(), qps, burst) if delay <= 0 { return nil } timer := time.NewTimer(delay) defer timer.Stop() select { case <-ctx.Done(): return ctx.Err() case <-timer.C: return nil } } func (l *localTokenBucketLimiter) bucket(key string) *tokenBucket { l.mu.Lock() defer l.mu.Unlock() bucket := l.buckets[key] if bucket == nil { bucket = &tokenBucket{} l.buckets[key] = bucket } return bucket } type tokenBucket struct { mu sync.Mutex tokens float64 last time.Time } func (b *tokenBucket) reserve(now time.Time, qps float64, burst int) time.Duration { if qps <= 0 { return 0 } capacity := float64(burst) if capacity <= 0 { capacity = defaultInboundUpstreamBurst } b.mu.Lock() defer b.mu.Unlock() if b.last.IsZero() { b.last = now b.tokens = capacity } if now.After(b.last) { elapsed := now.Sub(b.last).Seconds() b.tokens += elapsed * qps if b.tokens > capacity { b.tokens = capacity } b.last = now } if b.tokens >= 1 { b.tokens-- return 0 } waitSeconds := (1 - b.tokens) / qps wait := time.Duration(waitSeconds * float64(time.Second)) if wait <= 0 { wait = time.Nanosecond } b.tokens = 0 b.last = b.last.Add(wait) return b.last.Sub(now) } type redisTokenBucketLimiter struct { redis redisEvaler } const redisTokenBucketScript = ` local now = tonumber(ARGV[1]) local qps = tonumber(ARGV[2]) local burst = tonumber(ARGV[3]) local ttl = tonumber(ARGV[4]) if qps == nil or qps <= 0 then return {1, 0} end if burst == nil or burst <= 0 then burst = 1 end if ttl == nil or ttl <= 0 then ttl = 60000 end local values = redis.call("HMGET", KEYS[1], "tokens", "updated_at") local tokens = tonumber(values[1]) local updated_at = tonumber(values[2]) local capacity = burst if tokens == nil or updated_at == nil then tokens = capacity updated_at = now end if tokens > capacity then tokens = capacity end if now > updated_at then local elapsed = (now - updated_at) / 1000 tokens = math.min(capacity, tokens + elapsed * qps) updated_at = now end local wait = 0 if tokens >= 1 then tokens = tokens - 1 else wait = math.ceil((1 - tokens) / qps * 1000) if wait < 1 then wait = 1 end tokens = 0 updated_at = updated_at + wait wait = updated_at - now if wait < 0 then wait = 0 end end redis.call("HSET", KEYS[1], "tokens", tostring(tokens), "updated_at", tostring(updated_at)) redis.call("PEXPIRE", KEYS[1], ttl) return {1, wait} ` func (l *redisTokenBucketLimiter) Wait(ctx context.Context, key string, qps float64, burst int) error { wait, err := l.reserve(ctx, key, qps, burst) if err != nil { return err } if wait <= 0 { return nil } timer := time.NewTimer(wait) defer timer.Stop() select { case <-ctx.Done(): return ctx.Err() case <-timer.C: return nil } } func (l *redisTokenBucketLimiter) reserve(ctx context.Context, key string, qps float64, burst int) (time.Duration, error) { if l == nil || l.redis == nil || qps <= 0 { return 0, nil } if burst <= 0 { burst = defaultInboundUpstreamBurst } reply, err := l.redis.Eval(ctx, redisTokenBucketScript, 1, []string{key}, []any{ time.Now().UnixMilli(), strconv.FormatFloat(qps, 'f', -1, 64), burst, rateLimitRedisKeyTTLMillis(qps, burst), }) if err != nil { return 0, err } values := reply.Int64s() if len(values) < 2 || values[0] != 1 { return 0, fmt.Errorf("unexpected redis rate limiter reply: %v", reply.Val()) } waitMillis := values[1] if waitMillis <= 0 { return 0, nil } return time.Duration(waitMillis) * time.Millisecond, nil } func rateLimitRedisKeyTTLMillis(qps float64, burst int) int64 { if qps <= 0 { return int64(time.Minute / time.Millisecond) } if burst <= 0 { burst = defaultInboundUpstreamBurst } idleMillis := int64(math.Ceil(float64(burst) / qps * 1000 * 2)) minMillis := int64(time.Minute / time.Millisecond) if idleMillis < minMillis { return minMillis } return idleMillis }