Add LRU eviction to UberRateLimiter

README.md

@@ -518,6 +518,8 @@ See the `examples/middleware` directory for examples of using IP configuration.
 
 SRouter provides a flexible rate limiting system using `common.RateLimitConfig` (defined in `pkg/common/types.go`) that can be configured at the global, sub-router, or route level. Rate limits can be based on IP address, authenticated user, or custom criteria. Under the hood, SRouter uses [Uber's ratelimit library](https://github.com/uber-go/ratelimit) via the `middleware.RateLimit` function for efficient and smooth rate limiting with a leaky bucket algorithm.
 
+`UberRateLimiter` stores buckets in an in-memory LRU cache to bound memory usage. The cache size defaults to `10,000` entries and can be customized via `NewUberRateLimiterWithMax`.
+
 #### Rate Limiting Configuration
 
 ```go

docs/security.md

@@ -79,13 +79,14 @@ The rate limiting middleware in SRouter supports several strategies for identify
 
 ### Memory Consumption Considerations
 
-SRouter's default rate limiter, `UberRateLimiter` (which uses `go.uber.org/ratelimit`), stores rate limiting state in memory for each unique key (e.g., each IP address or user ID).
+SRouter's default rate limiter, `UberRateLimiter` (which uses `go.uber.org/ratelimit`), stores rate limiting state in memory for each unique key (e.g., each IP address or user ID). To prevent unbounded growth, it now keeps these limiters in an in-memory LRU cache with a configurable maximum size (default `10,000` entries).
 
 *   **High Cardinality Warning:** If your application anticipates a very large number of unique keys (e.g., millions of different IP addresses making requests over a short period, or a vast number of user accounts being actively rate-limited simultaneously) without application restarts, this in-memory storage can lead to significant memory consumption.
 *   **Suitability:** The current `UberRateLimiter` is well-suited for many common use cases, especially for applications with a moderate number of active users or for services deployed behind a load balancer that already handles a significant portion of traffic.
 *   **Alternatives for Extreme Scale:** For extremely high-traffic scenarios with very high cardinality of rate limiting keys, or if you require persistence of rate limiting state across application restarts or a distributed environment, consider:
+    *   Increasing or decreasing the LRU cache size via `NewUberRateLimiterWithMax` when memory constraints or key cardinality require it.
     *   Using an external rate-limiting solution (e.g., Redis-based, or dedicated proxy/API gateway rate limiters).
-    *   Implementing a custom `RateLimiter` for SRouter that uses a backend store with eviction policies (e.g., LRU cache, Redis with TTLs) to manage memory usage.
+    *   Implementing a custom `RateLimiter` for SRouter that uses a backend store with eviction policies (e.g., Redis with TTLs) to manage memory usage.
 
 ### Standard Rate Limit Headers
 

pkg/middleware/ratelimit.go

@@ -2,6 +2,7 @@
 package middleware
 
 import (
+	"container/list"
 	"errors" // Added for error handling
 	"fmt"
 	"net/http"
@@ -22,36 +23,68 @@
 // by allowing a steady flow of requests while preventing bursts.
 // The implementation maintains a map of rate limiters, one per unique key.
 type UberRateLimiter struct {
-	limiters sync.Map // map[string]ratelimit.Limiter
+	mu         sync.Mutex
+	limiters   map[string]*list.Element // LRU cache mapping
+	lru        *list.List               // order of keys, most recent at front
+	maxEntries int
 }
 
-// NewUberRateLimiter creates a new UberRateLimiter instance.
-// The returned limiter uses the leaky bucket algorithm to enforce rate limits.
-// It maintains separate rate limiters for different keys (e.g., different IPs or users).
+type limiterEntry struct {
+	key     string
+	limiter ratelimit.Limiter
+}
+
+// NewUberRateLimiter creates a new UberRateLimiter instance using a default cache size.
+// The returned limiter uses the leaky bucket algorithm to enforce rate limits and
+// keeps recently used limiters in an LRU cache to bound memory usage.
+const defaultLimiterCacheSize = 10000
+
 func NewUberRateLimiter() *UberRateLimiter {
-	return &UberRateLimiter{}
+	return NewUberRateLimiterWithMax(defaultLimiterCacheSize)
+}
+
+// NewUberRateLimiterWithMax creates a new UberRateLimiter instance with a
+// maximum number of limiter entries to keep in memory. If maxEntries is <= 0
+// a reasonable default is used.
+func NewUberRateLimiterWithMax(maxEntries int) *UberRateLimiter {
+	if maxEntries <= 0 {
+		maxEntries = defaultLimiterCacheSize
+	}
+	return &UberRateLimiter{
+		limiters:   make(map[string]*list.Element),
+		lru:        list.New(),
+		maxEntries: maxEntries,
+	}
 }
 
 // getLimiter gets or creates a limiter for the given key and rate (requests per second).
 // It uses a composite key including the RPS to handle different rate limits for the same base key.
 func (u *UberRateLimiter) getLimiter(key string, rps int) ratelimit.Limiter {
-	compositeKey := fmt.Sprintf("%s-%d", key, rps) // Combine key and rps
+	compositeKey := fmt.Sprintf("%s-%d", key, rps)
+
+	u.mu.Lock()
+	defer u.mu.Unlock()
 
-	// Fast path: Check if limiter already exists.
-	if limiter, ok := u.limiters.Load(compositeKey); ok {
-		return limiter.(ratelimit.Limiter)
+	if elem, ok := u.limiters[compositeKey]; ok {
+		u.lru.MoveToFront(elem)
+		return elem.Value.(*limiterEntry).limiter
 	}
 
-	// Slow path: Limiter doesn't exist, create a new one.
 	newLimiter := ratelimit.New(rps)
 
-	// Atomically load or store.
-	// - If compositeKey already exists (due to concurrent creation), LoadOrStore loads and returns the existing value.
-	// - If compositeKey doesn't exist, LoadOrStore stores newLimiter and returns it.
-	actualLimiter, _ := u.limiters.LoadOrStore(compositeKey, newLimiter)
+	// Evict oldest if we exceed the maximum size
+	if u.lru.Len() >= u.maxEntries {
+		if back := u.lru.Back(); back != nil {
+			ent := back.Value.(*limiterEntry)
+			delete(u.limiters, ent.key)
+			u.lru.Remove(back)
+		}
+	}
 
-	// Return the actual limiter stored in the map (either the existing one or the new one).
-	return actualLimiter.(ratelimit.Limiter)
+	ent := &limiterEntry{key: compositeKey, limiter: newLimiter}
+	elem := u.lru.PushFront(ent)
+	u.limiters[compositeKey] = elem
+	return newLimiter
 }
 
 // Ensure UberRateLimiter implements the common.RateLimiter interface.

pkg/middleware/ratelimit_coverage_test.go

@@ -66,16 +66,16 @@ func TestGetLimiter_Coverage(t *testing.T) {
 				// Check if the same limiter instance is returned across goroutines
 				// This indirectly tests that the double-check prevents creating multiple limiters
 				compositeKey := fmt.Sprintf("%s-%d", key, rps)
-				val, ok := limiter.limiters.Load(compositeKey)
-				// Explicitly check 'ok' before proceeding to prevent panic on nil interface conversion
+				limiter.mu.Lock()
+				elem, ok := limiter.limiters[compositeKey]
+				var currentLimiter ratelimit.Limiter
+				if ok {
+					currentLimiter = elem.Value.(*limiterEntry).limiter
+				}
+				limiter.mu.Unlock()
 				if !assert.True(t, ok, "Limiter should exist in the map for key %s", compositeKey) {
-					// If the assertion fails (ok is false), skip the rest of the checks for this iteration
-					// as val will be nil, causing a panic on type assertion.
-					// This indicates a potential timing issue or problem in limiter creation/storage.
-					continue // Skip to the next iteration of the inner loop
+					continue
 				}
-				// Only proceed if ok is true
-				currentLimiter := val.(ratelimit.Limiter)
 
 				once.Do(func() {
 					firstLimiter = currentLimiter // Capture the first successfully retrieved limiter
@@ -91,12 +91,13 @@ func TestGetLimiter_Coverage(t *testing.T) {
 
 	// Final check that only one limiter was created for the key/rps combination
 	count := 0
-	limiter.limiters.Range(func(k, v interface{}) bool {
+	limiter.mu.Lock()
+	for k := range limiter.limiters {
 		if k == fmt.Sprintf("%s-%d", key, rps) {
 			count++
 		}
-		return true
-	})
+	}
+	limiter.mu.Unlock()
 	assert.Equal(t, 1, count, "Expected exactly one limiter instance for the key/rps")
 }
 
@@ -227,7 +228,7 @@ func TestUberRateLimiter_ForceNegativeRemaining(t *testing.T) {
 
 	// Override the Allow method to force a negative remaining calculation
 	// This directly tests the if remaining < 0 { remaining = 0 } code path
-	mockLimiter := &mockLimiterWithNegativeRemaining{}
+	mockLimiter := &mockLimiterWithNegativeRemaining{*NewUberRateLimiter()}
 
 	// Create a test scenario where the calculation would result in negative remaining
 	key := "negative-remaining-test"

pkg/middleware/ratelimit_test.go

@@ -41,7 +41,9 @@ func TestUberRateLimiter(t *testing.T) {
 		rps1 = 1
 	}
 	compositeKey1 := fmt.Sprintf("%s-%d", key, rps1)
-	_, exists := limiter.limiters.Load(compositeKey1)
+	limiter.mu.Lock()
+	_, exists := limiter.limiters[compositeKey1]
+	limiter.mu.Unlock()
 	if !exists {
 		t.Errorf("Expected limiter to be stored for composite key %s", compositeKey1)
 	}
@@ -63,7 +65,9 @@ func TestUberRateLimiter(t *testing.T) {
 		rps2 = 1
 	}
 	compositeKey2 := fmt.Sprintf("%s-%d", otherKey, rps2)
-	_, exists = limiter.limiters.Load(compositeKey2)
+	limiter.mu.Lock()
+	_, exists = limiter.limiters[compositeKey2]
+	limiter.mu.Unlock()
 	if !exists {
 		t.Errorf("Expected limiter to be stored for composite key %s", compositeKey2)
 	}
@@ -85,7 +89,9 @@ func TestUberRateLimiter(t *testing.T) {
 		rps3 = 1
 	}
 	compositeKey3 := fmt.Sprintf("%s-%d", key, rps3)
-	_, exists = limiter.limiters.Load(compositeKey3)
+	limiter.mu.Lock()
+	_, exists = limiter.limiters[compositeKey3]
+	limiter.mu.Unlock()
 	if !exists {
 		t.Errorf("Expected limiter to be stored for composite key %s (different limit/window)", compositeKey3)
 	}
@@ -770,3 +776,45 @@ func TestRateLimitMiddlewareDefaultStrategy(t *testing.T) {
 		t.Errorf("Expected Retry-After header to be set")
 	}
 }
+
+func TestUberRateLimiterEviction(t *testing.T) {
+	limiter := NewUberRateLimiterWithMax(2)
+	key1 := "k1"
+	key2 := "k2"
+	key3 := "k3"
+	limit := 1
+	window := time.Second
+
+	limiter.Allow(key1, limit, window)
+	limiter.Allow(key2, limit, window)
+
+	limiter.mu.Lock()
+	if len(limiter.limiters) != 2 {
+		t.Fatalf("expected 2 limiters, got %d", len(limiter.limiters))
+	}
+	limiter.mu.Unlock()
+
+	limiter.Allow(key3, limit, window)
+
+	rps := int(float64(limit) / window.Seconds())
+	c1 := fmt.Sprintf("%s-%d", key1, rps)
+	c2 := fmt.Sprintf("%s-%d", key2, rps)
+	c3 := fmt.Sprintf("%s-%d", key3, rps)
+
+	limiter.mu.Lock()
+	_, ok1 := limiter.limiters[c1]
+	_, ok2 := limiter.limiters[c2]
+	_, ok3 := limiter.limiters[c3]
+	count := len(limiter.limiters)
+	limiter.mu.Unlock()
+
+	if count != 2 {
+		t.Fatalf("limiter count expected 2, got %d", count)
+	}
+	if ok1 {
+		t.Errorf("oldest entry was not evicted")
+	}
+	if !ok2 || !ok3 {
+		t.Errorf("expected newest entries to remain")
+	}
+}