fix: suppress canceled rate limit logs (#21)

Author: garrettladley

internal/middleware/ratelimit.go

@@ -1,8 +1,10 @@
 package middleware
 
 import (
+	"context"
 	"crypto/sha256"
 	"encoding/hex"
+	"errors"
 	"fmt"
 	"log/slog"
 	"net"
@@ -39,6 +41,10 @@ func RateLimit(store kv.Store, cfg config.RateLimit, logger *slog.Logger) Middle
 			key := rateLimitKey(ip, cfg.Window, now)
 			count, err := store.Increment(r.Context(), key, cfg.Window+time.Second)
 			if err != nil {
+				if errors.Is(err, context.Canceled) {
+					trace.SpanFromContext(r.Context()).SetAttributes(observability.RateLimitAttrs{Outcome: observability.RateLimitOutcomeCanceled, Limit: cfg.Requests, Window: cfg.Window}.Attributes()...)
+					return
+				}
 				trace.SpanFromContext(r.Context()).SetAttributes(observability.RateLimitAttrs{Outcome: observability.RateLimitOutcomeStoreError, Limit: cfg.Requests, Window: cfg.Window}.Attributes()...)
 				logger.ErrorContext(r.Context(), "rate limit check failed", slog.Any("error", err), slog.String("client_ip", ip))
 				http.Error(w, http.StatusText(http.StatusServiceUnavailable), http.StatusServiceUnavailable)

internal/middleware/ratelimit_test.go

@@ -3,9 +3,11 @@ package middleware
 import (
 	"context"
 	"errors"
+	"log/slog"
 	"net/http"
 	"net/http/httptest"
 	"strconv"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -78,6 +80,33 @@ func TestRateLimitFailsClosedOnStoreError(t *testing.T) {
 	}
 }
 
+func TestRateLimitDoesNotLogCanceledRequestsAsStoreErrors(t *testing.T) {
+	t.Parallel()
+
+	var errorLogs atomic.Int64
+	logger := slog.New(countingErrorHandler{count: &errorLogs})
+	store := incrementFunc(func(ctx context.Context, _ string, _ time.Duration) (int64, error) {
+		return 0, ctx.Err()
+	})
+	handler := RateLimit(store, config.RateLimit{Enabled: true, Requests: 2, Window: time.Minute}, logger)(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusNoContent)
+	}))
+	ctx, cancel := context.WithCancel(t.Context())
+	cancel()
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, "http://example.test/mcp", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	req.RemoteAddr = "203.0.113.10:1234"
+
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	if got := errorLogs.Load(); got != 0 {
+		t.Fatalf("error logs = %d, want 0", got)
+	}
+}
+
 func TestClientIPPrefersFlyHeaderAndNormalizes(t *testing.T) {
 	t.Parallel()
 
@@ -114,3 +143,28 @@ func (f incrementFunc) Set(context.Context, string, []byte, time.Duration) error
 func (f incrementFunc) Increment(ctx context.Context, key string, ttl time.Duration) (int64, error) {
 	return f(ctx, key, ttl)
 }
+
+type countingErrorHandler struct {
+	count *atomic.Int64
+}
+
+var _ slog.Handler = countingErrorHandler{}
+
+func (h countingErrorHandler) Enabled(_ context.Context, level slog.Level) bool {
+	return level >= slog.LevelError
+}
+
+func (h countingErrorHandler) Handle(_ context.Context, record slog.Record) error {
+	if record.Level >= slog.LevelError {
+		h.count.Add(1)
+	}
+	return nil
+}
+
+func (h countingErrorHandler) WithAttrs([]slog.Attr) slog.Handler {
+	return h
+}
+
+func (h countingErrorHandler) WithGroup(string) slog.Handler {
+	return h
+}

internal/observability/attrs.go

@@ -128,6 +128,7 @@ const (
 	RateLimitOutcomeAllowed    RateLimitOutcome = "allowed"
 	RateLimitOutcomeLimited    RateLimitOutcome = "limited"
 	RateLimitOutcomeStoreError RateLimitOutcome = "store_error"
+	RateLimitOutcomeCanceled   RateLimitOutcome = "canceled"
 	RateLimitOutcomeDisabled   RateLimitOutcome = "disabled"
 	RateLimitOutcomeSkipped    RateLimitOutcome = "skipped"
 )