mirror of
https://github.com/alibaba/higress.git
synced 2026-06-09 12:47:28 +08:00
feat: add ratelimit metrics in the ai-token-ratelimit plugin (#1918)
This commit is contained in:
@@ -6,6 +6,7 @@ import (
|
|||||||
"strings"
|
"strings"
|
||||||
|
|
||||||
"github.com/alibaba/higress/plugins/wasm-go/pkg/wrapper"
|
"github.com/alibaba/higress/plugins/wasm-go/pkg/wrapper"
|
||||||
|
"github.com/higress-group/proxy-wasm-go-sdk/proxywasm"
|
||||||
"github.com/tidwall/gjson"
|
"github.com/tidwall/gjson"
|
||||||
re "github.com/wasilibs/go-re2"
|
re "github.com/wasilibs/go-re2"
|
||||||
"github.com/zmap/go-iptree/iptree"
|
"github.com/zmap/go-iptree/iptree"
|
||||||
@@ -59,6 +60,7 @@ type ClusterKeyRateLimitConfig struct {
|
|||||||
rejectedCode uint32 // 当请求超过阈值被拒绝时,返回的HTTP状态码
|
rejectedCode uint32 // 当请求超过阈值被拒绝时,返回的HTTP状态码
|
||||||
rejectedMsg string // 当请求超过阈值被拒绝时,返回的响应体
|
rejectedMsg string // 当请求超过阈值被拒绝时,返回的响应体
|
||||||
redisClient wrapper.RedisClient
|
redisClient wrapper.RedisClient
|
||||||
|
counterMetrics map[string]proxywasm.MetricCounter // Metrics
|
||||||
}
|
}
|
||||||
|
|
||||||
type LimitRuleItem struct {
|
type LimitRuleItem struct {
|
||||||
|
|||||||
@@ -65,6 +65,8 @@ const (
|
|||||||
RateLimitLimitHeader string = "X-TokenRateLimit-Limit" // 限制的总请求数
|
RateLimitLimitHeader string = "X-TokenRateLimit-Limit" // 限制的总请求数
|
||||||
RateLimitRemainingHeader string = "X-TokenRateLimit-Remaining" // 剩余还可以发送的请求数
|
RateLimitRemainingHeader string = "X-TokenRateLimit-Remaining" // 剩余还可以发送的请求数
|
||||||
RateLimitResetHeader string = "X-TokenRateLimit-Reset" // 限流重置时间(触发限流时返回)
|
RateLimitResetHeader string = "X-TokenRateLimit-Reset" // 限流重置时间(触发限流时返回)
|
||||||
|
|
||||||
|
TokenRateLimitCount = "token_ratelimit_count" // metric name
|
||||||
)
|
)
|
||||||
|
|
||||||
type LimitContext struct {
|
type LimitContext struct {
|
||||||
@@ -88,6 +90,8 @@ func parseConfig(json gjson.Result, config *ClusterKeyRateLimitConfig, log wrapp
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
// Metric settings
|
||||||
|
config.counterMetrics = make(map[string]proxywasm.MetricCounter)
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -304,9 +308,54 @@ func getDownStreamIp(rule LimitRuleItem) (net.IP, error) {
|
|||||||
return realIP, nil
|
return realIP, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (config *ClusterKeyRateLimitConfig) incrementCounter(metricName string, inc uint64) {
|
||||||
|
if inc == 0 {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
counter, ok := config.counterMetrics[metricName]
|
||||||
|
if !ok {
|
||||||
|
counter = proxywasm.DefineCounterMetric(metricName)
|
||||||
|
config.counterMetrics[metricName] = counter
|
||||||
|
}
|
||||||
|
counter.Increment(inc)
|
||||||
|
}
|
||||||
|
|
||||||
|
func generateMetricName(route, cluster, model, consumer, metricName string) string {
|
||||||
|
return fmt.Sprintf("route.%s.upstream.%s.model.%s.consumer.%s.metric.%s", route, cluster, model, consumer, metricName)
|
||||||
|
}
|
||||||
|
|
||||||
|
func getRouteName() (string, error) {
|
||||||
|
if raw, err := proxywasm.GetProperty([]string{"route_name"}); err != nil {
|
||||||
|
return "-", err
|
||||||
|
} else {
|
||||||
|
return string(raw), nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func getClusterName() (string, error) {
|
||||||
|
if raw, err := proxywasm.GetProperty([]string{"cluster_name"}); err != nil {
|
||||||
|
return "-", err
|
||||||
|
} else {
|
||||||
|
return string(raw), nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func getConsumer() (string, error) {
|
||||||
|
if consumer, err := proxywasm.GetHttpRequestHeader(ConsumerHeader); err != nil {
|
||||||
|
return "none", err
|
||||||
|
} else {
|
||||||
|
return consumer, nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func rejected(config ClusterKeyRateLimitConfig, context LimitContext) {
|
func rejected(config ClusterKeyRateLimitConfig, context LimitContext) {
|
||||||
headers := make(map[string][]string)
|
headers := make(map[string][]string)
|
||||||
headers[RateLimitResetHeader] = []string{strconv.Itoa(context.reset)}
|
headers[RateLimitResetHeader] = []string{strconv.Itoa(context.reset)}
|
||||||
_ = proxywasm.SendHttpResponseWithDetail(
|
_ = proxywasm.SendHttpResponseWithDetail(
|
||||||
config.rejectedCode, "ai-token-ratelimit.rejected", reconvertHeaders(headers), []byte(config.rejectedMsg), -1)
|
config.rejectedCode, "ai-token-ratelimit.rejected", reconvertHeaders(headers), []byte(config.rejectedMsg), -1)
|
||||||
|
|
||||||
|
route, _ := getRouteName()
|
||||||
|
cluster, _ := getClusterName()
|
||||||
|
consumer, _ := getConsumer()
|
||||||
|
config.incrementCounter(generateMetricName(route, cluster, "none", consumer, TokenRateLimitCount), 1)
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user