Merge pull request #9 from StudioLambda/fix/rate-limit-memory-eviction

Fix rate limiter memory leak and update defaults

Author: ConsoleTVs

AGENTS.md

@@ -133,7 +133,7 @@ Available in `framework/middleware`:
 - `CSRF(origins...)` / `CSRFWith(csrf, problem)` — cross-origin protection
 - `CORS(CORSOptions)` — configurable CORS headers
 - `SecureHeaders()` / `SecureHeadersWith(opts)` — security headers (HSTS, CSP, X-Frame-Options)
-- `RateLimit()` / `RateLimitWith(opts)` — per-key token bucket (default 10 req/s, burst 20)
+- `RateLimit()` / `RateLimitWith(opts)` — per-key token bucket (default 15 req/s, burst 30, idle eviction after 5m)
 - `Provide(key, value)` / `ProvideWith(fn)` — context injection
 - `HTTP(func(http.Handler) http.Handler)` — stdlib middleware adapter
 

framework/middleware/rate_limit.go

@@ -3,6 +3,7 @@ package middleware
 import (
 	"net/http"
 	"sync"
+	"time"
 
 	"github.com/studiolambda/cosmos/framework"
 	"github.com/studiolambda/cosmos/problem"
@@ -21,11 +22,11 @@ var ErrRateLimited = problem.Problem{
 // RateLimitOptions configures the rate limiter middleware.
 type RateLimitOptions struct {
 	// RequestsPerSecond is the sustained request rate allowed
-	// per key (typically per IP). Defaults to 10.
+	// per key (typically per IP). Defaults to 15.
 	RequestsPerSecond float64
 
 	// Burst is the maximum number of requests allowed in a
-	// single burst above the sustained rate. Defaults to 20.
+	// single burst above the sustained rate. Defaults to 30.
 	Burst int
 
 	// KeyFunc extracts the rate-limit key from a request.
@@ -35,17 +36,28 @@ type RateLimitOptions struct {
 	// ErrorResponse is the problem returned when a request is
 	// rate-limited. Defaults to [ErrRateLimited].
 	ErrorResponse problem.Problem
+
+	// CleanupInterval is how often the registry sweeps for
+	// idle entries. Defaults to 1 minute.
+	CleanupInterval time.Duration
+
+	// MaxIdleTime is how long an entry can be idle before
+	// being evicted. Defaults to 5 minutes.
+	MaxIdleTime time.Duration
 }
 
-// DefaultRateLimitOptions holds sensible defaults: 10 req/s
-// sustained with a burst of 20, keyed by remote address.
+// DefaultRateLimitOptions holds sensible defaults: 15 req/s
+// sustained with a burst of 30, keyed by remote address.
+// Idle entries are evicted after 5 minutes of inactivity.
 var DefaultRateLimitOptions = RateLimitOptions{
-	RequestsPerSecond: 10,
-	Burst:             20,
+	RequestsPerSecond: 15,
+	Burst:             30,
 	KeyFunc: func(r *http.Request) string {
 		return r.RemoteAddr
 	},
-	ErrorResponse: ErrRateLimited,
+	ErrorResponse:   ErrRateLimited,
+	CleanupInterval: 1 * time.Minute,
+	MaxIdleTime:     5 * time.Minute,
 }
 
 // withDefaults returns a copy of the options with zero values
@@ -67,18 +79,38 @@ func (options RateLimitOptions) withDefaults() RateLimitOptions {
 		options.ErrorResponse = DefaultRateLimitOptions.ErrorResponse
 	}
 
+	if options.CleanupInterval == 0 {
+		options.CleanupInterval = DefaultRateLimitOptions.CleanupInterval
+	}
+
+	if options.MaxIdleTime == 0 {
+		options.MaxIdleTime = DefaultRateLimitOptions.MaxIdleTime
+	}
+
 	return options
 }
 
-// rateLimitRegistry manages per-key token bucket rate limiters.
-// Each unique key gets its own [rate.Limiter] instance, created
-// on first access and reused for subsequent requests.
+// rateLimitEntry pairs a token bucket limiter with the time it
+// was last accessed. Entries idle longer than [RateLimitOptions.MaxIdleTime]
+// are evicted by the cleanup goroutine.
+type rateLimitEntry struct {
+	limiter  *rate.Limiter
+	lastSeen time.Time
+}
+
+// rateLimitRegistry manages per-key token bucket rate limiters
+// with automatic eviction of idle entries. Each unique key gets
+// its own [rate.Limiter] instance, created on first access and
+// reused for subsequent requests. A background goroutine
+// periodically removes entries that have been idle longer than
+// the configured maximum idle time.
 type rateLimitRegistry struct {
-	// mu protects concurrent access to the limiters map.
+	// mu protects concurrent access to the entries map.
 	mu sync.Mutex
 
-	// limiters maps rate-limit keys to their token bucket.
-	limiters map[string]*rate.Limiter
+	// entries maps rate-limit keys to their limiter and
+	// last-seen timestamp.
+	entries map[string]*rateLimitEntry
 
 	// rps is the sustained requests-per-second rate for
 	// newly created limiters.
@@ -87,46 +119,120 @@ type rateLimitRegistry struct {
 	// burst is the maximum burst size for newly created
 	// limiters.
 	burst int
+
+	// stop signals the cleanup goroutine to exit.
+	stop chan struct{}
 }
 
 // newRateLimitRegistry creates a registry that produces limiters
-// with the given sustained rate and burst size.
-func newRateLimitRegistry(rps float64, burst int) *rateLimitRegistry {
-	return &rateLimitRegistry{
-		limiters: make(map[string]*rate.Limiter),
-		rps:      rps,
-		burst:    burst,
+// with the given sustained rate and burst size. It starts a
+// background goroutine that evicts entries idle longer than
+// maxIdle at the given interval. Call [rateLimitRegistry.close]
+// to stop the goroutine.
+func newRateLimitRegistry(
+	rps float64,
+	burst int,
+	cleanupInterval time.Duration,
+	maxIdle time.Duration,
+) *rateLimitRegistry {
+	registry := &rateLimitRegistry{
+		entries: make(map[string]*rateLimitEntry),
+		rps:     rps,
+		burst:   burst,
+		stop:    make(chan struct{}),
 	}
+
+	go registry.cleanup(cleanupInterval, maxIdle)
+
+	return registry
 }
 
 // get returns the limiter for key, creating one if it does not
-// already exist.
+// already exist. It updates the entry's last-seen timestamp on
+// every access.
 func (registry *rateLimitRegistry) get(key string) *rate.Limiter {
 	registry.mu.Lock()
 	defer registry.mu.Unlock()
 
-	if limiter, ok := registry.limiters[key]; ok {
-		return limiter
+	if entry, ok := registry.entries[key]; ok {
+		entry.lastSeen = time.Now()
+
+		return entry.limiter
 	}
 
 	limiter := rate.NewLimiter(rate.Limit(registry.rps), registry.burst)
-	registry.limiters[key] = limiter
+
+	registry.entries[key] = &rateLimitEntry{
+		limiter:  limiter,
+		lastSeen: time.Now(),
+	}
 
 	return limiter
 }
 
-// RateLimit returns middleware that limits requests to 10 req/s
-// per IP with a burst of 20 using [DefaultRateLimitOptions].
+// size returns the number of entries in the registry.
+func (registry *rateLimitRegistry) size() int {
+	registry.mu.Lock()
+	defer registry.mu.Unlock()
+
+	return len(registry.entries)
+}
+
+// cleanup periodically removes entries that have been idle
+// longer than maxIdle. It runs until [rateLimitRegistry.close]
+// is called.
+func (registry *rateLimitRegistry) cleanup(
+	interval time.Duration,
+	maxIdle time.Duration,
+) {
+	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-registry.stop:
+			return
+		case now := <-ticker.C:
+			registry.mu.Lock()
+
+			for key, entry := range registry.entries {
+				if now.Sub(entry.lastSeen) > maxIdle {
+					delete(registry.entries, key)
+				}
+			}
+
+			registry.mu.Unlock()
+		}
+	}
+}
+
+// close stops the background cleanup goroutine.
+func (registry *rateLimitRegistry) close() {
+	close(registry.stop)
+}
+
+// RateLimit returns middleware that limits requests to 15 req/s
+// per IP with a burst of 30 using [DefaultRateLimitOptions].
+// Idle entries are automatically evicted after 5 minutes.
 func RateLimit() framework.Middleware {
 	return RateLimitWith(DefaultRateLimitOptions)
 }
 
 // RateLimitWith returns middleware that limits requests using
 // the provided options. It uses a per-key token bucket algorithm
-// backed by [golang.org/x/time/rate].
+// backed by [golang.org/x/time/rate]. A background goroutine
+// periodically evicts entries that have been idle longer than
+// [RateLimitOptions.MaxIdleTime] to prevent unbounded memory
+// growth.
 func RateLimitWith(opts RateLimitOptions) framework.Middleware {
 	opts = opts.withDefaults()
-	registry := newRateLimitRegistry(opts.RequestsPerSecond, opts.Burst)
+
+	registry := newRateLimitRegistry(
+		opts.RequestsPerSecond,
+		opts.Burst,
+		opts.CleanupInterval,
+		opts.MaxIdleTime,
+	)
 
 	return func(next framework.Handler) framework.Handler {
 		return func(w http.ResponseWriter, r *http.Request) error {