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

@@ -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) {

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)
+	}
+}