diff --git a/docs/observability.md b/docs/observability.md index 3a11cd3d48..c695087b74 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -248,6 +248,48 @@ Number of currently active MCP connections. | `transport` | string | Backend transport type | | `connection_type` | string | `"sse"` (only present for SSE connections) | +### Rate Limit Metrics + +These metrics are emitted for Redis-backed rate limit checks used by MCPServer +and VirtualMCPServer. Prometheus appends `_total` to counter names. The latency +histogram is exported with the `_seconds` unit suffix and the standard +`_bucket`, `_sum`, and `_count` series suffixes. + +#### `toolhive_rate_limit_decisions` (Counter) + +Total number of rate limit bucket decisions. An allowed request increments once +for every applicable bucket. A rejected request increments only for the first +bucket rejected by the atomic Redis check. Requests with no applicable bucket +do not increment this counter. + +| Attribute | Type | Description | +|-----------|------|-------------| +| `namespace` | string | Kubernetes namespace associated with the server | +| `server` | string | MCPServer or VirtualMCPServer name | +| `decision` | string | `"allowed"` or `"rejected"` | +| `scope` | string | `"shared"` or `"per_user"` | +| `operation_type` | string | `"server"` or `"tool"` | + +#### `toolhive_rate_limit_redis_errors` (Counter) + +Total number of Redis errors encountered while checking rate limits. + +| Attribute | Type | Description | +|-----------|------|-------------| +| `namespace` | string | Kubernetes namespace associated with the server | +| `server` | string | MCPServer or VirtualMCPServer name | +| `error_type` | string | `"timeout"`, `"connection"`, `"auth"`, or `"other"` | + +#### `toolhive_rate_limit_check_latency` (Histogram, seconds) + +Duration of each attempted atomic Redis Lua rate limit check, including failed +checks. + +| Attribute | Type | Description | +|-----------|------|-------------| +| `namespace` | string | Kubernetes namespace associated with the server | +| `server` | string | MCPServer or VirtualMCPServer name | + ## Span Attributes ### HTTP Attributes diff --git a/go.mod b/go.mod index d52fc46a2b..585a44e201 100644 --- a/go.mod +++ b/go.mod @@ -44,6 +44,8 @@ require ( github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c github.com/pressly/goose/v3 v3.27.1 github.com/prometheus/client_golang v1.23.2 + github.com/prometheus/client_model v0.6.2 + github.com/prometheus/common v0.67.5 github.com/redis/go-redis/v9 v9.21.0 github.com/shirou/gopsutil/v4 v4.26.5 github.com/spf13/viper v1.21.0 @@ -233,8 +235,6 @@ require ( github.com/ory/x v0.0.665 // indirect github.com/pjbgf/sha1cd v0.6.0 // indirect github.com/power-devops/perfstat v0.0.0-20240221224432-82ca36839d55 // indirect - github.com/prometheus/client_model v0.6.2 // indirect - github.com/prometheus/common v0.67.5 // indirect github.com/prometheus/otlptranslator v1.0.0 // indirect github.com/prometheus/procfs v0.20.1 // indirect github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect diff --git a/pkg/ratelimit/limiter.go b/pkg/ratelimit/limiter.go index ad47764dbd..14fca1ce87 100644 --- a/pkg/ratelimit/limiter.go +++ b/pkg/ratelimit/limiter.go @@ -13,6 +13,8 @@ import ( "time" "github.com/redis/go-redis/v9" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/metric" v1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/pkg/auth" @@ -65,11 +67,29 @@ func Allow(ctx context.Context, limiter Limiter, identity *auth.Identity, toolNa // Returns a no-op limiter (always allows) when crd is nil. // namespace and name identify the MCP server for Redis key derivation. func NewLimiter(client redis.Cmdable, namespace, name string, crd *v1beta1.RateLimitConfig) (Limiter, error) { + return newLimiter(client, namespace, name, crd, otel.GetMeterProvider()) +} + +func newLimiter( + client redis.Cmdable, + namespace string, + name string, + crd *v1beta1.RateLimitConfig, + meterProvider metric.MeterProvider, +) (Limiter, error) { if crd == nil { return noopLimiter{}, nil } - l := &limiter{client: client} + telemetry, err := newRateLimitTelemetry(meterProvider, namespace, name) + if err != nil { + return nil, fmt.Errorf("rate limit telemetry: %w", err) + } + + l := &limiter{ + client: client, + telemetry: telemetry, + } if crd.Shared != nil { b, err := newBucket(namespace, name, "shared", crd.Shared) @@ -122,9 +142,17 @@ type bucketSpec struct { refillPeriod time.Duration } +// limitCheck keeps a bucket paired with its metric dimensions. +type limitCheck struct { + bucket *bucket.TokenBucket + scope string + operationType string +} + // limiter is the concrete implementation of Limiter. type limiter struct { client redis.Cmdable + telemetry *rateLimitTelemetry serverBucket *bucket.TokenBucket // nil when no shared server limit toolBuckets map[string]*bucket.TokenBucket // tool name -> shared bucket perUserSpec *bucketSpec // nil when no server-level per-user limit @@ -136,13 +164,21 @@ type limiter struct { // a rejected per-tool or per-user call from draining other budgets. func (l *limiter) Allow(ctx context.Context, toolName, userID string) (*Decision, error) { // Collect applicable buckets in priority order. - var buckets []*bucket.TokenBucket + var checks []limitCheck if l.serverBucket != nil { - buckets = append(buckets, l.serverBucket) + checks = append(checks, limitCheck{ + bucket: l.serverBucket, + scope: rateLimitScopeShared, + operationType: rateLimitOperationServer, + }) } if toolName != "" && l.toolBuckets != nil { if tb, ok := l.toolBuckets[toolName]; ok { - buckets = append(buckets, tb) + checks = append(checks, limitCheck{ + bucket: tb, + scope: rateLimitScopeShared, + operationType: rateLimitOperationTool, + }) } } @@ -157,40 +193,58 @@ func (l *limiter) Allow(ctx context.Context, toolName, userID string) (*Decision if userID != "" { if l.perUserSpec != nil { s := l.perUserSpec - buckets = append(buckets, bucket.New( - s.namespace, s.serverName, - "user:"+userID, - s.maxTokens, s.refillPeriod, - )) + checks = append(checks, limitCheck{ + bucket: bucket.New( + s.namespace, s.serverName, + "user:"+userID, + s.maxTokens, s.refillPeriod, + ), + scope: rateLimitScopePerUser, + operationType: rateLimitOperationServer, + }) } if toolName != "" && l.perUserTools != nil { if s, ok := l.perUserTools[toolName]; ok { // Key prefix "user-tool:" is distinct from "user:" to prevent // collisions when a userID contains delimiter characters. - buckets = append(buckets, bucket.New( - s.namespace, s.serverName, - "user-tool:"+toolName+":"+userID, - s.maxTokens, s.refillPeriod, - )) + checks = append(checks, limitCheck{ + bucket: bucket.New( + s.namespace, s.serverName, + "user-tool:"+toolName+":"+userID, + s.maxTokens, s.refillPeriod, + ), + scope: rateLimitScopePerUser, + operationType: rateLimitOperationTool, + }) } } } - if len(buckets) == 0 { + if len(checks) == 0 { return &Decision{Allowed: true}, nil } + buckets := make([]*bucket.TokenBucket, len(checks)) + for i, check := range checks { + buckets[i] = check.bucket + } + + start := time.Now() rejectedIdx, err := bucket.ConsumeAll(ctx, l.client, buckets) + l.telemetry.recordCheckLatency(ctx, time.Since(start)) if err != nil { + l.telemetry.recordRedisError(ctx, err) return nil, fmt.Errorf("rate limit check: %w", err) } if rejectedIdx >= 0 { + l.telemetry.recordRejected(ctx, checks[rejectedIdx]) return &Decision{ Allowed: false, RetryAfter: buckets[rejectedIdx].RetryAfter(), }, nil } + l.telemetry.recordAllowed(ctx, checks) return &Decision{Allowed: true}, nil } diff --git a/pkg/ratelimit/observability.go b/pkg/ratelimit/observability.go new file mode 100644 index 0000000000..447ed64215 --- /dev/null +++ b/pkg/ratelimit/observability.go @@ -0,0 +1,175 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package ratelimit + +import ( + "context" + "errors" + "fmt" + "io" + "net" + "strings" + "time" + + "github.com/redis/go-redis/v9" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" + + "github.com/stacklok/toolhive/pkg/telemetry" +) + +const ( + rateLimitInstrumentationName = "github.com/stacklok/toolhive/pkg/ratelimit" + + rateLimitDecisionAllowed = "allowed" + rateLimitDecisionRejected = "rejected" + + rateLimitScopeShared = "shared" + rateLimitScopePerUser = "per_user" + + rateLimitOperationServer = "server" + rateLimitOperationTool = "tool" + + redisErrorTypeTimeout = "timeout" + redisErrorTypeConnection = "connection" + redisErrorTypeAuth = "auth" + redisErrorTypeOther = "other" +) + +type rateLimitTelemetry struct { + namespace string + serverName string + + decisions metric.Int64Counter + redisErrors metric.Int64Counter + checkLatency metric.Float64Histogram +} + +func newRateLimitTelemetry( + meterProvider metric.MeterProvider, + namespace string, + serverName string, +) (*rateLimitTelemetry, error) { + if meterProvider == nil { + return nil, nil + } + + meter := meterProvider.Meter(rateLimitInstrumentationName) + + decisions, err := meter.Int64Counter( + "toolhive_rate_limit_decisions", + metric.WithDescription("Total number of rate limit bucket decisions"), + ) + if err != nil { + return nil, fmt.Errorf("create decision counter: %w", err) + } + + redisErrors, err := meter.Int64Counter( + "toolhive_rate_limit_redis_errors", + metric.WithDescription("Total number of Redis errors during rate limit checks"), + ) + if err != nil { + return nil, fmt.Errorf("create Redis error counter: %w", err) + } + + checkLatency, err := meter.Float64Histogram( + "toolhive_rate_limit_check_latency", + metric.WithDescription("Duration of Redis Lua rate limit checks in seconds"), + metric.WithUnit("s"), + metric.WithExplicitBucketBoundaries(telemetry.MCPHistogramBuckets...), + ) + if err != nil { + return nil, fmt.Errorf("create check latency histogram: %w", err) + } + + return &rateLimitTelemetry{ + namespace: namespace, + serverName: serverName, + decisions: decisions, + redisErrors: redisErrors, + checkLatency: checkLatency, + }, nil +} + +func (t *rateLimitTelemetry) recordAllowed(ctx context.Context, checks []limitCheck) { + if t == nil { + return + } + for _, check := range checks { + t.recordDecision(ctx, rateLimitDecisionAllowed, check) + } +} + +func (t *rateLimitTelemetry) recordRejected(ctx context.Context, check limitCheck) { + if t == nil { + return + } + t.recordDecision(ctx, rateLimitDecisionRejected, check) +} + +func (t *rateLimitTelemetry) recordDecision(ctx context.Context, decision string, check limitCheck) { + t.decisions.Add(ctx, 1, metric.WithAttributes( + attribute.String("namespace", t.namespace), + attribute.String("server", t.serverName), + attribute.String("decision", decision), + attribute.String("scope", check.scope), + attribute.String("operation_type", check.operationType), + )) +} + +func (t *rateLimitTelemetry) recordRedisError(ctx context.Context, err error) { + if t == nil { + return + } + t.redisErrors.Add(ctx, 1, metric.WithAttributes( + attribute.String("namespace", t.namespace), + attribute.String("server", t.serverName), + attribute.String("error_type", classifyRedisError(err)), + )) +} + +func (t *rateLimitTelemetry) recordCheckLatency(ctx context.Context, duration time.Duration) { + if t == nil { + return + } + t.checkLatency.Record(ctx, duration.Seconds(), metric.WithAttributes( + attribute.String("namespace", t.namespace), + attribute.String("server", t.serverName), + )) +} + +func classifyRedisError(err error) string { + if redis.IsAuthError(err) { + return redisErrorTypeAuth + } + + if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, redis.ErrPoolTimeout) { + return redisErrorTypeTimeout + } + var netErr net.Error + if errors.As(err, &netErr) && netErr.Timeout() { + return redisErrorTypeTimeout + } + + if errors.Is(err, redis.ErrClosed) || + errors.Is(err, net.ErrClosed) || + errors.Is(err, io.EOF) || + errors.Is(err, io.ErrUnexpectedEOF) { + return redisErrorTypeConnection + } + var opErr *net.OpError + if errors.As(err, &opErr) { + return redisErrorTypeConnection + } + + message := strings.ToLower(err.Error()) + if strings.Contains(message, "connection refused") || + strings.Contains(message, "connection reset") || + strings.Contains(message, "broken pipe") || + strings.Contains(message, "no such host") { + return redisErrorTypeConnection + } + + return redisErrorTypeOther +} diff --git a/pkg/ratelimit/observability_test.go b/pkg/ratelimit/observability_test.go new file mode 100644 index 0000000000..8eb2ba58c8 --- /dev/null +++ b/pkg/ratelimit/observability_test.go @@ -0,0 +1,349 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package ratelimit + +import ( + "context" + "errors" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel/attribute" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" + "go.opentelemetry.io/otel/sdk/metric/metricdata" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + + v1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" +) + +func TestRateLimitMetrics_SharedDecisionsAndLatency(t *testing.T) { + t.Parallel() + client, _ := newTestClient(t) + reader, meterProvider := newRateLimitMeterProvider() + + limiter, err := newLimiter(client, "test-ns", "test-server", &v1beta1.RateLimitConfig{ + Shared: &v1beta1.RateLimitBucket{ + MaxTokens: 10, + RefillPeriod: metav1.Duration{Duration: time.Minute}, + }, + Tools: []v1beta1.ToolRateLimitConfig{ + { + Name: "search", + Shared: &v1beta1.RateLimitBucket{ + MaxTokens: 1, + RefillPeriod: metav1.Duration{Duration: time.Minute}, + }, + }, + }, + }, meterProvider) + require.NoError(t, err) + + first, err := limiter.Allow(t.Context(), "search", "") + require.NoError(t, err) + require.True(t, first.Allowed) + + second, err := limiter.Allow(t.Context(), "search", "") + require.NoError(t, err) + require.False(t, second.Allowed) + + metrics := collectRateLimitMetrics(t, reader) + decisions := requireRateLimitMetric(t, metrics, "toolhive_rate_limit_decisions") + assert.Equal(t, int64(1), counterValueWithAttributes(t, decisions, map[string]string{ + "namespace": "test-ns", + "server": "test-server", + "decision": rateLimitDecisionAllowed, + "scope": rateLimitScopeShared, + "operation_type": rateLimitOperationServer, + })) + assert.Equal(t, int64(1), counterValueWithAttributes(t, decisions, map[string]string{ + "namespace": "test-ns", + "server": "test-server", + "decision": rateLimitDecisionAllowed, + "scope": rateLimitScopeShared, + "operation_type": rateLimitOperationTool, + })) + assert.Equal(t, int64(1), counterValueWithAttributes(t, decisions, map[string]string{ + "namespace": "test-ns", + "server": "test-server", + "decision": rateLimitDecisionRejected, + "scope": rateLimitScopeShared, + "operation_type": rateLimitOperationTool, + })) + + latency := requireRateLimitMetric(t, metrics, "toolhive_rate_limit_check_latency") + assert.Equal(t, uint64(2), histogramCountWithAttributes(t, latency, map[string]string{ + "namespace": "test-ns", + "server": "test-server", + })) +} + +func TestRateLimitMetrics_PerUserDecisions(t *testing.T) { + t.Parallel() + client, _ := newTestClient(t) + reader, meterProvider := newRateLimitMeterProvider() + + limiter, err := newLimiter(client, "test-ns", "test-server", &v1beta1.RateLimitConfig{ + PerUser: &v1beta1.RateLimitBucket{ + MaxTokens: 10, + RefillPeriod: metav1.Duration{Duration: time.Minute}, + }, + Tools: []v1beta1.ToolRateLimitConfig{ + { + Name: "search", + PerUser: &v1beta1.RateLimitBucket{ + MaxTokens: 1, + RefillPeriod: metav1.Duration{Duration: time.Minute}, + }, + }, + }, + }, meterProvider) + require.NoError(t, err) + + first, err := limiter.Allow(t.Context(), "search", "alice") + require.NoError(t, err) + require.True(t, first.Allowed) + + second, err := limiter.Allow(t.Context(), "search", "alice") + require.NoError(t, err) + require.False(t, second.Allowed) + + metrics := collectRateLimitMetrics(t, reader) + decisions := requireRateLimitMetric(t, metrics, "toolhive_rate_limit_decisions") + assert.Equal(t, int64(1), counterValueWithAttributes(t, decisions, map[string]string{ + "namespace": "test-ns", + "server": "test-server", + "decision": rateLimitDecisionAllowed, + "scope": rateLimitScopePerUser, + "operation_type": rateLimitOperationServer, + })) + assert.Equal(t, int64(1), counterValueWithAttributes(t, decisions, map[string]string{ + "namespace": "test-ns", + "server": "test-server", + "decision": rateLimitDecisionAllowed, + "scope": rateLimitScopePerUser, + "operation_type": rateLimitOperationTool, + })) + assert.Equal(t, int64(1), counterValueWithAttributes(t, decisions, map[string]string{ + "namespace": "test-ns", + "server": "test-server", + "decision": rateLimitDecisionRejected, + "scope": rateLimitScopePerUser, + "operation_type": rateLimitOperationTool, + })) +} + +func TestRateLimitMetrics_RedisErrorAndFailedLatency(t *testing.T) { + t.Parallel() + client, redisServer := newTestClient(t) + reader, meterProvider := newRateLimitMeterProvider() + + limiter, err := newLimiter(client, "test-ns", "test-server", &v1beta1.RateLimitConfig{ + Shared: &v1beta1.RateLimitBucket{ + MaxTokens: 1, + RefillPeriod: metav1.Duration{Duration: time.Minute}, + }, + }, meterProvider) + require.NoError(t, err) + + redisServer.Close() + _, err = limiter.Allow(t.Context(), "", "") + require.Error(t, err) + + metrics := collectRateLimitMetrics(t, reader) + redisErrors := requireRateLimitMetric(t, metrics, "toolhive_rate_limit_redis_errors") + assert.Equal(t, int64(1), counterValueWithAttributes(t, redisErrors, map[string]string{ + "namespace": "test-ns", + "server": "test-server", + "error_type": redisErrorTypeConnection, + })) + + latency := requireRateLimitMetric(t, metrics, "toolhive_rate_limit_check_latency") + assert.Equal(t, uint64(1), histogramCountWithAttributes(t, latency, map[string]string{ + "namespace": "test-ns", + "server": "test-server", + })) + assert.Nil(t, findRateLimitMetric(metrics, "toolhive_rate_limit_decisions")) +} + +func TestRateLimitMetrics_NoApplicableBucketRecordsNothing(t *testing.T) { + t.Parallel() + client, _ := newTestClient(t) + reader, meterProvider := newRateLimitMeterProvider() + + limiter, err := newLimiter(client, "test-ns", "test-server", &v1beta1.RateLimitConfig{ + Tools: []v1beta1.ToolRateLimitConfig{ + { + Name: "search", + Shared: &v1beta1.RateLimitBucket{ + MaxTokens: 1, + RefillPeriod: metav1.Duration{Duration: time.Minute}, + }, + }, + }, + }, meterProvider) + require.NoError(t, err) + + decision, err := limiter.Allow(t.Context(), "other-tool", "") + require.NoError(t, err) + require.True(t, decision.Allowed) + + metrics := collectRateLimitMetrics(t, reader) + assert.Empty(t, metrics.ScopeMetrics) +} + +func TestRateLimitMetrics_NilMeterProviderIsNoOp(t *testing.T) { + t.Parallel() + client, _ := newTestClient(t) + + limiter, err := newLimiter(client, "test-ns", "test-server", &v1beta1.RateLimitConfig{ + Shared: &v1beta1.RateLimitBucket{ + MaxTokens: 1, + RefillPeriod: metav1.Duration{Duration: time.Minute}, + }, + }, nil) + require.NoError(t, err) + + decision, err := limiter.Allow(t.Context(), "", "") + require.NoError(t, err) + assert.True(t, decision.Allowed) +} + +func TestClassifyRedisError(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + err error + want string + }{ + { + name: "deadline exceeded", + err: context.DeadlineExceeded, + want: redisErrorTypeTimeout, + }, + { + name: "authentication", + err: errors.New("NOAUTH Authentication required"), + want: redisErrorTypeAuth, + }, + { + name: "connection refused", + err: errors.New("dial tcp 127.0.0.1:6379: connection refused"), + want: redisErrorTypeConnection, + }, + { + name: "other", + err: errors.New("ERR script failed"), + want: redisErrorTypeOther, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, tt.want, classifyRedisError(tt.err)) + }) + } +} + +func TestRateLimitMetricNamesUseToolHivePrefix(t *testing.T) { + t.Parallel() + reader, meterProvider := newRateLimitMeterProvider() + + telemetry, err := newRateLimitTelemetry(meterProvider, "test-ns", "test-server") + require.NoError(t, err) + telemetry.recordRejected(t.Context(), limitCheck{ + scope: rateLimitScopeShared, + operationType: rateLimitOperationServer, + }) + telemetry.recordRedisError(t.Context(), errors.New("ERR script failed")) + telemetry.recordCheckLatency(t.Context(), time.Millisecond) + + metrics := collectRateLimitMetrics(t, reader) + require.NotEmpty(t, metrics.ScopeMetrics) + for _, scopeMetrics := range metrics.ScopeMetrics { + for _, measured := range scopeMetrics.Metrics { + assert.True(t, strings.HasPrefix(measured.Name, "toolhive_"), measured.Name) + } + } +} + +func newRateLimitMeterProvider() (*sdkmetric.ManualReader, *sdkmetric.MeterProvider) { + reader := sdkmetric.NewManualReader() + return reader, sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader)) +} + +func collectRateLimitMetrics(t *testing.T, reader *sdkmetric.ManualReader) metricdata.ResourceMetrics { + t.Helper() + var metrics metricdata.ResourceMetrics + require.NoError(t, reader.Collect(context.Background(), &metrics)) + return metrics +} + +func findRateLimitMetric(metrics metricdata.ResourceMetrics, name string) *metricdata.Metrics { + for _, scopeMetrics := range metrics.ScopeMetrics { + for i := range scopeMetrics.Metrics { + if scopeMetrics.Metrics[i].Name == name { + return &scopeMetrics.Metrics[i] + } + } + } + return nil +} + +func requireRateLimitMetric( + t *testing.T, + metrics metricdata.ResourceMetrics, + name string, +) *metricdata.Metrics { + t.Helper() + measured := findRateLimitMetric(metrics, name) + require.NotNil(t, measured, "metric %q not found", name) + return measured +} + +func counterValueWithAttributes( + t *testing.T, + measured *metricdata.Metrics, + want map[string]string, +) int64 { + t.Helper() + sum, ok := measured.Data.(metricdata.Sum[int64]) + require.True(t, ok, "metric %q is not an int64 sum", measured.Name) + for _, point := range sum.DataPoints { + if attributesMatch(point.Attributes, want) { + return point.Value + } + } + return 0 +} + +func histogramCountWithAttributes( + t *testing.T, + measured *metricdata.Metrics, + want map[string]string, +) uint64 { + t.Helper() + histogram, ok := measured.Data.(metricdata.Histogram[float64]) + require.True(t, ok, "metric %q is not a float64 histogram", measured.Name) + for _, point := range histogram.DataPoints { + if attributesMatch(point.Attributes, want) { + return point.Count + } + } + return 0 +} + +func attributesMatch(attributes attribute.Set, want map[string]string) bool { + for key, wantValue := range want { + value, ok := attributes.Value(attribute.Key(key)) + if !ok || value.AsString() != wantValue { + return false + } + } + return true +} diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_rate_limiting_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_rate_limiting_test.go index 981d3fae5f..75a8f1f2ed 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_rate_limiting_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_rate_limiting_test.go @@ -17,6 +17,9 @@ import ( "github.com/mark3labs/mcp-go/mcp" "github.com/onsi/ginkgo/v2" "github.com/onsi/gomega" + dto "github.com/prometheus/client_model/go" + "github.com/prometheus/common/expfmt" + "github.com/prometheus/common/model" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" @@ -40,6 +43,7 @@ var _ = ginkgo.Describe("VirtualMCPServer Rate Limiting", ginkgo.Ordered, func() vmcpName string redisName string oidcName string + telemetryName string vmcpLocalPort int oidcLocalPort int vmcpPortForwardCleanup func() @@ -54,6 +58,7 @@ var _ = ginkgo.Describe("VirtualMCPServer Rate Limiting", ginkgo.Ordered, func() vmcpName = fmt.Sprintf("e2e-rl-vmcp-%d", ts) redisName = fmt.Sprintf("e2e-rl-redis-%d", ts) oidcName = fmt.Sprintf("e2e-rl-oidc-%d", ts) + telemetryName = fmt.Sprintf("e2e-rl-telemetry-%d", ts) ginkgo.By("Deploying Redis") deployRedis(redisName) @@ -112,10 +117,30 @@ var _ = ginkgo.Describe("VirtualMCPServer Rate Limiting", ginkgo.Ordered, func() return nil }, timeout, pollInterval).Should(gomega.Succeed()) + ginkgo.By("Creating Prometheus telemetry configuration") + gomega.Expect(k8sClient.Create(ctx, &mcpv1beta1.MCPTelemetryConfig{ + ObjectMeta: metav1.ObjectMeta{ + Name: telemetryName, + Namespace: defaultNamespace, + }, + Spec: mcpv1beta1.MCPTelemetryConfigSpec{ + Prometheus: &mcpv1beta1.PrometheusConfig{Enabled: true}, + }, + })).To(gomega.Succeed()) + gomega.Eventually(func() bool { + telemetryConfig := &mcpv1beta1.MCPTelemetryConfig{} + err := k8sClient.Get(ctx, types.NamespacedName{ + Name: telemetryName, + Namespace: defaultNamespace, + }, telemetryConfig) + return err == nil && telemetryConfig.Status.ConfigHash != "" + }, timeout, pollInterval).Should(gomega.BeTrue()) + redisAddr := fmt.Sprintf("%s.%s.svc.cluster.local:6379", redisName, defaultNamespace) ginkgo.By("Creating VirtualMCPServer with per-user rate limiting") gomega.Expect(k8sClient.Create(ctx, v1beta1test.NewVirtualMCPServer(vmcpName, defaultNamespace, v1beta1test.WithVMCPGroupRef(mcpGroupName), + v1beta1test.WithVMCPTelemetryConfigRef(telemetryName), v1beta1test.WithVMCPConfig(vmcpconfig.Config{ Group: mcpGroupName, RateLimiting: &mcpv1beta1.RateLimitConfig{ @@ -166,6 +191,9 @@ var _ = ginkgo.Describe("VirtualMCPServer Rate Limiting", ginkgo.Ordered, func() _ = k8sClient.Delete(ctx, &mcpv1beta1.MCPOIDCConfig{ ObjectMeta: metav1.ObjectMeta{Name: oidcName, Namespace: defaultNamespace}, }) + _ = k8sClient.Delete(ctx, &mcpv1beta1.MCPTelemetryConfig{ + ObjectMeta: metav1.ObjectMeta{Name: telemetryName, Namespace: defaultNamespace}, + }) cleanupRedis(redisName) }) @@ -198,9 +226,100 @@ var _ = ginkgo.Describe("VirtualMCPServer Rate Limiting", ginkgo.Ordered, func() data, ok := structured["data"].(map[string]any) gomega.Expect(ok).To(gomega.BeTrue()) gomega.Expect(data["retryAfterSeconds"]).To(gomega.BeNumerically(">", 0)) + + ginkgo.By("Scraping non-zero rate limit metrics") + gomega.Eventually(func() error { + metricFamilies, scrapeErr := scrapeRateLimitMetrics(vmcpLocalPort) + if scrapeErr != nil { + return scrapeErr + } + if !rateLimitCounterIsNonZero(metricFamilies, "toolhive_rate_limit_decisions", map[string]string{ + "decision": "allowed", + "scope": "per_user", + "operation_type": "server", + }) { + return fmt.Errorf("allowed rate limit decision counter is zero") + } + if !rateLimitCounterIsNonZero(metricFamilies, "toolhive_rate_limit_decisions", map[string]string{ + "decision": "rejected", + "scope": "per_user", + "operation_type": "server", + }) { + return fmt.Errorf("rejected rate limit decision counter is zero") + } + if !rateLimitHistogramIsNonZero(metricFamilies, "toolhive_rate_limit_check_latency") { + return fmt.Errorf("rate limit check latency histogram is empty") + } + return nil + }, 30*time.Second, time.Second).Should(gomega.Succeed()) }) }) +func scrapeRateLimitMetrics(port int) (map[string]*dto.MetricFamily, error) { + url := fmt.Sprintf("http://localhost:%d/metrics", port) + resp, err := http.Get(url) //nolint:noctx + if err != nil { + return nil, err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("metrics endpoint returned HTTP %d", resp.StatusCode) + } + + parser := expfmt.NewTextParser(model.UTF8Validation) + return parser.TextToMetricFamilies(resp.Body) +} + +func rateLimitCounterIsNonZero( + families map[string]*dto.MetricFamily, + namePrefix string, + wantLabels map[string]string, +) bool { + for name, family := range families { + if !strings.HasPrefix(name, namePrefix) { + continue + } + for _, metric := range family.Metric { + if prometheusLabelsMatch(metric, wantLabels) && + metric.Counter != nil && + metric.Counter.GetValue() > 0 { + return true + } + } + } + return false +} + +func rateLimitHistogramIsNonZero( + families map[string]*dto.MetricFamily, + namePrefix string, +) bool { + for name, family := range families { + if !strings.HasPrefix(name, namePrefix) { + continue + } + for _, metric := range family.Metric { + if metric.Histogram != nil && metric.Histogram.GetSampleCount() > 0 { + return true + } + } + } + return false +} + +func prometheusLabelsMatch(metric *dto.Metric, want map[string]string) bool { + labels := make(map[string]string, len(metric.Label)) + for _, pair := range metric.Label { + labels[pair.GetName()] = pair.GetValue() + } + for name, value := range want { + if labels[name] != value { + return false + } + } + return true +} + func fetchRateLimitOIDCToken(oidcPort int, subject string) string { url := fmt.Sprintf("http://localhost:%d/token?subject=%s", oidcPort, subject) resp, err := http.Post(url, "application/x-www-form-urlencoded", nil) //nolint:noctx