Skip to content

Commit 358b21a

Browse files
tadasantclaude
andcommitted
Add metrics for rate-limited requests
Add mcp_registry.http.rate_limited counter to track when requests are blocked by rate limiting. This enables monitoring rate limit activity over time via Prometheus. - Add RateLimitedRequests counter to telemetry.Metrics - Add OnRateLimited callback to rate limiter config - Wire up metrics in server initialization - Add test for callback behavior 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
1 parent bc2a870 commit 358b21a

File tree

4 files changed

+96
-4
lines changed

4 files changed

+96
-4
lines changed

internal/api/ratelimit/ratelimit.go

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,10 @@ import (
1313
"golang.org/x/time/rate"
1414
)
1515

16+
// OnRateLimitedFunc is a callback invoked when a request is rate limited.
17+
// It receives the client IP that was blocked.
18+
type OnRateLimitedFunc func(ip string)
19+
1620
// Config holds the rate limiting configuration
1721
type Config struct {
1822
// RequestsPerMinute is the maximum number of requests allowed per minute per IP
@@ -26,6 +30,9 @@ type Config struct {
2630
// MaxVisitors is the maximum number of visitor entries to track (memory protection).
2731
// When exceeded, oldest entries are evicted. Default: 100000.
2832
MaxVisitors int
33+
// OnRateLimited is an optional callback invoked when a request is rate limited.
34+
// Used for recording metrics.
35+
OnRateLimited OnRateLimitedFunc
2936
}
3037

3138
// DefaultConfig returns the default rate limiting configuration
@@ -260,6 +267,11 @@ func (rl *RateLimiter) Middleware(next http.Handler) http.Handler {
260267
ip := getClientIP(r)
261268

262269
if !rl.Allow(ip) {
270+
// Record the rate-limited request if callback is configured
271+
if rl.config.OnRateLimited != nil {
272+
rl.config.OnRateLimited(ip)
273+
}
274+
263275
w.Header().Set("Content-Type", "application/problem+json")
264276
w.Header().Set("Retry-After", "60")
265277
w.WriteHeader(http.StatusTooManyRequests)

internal/api/ratelimit/ratelimit_test.go

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -419,3 +419,66 @@ func TestDefaultConfig(t *testing.T) {
419419
t.Errorf("missing skip paths: %v", expectedSkipPaths)
420420
}
421421
}
422+
423+
func TestOnRateLimitedCallback(t *testing.T) {
424+
var callbackCount int
425+
var lastIP string
426+
427+
cfg := ratelimit.Config{
428+
RequestsPerMinute: 1,
429+
RequestsPerHour: 100,
430+
CleanupInterval: time.Hour,
431+
SkipPaths: []string{},
432+
MaxVisitors: 1000,
433+
OnRateLimited: func(ip string) {
434+
callbackCount++
435+
lastIP = ip
436+
},
437+
}
438+
rl := ratelimit.New(cfg)
439+
defer rl.Stop()
440+
441+
handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
442+
w.WriteHeader(http.StatusOK)
443+
})
444+
445+
middleware := rl.Middleware(handler)
446+
447+
// First request should succeed, no callback
448+
req := httptest.NewRequest(http.MethodGet, "/test", nil)
449+
req.RemoteAddr = "192.168.1.100:12345"
450+
w := httptest.NewRecorder()
451+
middleware.ServeHTTP(w, req)
452+
453+
if callbackCount != 0 {
454+
t.Errorf("callback should not be called on allowed request, got %d calls", callbackCount)
455+
}
456+
457+
// Second request should be blocked, callback should fire
458+
req2 := httptest.NewRequest(http.MethodGet, "/test", nil)
459+
req2.RemoteAddr = "192.168.1.100:12346"
460+
w2 := httptest.NewRecorder()
461+
middleware.ServeHTTP(w2, req2)
462+
463+
if w2.Code != http.StatusTooManyRequests {
464+
t.Errorf("expected status %d, got %d", http.StatusTooManyRequests, w2.Code)
465+
}
466+
467+
if callbackCount != 1 {
468+
t.Errorf("callback should be called once, got %d calls", callbackCount)
469+
}
470+
471+
if lastIP != "192.168.1.100" {
472+
t.Errorf("expected IP 192.168.1.100, got %s", lastIP)
473+
}
474+
475+
// Third request also blocked, callback should fire again
476+
req3 := httptest.NewRequest(http.MethodGet, "/test", nil)
477+
req3.RemoteAddr = "192.168.1.100:12347"
478+
w3 := httptest.NewRecorder()
479+
middleware.ServeHTTP(w3, req3)
480+
481+
if callbackCount != 2 {
482+
t.Errorf("callback should be called twice, got %d calls", callbackCount)
483+
}
484+
}

internal/api/server.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,11 @@ func NewServer(cfg *config.Config, registryService service.RegistryService, metr
8181
CleanupInterval: 10 * time.Minute,
8282
SkipPaths: []string{"/health", "/ping", "/metrics"},
8383
MaxVisitors: 100000,
84+
OnRateLimited: func(_ string) {
85+
if metrics != nil {
86+
metrics.RateLimitedRequests.Add(context.Background(), 1)
87+
}
88+
},
8489
}
8590
rateLimiter = ratelimit.New(rateLimitConfig)
8691
handler = rateLimiter.Middleware(handler)

internal/telemetry/metrics.go

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,9 @@ type Metrics struct {
3232

3333
// Up tracks the health of the service
3434
Up metric.Int64Gauge
35+
36+
// RateLimitedRequests tracks requests blocked by rate limiting
37+
RateLimitedRequests metric.Int64Counter
3538
}
3639

3740
// ShutdownFunc is a delegate that shuts down the OpenTelemetry components.
@@ -73,11 +76,20 @@ func NewMetrics(meter metric.Meter) (*Metrics, error) {
7376
return nil, fmt.Errorf("failed to create service up gauge: %w", err)
7477
}
7578

79+
rateLimited, err := meter.Int64Counter(
80+
Namespace+".http.rate_limited",
81+
metric.WithDescription("Total number of requests blocked by rate limiting"),
82+
)
83+
if err != nil {
84+
return nil, fmt.Errorf("failed to create rate limited counter: %w", err)
85+
}
86+
7687
return &Metrics{
77-
Requests: req,
78-
RequestDuration: reqDuration,
79-
ErrorCount: errCount,
80-
Up: up,
88+
Requests: req,
89+
RequestDuration: reqDuration,
90+
ErrorCount: errCount,
91+
Up: up,
92+
RateLimitedRequests: rateLimited,
8193
}, nil
8294
}
8395

0 commit comments

Comments
 (0)