diff --git a/server/internal/operational/memory/limiter.go b/server/internal/operational/memory/limiter.go index acbc7f31ee..6f60f72c37 100644 --- a/server/internal/operational/memory/limiter.go +++ b/server/internal/operational/memory/limiter.go @@ -90,7 +90,11 @@ func (s *Store) TakeRateLimitTokens(ctx context.Context, key string, limit int, if !allowed { result.RetryAfter = time.Duration(math.Ceil((float64(cost)-v.tokens)*millis/float64(limit))) * time.Millisecond } - v.updated, v.expires = now, now.Add(time.Duration(period.Milliseconds()*2)*time.Millisecond) + // Like the PostgreSQL adapter, never move the refill timestamp backwards. + if now.After(v.updated) { + v.updated = now + } + v.expires = now.Add(time.Duration(period.Milliseconds()*2) * time.Millisecond) s.buckets[strings.Clone(key)] = v return result, nil } diff --git a/server/internal/operational/memory/limiter_test.go b/server/internal/operational/memory/limiter_test.go index 32376e7605..7a025664e8 100644 --- a/server/internal/operational/memory/limiter_test.go +++ b/server/internal/operational/memory/limiter_test.go @@ -66,3 +66,35 @@ func TestClassifyTokenBucketRequestOversizedLimit(t *testing.T) { t.Fatalf("classifyTokenBucketRequest with oversized limit = %s, want %s", got, tokenBucketRejectedParameter) } } + +func TestTakeRateLimitTokensClockRegressionDoesNotRefillTwice(t *testing.T) { + now := time.Unix(1000, 0) + s, err := New(Config{Clock: func() time.Time { return now }}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = s.Close() }) + take := func(cost int) bool { + t.Helper() + r, err := s.TakeRateLimitTokens(t.Context(), "subject", 10, 10*time.Second, cost) + if err != nil { + t.Fatal(err) + } + return r.Allowed + } + if !take(10) { + t.Fatal("initial full-bucket spend denied") + } + now = now.Add(5 * time.Second) + if !take(5) { + t.Fatal("refilled spend denied") + } + now = now.Add(-5 * time.Second) + if take(1) { + t.Fatal("empty bucket allowed a token after clock regression") + } + now = now.Add(5 * time.Second) + if take(1) { + t.Fatal("clock regression refilled the same elapsed interval twice") + } +}