fix: Optimization of Rate Limiting Logic for Cluster, AI Token and WASM Plugin (#2997)

This commit is contained in:
韩贤涛
2025-10-15 17:24:42 +08:00
committed by GitHub
parent b026455701
commit 1f301be851
4 changed files with 109 additions and 53 deletions

View File

@@ -46,17 +46,30 @@ func init() {
const (
// RedisKeyPrefix 集群限流插件在 Redis 中 key 的统一前缀
RedisKeyPrefix = "higress-cluster-key-rate-limit"
// ClusterGlobalRateLimitFormat 全局限流模式 redis key 为 RedisKeyPrefix:限流规则名称:global_threshold:时间窗口:窗口内限流数
ClusterGlobalRateLimitFormat = RedisKeyPrefix + ":%s:global_threshold:%d:%d"
// ClusterRateLimitFormat 规则限流模式 redis key 为 RedisKeyPrefix:限流规则名称:限流类型:时间窗口:窗口内限流数:限流key名称:限流key对应的实际值
ClusterRateLimitFormat = RedisKeyPrefix + ":%s:%s:%d:%d:%s:%s"
// ClusterGlobalRateLimitFormat 全局限流模式 redis key 为 RedisKeyPrefix:限流规则名称:global_threshold:时间窗口
ClusterGlobalRateLimitFormat = RedisKeyPrefix + ":%s:global_threshold:%d"
// ClusterRateLimitFormat 规则限流模式 redis key 为 RedisKeyPrefix:限流规则名称:限流类型:时间窗口:限流key名称:限流key对应的实际值
ClusterRateLimitFormat = RedisKeyPrefix + ":%s:%s:%d:%s:%s"
FixedWindowScript = `
local ttl = redis.call('ttl', KEYS[1])
if ttl < 0 then
redis.call('set', KEYS[1], ARGV[1] - 1, 'EX', ARGV[2])
return {ARGV[1], ARGV[1] - 1, ARGV[2]}
end
return {ARGV[1], redis.call('incrby', KEYS[1], -1), ttl}
local key = KEYS[1]
local threshold = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local current = tonumber(redis.call('get', key) or "0")
-- 只有超过阈值时才停止累加,达到阈值时仍允许(此时是最后一次允许)
if current > threshold then
return {threshold, current, redis.call('ttl', key)}
end
-- 计数未超过阈值,执行累加
current = redis.call('incr', key)
-- 第一次累加时设置过期时间
if current == 1 then
redis.call('expire', key, window)
end
return {threshold, current, redis.call('ttl', key)}
`
LimitContextKey = "LimitContext" // 限流上下文信息
@@ -92,7 +105,7 @@ func onHttpRequestHeaders(ctx wrapper.HttpContext, cfg config.ClusterKeyRateLimi
if cfg.GlobalThreshold != nil {
// 全局限流模式
limitKey = fmt.Sprintf(ClusterGlobalRateLimitFormat, cfg.RuleName, cfg.GlobalThreshold.TimeWindow, cfg.GlobalThreshold.Count)
limitKey = fmt.Sprintf(ClusterGlobalRateLimitFormat, cfg.RuleName, cfg.GlobalThreshold.TimeWindow)
count = cfg.GlobalThreshold.Count
timeWindow = cfg.GlobalThreshold.TimeWindow
} else {
@@ -103,7 +116,7 @@ func onHttpRequestHeaders(ctx wrapper.HttpContext, cfg config.ClusterKeyRateLimi
return types.ActionContinue
}
limitKey = fmt.Sprintf(ClusterRateLimitFormat, cfg.RuleName, ruleItem.LimitType, configItem.TimeWindow, configItem.Count, ruleItem.Key, val)
limitKey = fmt.Sprintf(ClusterRateLimitFormat, cfg.RuleName, ruleItem.LimitType, configItem.TimeWindow, ruleItem.Key, val)
count = configItem.Count
timeWindow = configItem.TimeWindow
}
@@ -118,12 +131,15 @@ func onHttpRequestHeaders(ctx wrapper.HttpContext, cfg config.ClusterKeyRateLimi
proxywasm.ResumeHttpRequest()
return
}
// 获取限流结果
threshold, current, ttl := resultArray[0].Integer(), resultArray[1].Integer(), resultArray[2].Integer()
context := LimitContext{
count: resultArray[0].Integer(),
remaining: resultArray[1].Integer(),
reset: resultArray[2].Integer(),
count: threshold,
remaining: threshold - current,
reset: ttl,
}
if context.remaining < 0 {
if current > threshold {
// 触发限流
rejected(cfg, context)
} else {