From d3e6fd6aa83f18624ce462f3886663f87e5f16eb Mon Sep 17 00:00:00 2001
From: ZverGuy <maximbel2003@gmail.com>
Date: Thu, 5 Nov 2026 00:25:07 +0700
Subject: [PATCH] fix(observability-proxy): keep the rate-limit buckets under
 their bound

When the bucket map was full and the sweep freed nothing, the new
user's bucket was added anyway, so the map grew past its bound and
every further newcomer paid a full scan under the limiter lock.

A newcomer is now refused with 429 until the first bucket refills,
and the Retry-After tells it when. Resetting the map instead would
hand every throttled user a fresh burst.

Assisted-by: LLM
Signed-off-by: ZverGuy <maximbel2003@gmail.com>
---
 internal/observabilityproxy/limits.go      | 21 +++++++++++++----
 internal/observabilityproxy/limits_test.go | 27 ++++++++++++++++++++++
 2 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/internal/observabilityproxy/limits.go b/internal/observabilityproxy/limits.go
index eaebe3104..06cc41d41 100644
--- a/internal/observabilityproxy/limits.go
+++ b/internal/observabilityproxy/limits.go
@@ -72,7 +72,12 @@ func (l *limiter) allowRate(user string) (time.Duration, bool) {
 	b, ok := l.buckets[user]
 	if !ok {
 		if len(l.buckets) >= ttlCacheMaxEntries {
-			l.sweepBuckets(now, burst)
+			// Every bucket still refilling means that many users within one
+			// refill period; a newcomer waits for the first to drop out
+			// rather than grow the map past its bound.
+			if soonest := l.sweepBuckets(now, burst); len(l.buckets) >= ttlCacheMaxEntries {
+				return soonest, false
+			}
 		}
 		b = &bucket{tokens: burst, last: now}
 		l.buckets[user] = b
@@ -91,13 +96,21 @@ func (l *limiter) allowRate(user string) (time.Duration, bool) {
 }
 
 // sweepBuckets drops buckets that have refilled completely; recreating one
-// later starts it full, which is the same state.
-func (l *limiter) sweepBuckets(now time.Time, burst float64) {
+// later starts it full, which is the same state. It returns how long until
+// the first remaining bucket is full.
+func (l *limiter) sweepBuckets(now time.Time, burst float64) time.Duration {
+	soonest := time.Duration(math.MaxInt64)
 	for user, b := range l.buckets {
-		if b.tokens+now.Sub(b.last).Seconds()*l.cfg.RatePerUser >= burst {
+		missing := burst - (b.tokens + now.Sub(b.last).Seconds()*l.cfg.RatePerUser)
+		if missing <= 0 {
 			delete(l.buckets, user)
+			continue
+		}
+		if wait, ok := scaleDuration(missing/l.cfg.RatePerUser, time.Second); ok {
+			soonest = min(soonest, max(wait, time.Nanosecond))
 		}
 	}
+	return soonest
 }
 
 func (l *limiter) acquire(user string) (func(), bool) {
diff --git a/internal/observabilityproxy/limits_test.go b/internal/observabilityproxy/limits_test.go
index f7ff45d3a..706ed71c7 100644
--- a/internal/observabilityproxy/limits_test.go
+++ b/internal/observabilityproxy/limits_test.go
@@ -17,6 +17,7 @@ limitations under the License.
 package observabilityproxy
 
 import (
+	"fmt"
 	"testing"
 	"time"
 )
@@ -116,3 +117,29 @@ func TestLimiterRateWaitAtTinyRate(t *testing.T) {
 		t.Fatalf("empty bucket at a tiny rate: admitted %v, wait %s, want refused with a positive wait", ok, wait)
 	}
 }
+
+// TestLimiterRateBucketsBounded fills the bucket map with users that are
+// all still refilling: a new user must be refused rather than grow it, and
+// admitted once a bucket can be dropped.
+func TestLimiterRateBucketsBounded(t *testing.T) {
+	now := time.Unix(1000, 0)
+	l := newLimiter(LimiterConfig{RatePerUser: 2, RateBurst: 2})
+	l.now = func() time.Time { return now }
+	for i := range ttlCacheMaxEntries {
+		if !allowed(l, fmt.Sprintf("user-%d", i)) {
+			t.Fatalf("user-%d refused", i)
+		}
+		now = now.Add(time.Microsecond)
+	}
+	wait, ok := l.allowRate("newcomer")
+	if ok || wait <= 0 {
+		t.Fatalf("newcomer on a full map: admitted %v, wait %s; want refused with a positive wait", ok, wait)
+	}
+	if n := len(l.buckets); n > ttlCacheMaxEntries {
+		t.Fatalf("bucket map grew to %d, above %d", n, ttlCacheMaxEntries)
+	}
+	now = now.Add(wait)
+	if !allowed(l, "newcomer") {
+		t.Fatalf("newcomer still refused after the advertised wait %s", wait)
+	}
+}