diff --git a/plugins/wasm-go/extensions/ai-token-ratelimit/config.go b/plugins/wasm-go/extensions/ai-token-ratelimit/config.go index d9b3482c00..1df66219e9 100644 --- a/plugins/wasm-go/extensions/ai-token-ratelimit/config.go +++ b/plugins/wasm-go/extensions/ai-token-ratelimit/config.go @@ -3,6 +3,7 @@ package main import ( "errors" "fmt" + "github.com/higress-group/proxy-wasm-go-sdk/proxywasm" "strings" "github.com/alibaba/higress/plugins/wasm-go/pkg/wrapper" @@ -59,6 +60,7 @@ type ClusterKeyRateLimitConfig struct { rejectedCode uint32 // 当请求超过阈值被拒绝时,返回的HTTP状态码 rejectedMsg string // 当请求超过阈值被拒绝时,返回的响应体 redisClient wrapper.RedisClient + counterMetrics map[string]proxywasm.MetricCounter // Metrics } type LimitRuleItem struct { diff --git a/plugins/wasm-go/extensions/ai-token-ratelimit/main.go b/plugins/wasm-go/extensions/ai-token-ratelimit/main.go index 3ebec532c2..706715f7e2 100644 --- a/plugins/wasm-go/extensions/ai-token-ratelimit/main.go +++ b/plugins/wasm-go/extensions/ai-token-ratelimit/main.go @@ -88,6 +88,8 @@ func parseConfig(json gjson.Result, config *ClusterKeyRateLimitConfig, log wrapp if err != nil { return err } + // Metric settings + config.counterMetrics = make(map[string]proxywasm.MetricCounter) return nil } @@ -304,9 +306,56 @@ func getDownStreamIp(rule LimitRuleItem) (net.IP, error) { 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 rejected(config ClusterKeyRateLimitConfig, context LimitContext) { headers := make(map[string][]string) headers[RateLimitResetHeader] = []string{strconv.Itoa(context.reset)} _ = proxywasm.SendHttpResponseWithDetail( config.rejectedCode, "ai-token-ratelimit.rejected", reconvertHeaders(headers), []byte(config.rejectedMsg), -1) + + route, err := getRouteName() + if err != nil { + route = "none" + } + cluster, err := getClusterName() + if err != nil { + cluster = "none" + } + consumer, err := proxywasm.GetHttpRequestHeader(ConsumerHeader) + if err != nil { + consumer = "none" + } + metricsName := generateMetricName(route, cluster, "none", consumer, "ai-token-ratelimit") + config.incrementCounter(metricsName, 1) }