chore: 初始开源发布
iotlcm-worker-svr:异步执行中枢(RabbitMQ 消费者 + 全量同步编排)。 - Go module:iotlcm-worker-svr → git.5iot.com/iotlcm-public/iotlcm-worker-svr(214 文件) - 凭据与内部标识脱敏:MySQL/Redis 凭据、私网与生产 IP、内网域名、 库名 xiujie_admin → iotlcm、哨兵名占位 - 159 处真实手机号 → 合法 11 位编造号(保留号段与后 4 位,维持互异, 不破坏「校验 MSISDN 格式」等测试断言) - 移除内部资产:CI 触发器、Jenkins 触发测试文档、预演门禁与镜像 digest 记录 - Dockerfile 改用公共基础镜像(golang:1.25-alpine → alpine:3.23), 移除写死的 git 凭据 - README 重写(架构/全量同步模型/配置/已知失败),补 Apache-2.0 LICENSE,LF 行尾 + .gitattributes 已知:3 个测试包在源仓同一提交上同样失败(client/upstream、 service/allocate、service/inbound),属既有缺陷,本仓未引入亦未修复。
This commit is contained in:
@@ -0,0 +1,738 @@
|
||||
package inbound
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"math"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"git.5iot.com/iotlcm-public/iotlcm-worker-svr/internal/model"
|
||||
"github.com/gogf/gf/v2/container/gvar"
|
||||
"github.com/gogf/gf/v2/frame/g"
|
||||
)
|
||||
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user