jiazhizhong/higress
1
0
Fork 0
mirror of https://github.com/alibaba/higress.git synced 2026-06-06 03:07:28 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: ai-token-ratelimit support setting global rate limit thresholds for routes​ (#2667)

Browse Source
This commit is contained in:
韩贤涛
2025-07-28 08:14:35 +08:00
committed by GitHub
parent e6e4193679
commit 6a1557f6ac
14 changed files with 981 additions and 564 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
plugins/wasm-go/extensions/cluster-key-rate-limit/config/config.go
View File

@@ -3,12 +3,11 @@ package config
import (
"errors"
"fmt"
re "regexp"
"strings"
"cluster-key-rate-limit/util"
re "regexp"
"github.com/higress-group/wasm-go/pkg/log"
"github.com/higress-group/wasm-go/pkg/wrapper"
"github.com/tidwall/gjson"
"github.com/zmap/go-iptree/iptree"
@@ -191,11 +190,25 @@ func initLimitRule(json gjson.Result, config *ClusterKeyRateLimitConfig) error {
}
var ruleItems []LimitRuleItem
// 用于记录已出现的LimitType和Key的组合
seenLimitRules := make(map[string]bool)
for _, item := range items {
ruleItem, err := parseLimitRuleItem(item)
if err != nil {
return fmt.Errorf("failed to parse rule_item in rule_items: %w", err)
}
// 构造LimitType和Key的唯一标识
ruleKey := string(ruleItem.LimitType) + ":" + ruleItem.Key
// 检查是否有重复的LimitType和Key组合
if seenLimitRules[ruleKey] {
log.Warnf("duplicate rule found: %s='%s' in rule_items", ruleItem.LimitType, ruleItem.Key)
} else {
seenLimitRules[ruleKey] = true
}
ruleItems = append(ruleItems, *ruleItem)
}
config.RuleItems = ruleItems
@@ -205,9 +218,13 @@ func initLimitRule(json gjson.Result, config *ClusterKeyRateLimitConfig) error {
func parseGlobalThreshold(item gjson.Result) (*GlobalThreshold, error) {
for timeWindowKey, duration := range timeWindows {
q := item.Get(timeWindowKey)
if q.Exists() && q.Int() > 0 {
if q.Exists() {
count := q.Int()
if count <= 0 {
return nil, fmt.Errorf("'%s' must be a positive integer, got %d", timeWindowKey, count)
}
return &GlobalThreshold{
Count: q.Int(),
Count: count,
TimeWindow: duration,
}, nil
}
@@ -276,7 +293,7 @@ func parseLimitRuleItem(item gjson.Result) (*LimitRuleItem, error) {
// 初始化configItems
err := initConfigItems(item, &ruleItem)
if err != nil {
return nil, fmt.Errorf("failed to init config items: %w", err)
return nil, err
}
return &ruleItem, nil
@@ -344,13 +361,17 @@ func initConfigItems(json gjson.Result, rule *LimitRuleItem) error {
func createConfigItemFromRate(item gjson.Result, itemType LimitConfigItemType, key string, ipNet *iptree.IPTree, regexp *re.Regexp) (*LimitConfigItem, error) {
for timeWindowKey, duration := range timeWindows {
q := item.Get(timeWindowKey)
if q.Exists() && q.Int() > 0 {
if q.Exists() {
count := q.Int()
if count <= 0 {
return nil, fmt.Errorf("'%s' must be a positive integer for key '%s', got %d", timeWindowKey, key, count)
}
return &LimitConfigItem{
ConfigType: itemType,
Key: key,
IpNet: ipNet,
Regexp: regexp,
Count: q.Int(),
Count: count,
TimeWindow: duration,
}, nil
}

25
plugins/wasm-go/extensions/cluster-key-rate-limit/config/config_test.go
View File

@@ -20,6 +20,16 @@ func TestParseClusterKeyRateLimitConfig(t *testing.T) {
json: `{}`,
expectedErr: errors.New("missing rule_name in config"),
},
{
name: "GlobalThreshold_InvalidThreshold",
json: `{
"rule_name": "invalid-threshold",
"global_threshold": {
"query_per_minute": -100
}
}`,
expectedErr: errors.New("failed to parse global_threshold: 'query_per_minute' must be a positive integer, got -100"),
},
{
name: "GlobalThreshold_QueryPerSecond",
json: `{
@@ -56,6 +66,21 @@ func TestParseClusterKeyRateLimitConfig(t *testing.T) {
RejectedMsg: DefaultRejectedMsg,
},
},
{
name: "RuleItems_InvalidThreshold",
json: `{
"rule_name": "invalid-threshold",
"rule_items": [
{
"limit_by_header": "x-test",
"limit_keys": [
{"key": "key1", "query_per_minute": -100}
]
}
]
}`,
expectedErr: errors.New("failed to parse rule_item in rule_items: 'query_per_minute' must be a positive integer for key 'key1', got -100"),
},
{
name: "RuleItems_SingleRule",
json: `{

2
plugins/wasm-go/extensions/cluster-key-rate-limit/go.sum
View File

@@ -6,8 +6,6 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20250611100342-5654e89a7a80 h1:xqmtTZI0JQ2O+Lg9/CE6c+Tw9KD6FnvWw8EpLVuuvfg=
github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20250611100342-5654e89a7a80/go.mod h1:tRI2LfMudSkKHhyv1uex3BWzcice2s/l8Ah8axporfA=
github.com/higress-group/wasm-go v1.0.0 h1:4Ik5n3FsJ5+r13KLQl2ky+8NuAE8dfWQwoKxXYD2KAw=
github.com/higress-group/wasm-go v1.0.0/go.mod h1:ODBV27sjmhIW8Cqv3R74EUcTnbdkE69bmXBQFuRkY1M=
github.com/higress-group/wasm-go v1.0.1 h1:T1m++qTEANp8+jwE0sxltwtaTKmrHCkLOp1m9N+YeqY=
github.com/higress-group/wasm-go v1.0.1/go.mod h1:9k7L730huS/q4V5iH9WLDgf5ZUHEtfhM/uXcegKDG/M=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=

30
plugins/wasm-go/extensions/cluster-key-rate-limit/main.go
View File

@@ -44,12 +44,12 @@ func init() {
}
const (
// ClusterKeyPrefix 集群限流插件在 Redis 中 key 的统一前缀
ClusterKeyPrefix = "higress-cluster-key-rate-limit"
// ClusterGlobalRateLimitFormat 全局限流模式 redis key 为 ClusterKeyPrefix:限流规则名称:global_threshold:时间窗口:窗口内限流数
ClusterGlobalRateLimitFormat = ClusterKeyPrefix + ":%s:global_threshold:%d:%d"
// ClusterRateLimitFormat 规则限流模式 redis key 为 ClusterKeyPrefix:限流规则名称:限流类型:时间窗口:窗口内限流数:限流key名称:限流key对应的实际值
ClusterRateLimitFormat = ClusterKeyPrefix + ":%s:%s:%d:%d:%s:%s"
// 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"
FixedWindowScript = `
local ttl = redis.call('ttl', KEYS[1])
if ttl < 0 then
@@ -86,24 +86,24 @@ func parseConfig(json gjson.Result, cfg *config.ClusterKeyRateLimitConfig) error
return nil
}
func onHttpRequestHeaders(ctx wrapper.HttpContext, config config.ClusterKeyRateLimitConfig) types.Action {
func onHttpRequestHeaders(ctx wrapper.HttpContext, cfg config.ClusterKeyRateLimitConfig) types.Action {
ctx.DisableReroute()
limitKey, count, timeWindow := "", int64(0), int64(0)
if config.GlobalThreshold != nil {
if cfg.GlobalThreshold != nil {
// 全局限流模式
limitKey = fmt.Sprintf(ClusterGlobalRateLimitFormat, config.RuleName, config.GlobalThreshold.TimeWindow, config.GlobalThreshold.Count)
count = config.GlobalThreshold.Count
timeWindow = config.GlobalThreshold.TimeWindow
limitKey = fmt.Sprintf(ClusterGlobalRateLimitFormat, cfg.RuleName, cfg.GlobalThreshold.TimeWindow, cfg.GlobalThreshold.Count)
count = cfg.GlobalThreshold.Count
timeWindow = cfg.GlobalThreshold.TimeWindow
} else {
// 规则限流模式
val, ruleItem, configItem := checkRequestAgainstLimitRule(ctx, config.RuleItems)
val, ruleItem, configItem := checkRequestAgainstLimitRule(ctx, cfg.RuleItems)
if ruleItem == nil || configItem == nil {
// 没有匹配到限流规则直接返回
return types.ActionContinue
}
limitKey = fmt.Sprintf(ClusterRateLimitFormat, config.RuleName, ruleItem.LimitType, configItem.TimeWindow, configItem.Count, ruleItem.Key, val)
limitKey = fmt.Sprintf(ClusterRateLimitFormat, cfg.RuleName, ruleItem.LimitType, configItem.TimeWindow, configItem.Count, ruleItem.Key, val)
count = configItem.Count
timeWindow = configItem.TimeWindow
}
@@ -111,7 +111,7 @@ func onHttpRequestHeaders(ctx wrapper.HttpContext, config config.ClusterKeyRateL
// 执行限流逻辑
keys := []interface{}{limitKey}
args := []interface{}{count, timeWindow}
err := config.RedisClient.Eval(FixedWindowScript, 1, keys, args, func(response resp.Value) {
err := cfg.RedisClient.Eval(FixedWindowScript, 1, keys, args, func(response resp.Value) {
resultArray := response.Array()
if len(resultArray) != 3 {
log.Errorf("redis response parse error, response: %v", response)
@@ -125,7 +125,7 @@ func onHttpRequestHeaders(ctx wrapper.HttpContext, config config.ClusterKeyRateL
}
if context.remaining < 0 {
// 触发限流
rejected(config, context)
rejected(cfg, context)
} else {
ctx.SetContext(LimitContextKey, context)
proxywasm.ResumeHttpRequest()