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