x/time/rate: correctly handle 0 limits
Decrementing the burst in the reserveN method will frequently lead to us
setting the burst to 0 which makes the limiter mostly unusable.
This code was originally added in https://go.dev/cl/323429 to fix #39984
but the implementation introduced a different bug. To avoid regressing
to the behaviour described in #39984, pre-fill the limiter to the burst
value in the constructor.
Fixes #68541
Change-Id: Iab3b85d548a44fcb2d058336e5bbf11b19ea67b1
Reviewed-on: https://go-review.googlesource.com/c/time/+/600876
Reviewed-by: Sameer Ajmani <sameer@golang.org>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Michael Knyszek <mknyszek@google.com>
Auto-Submit: Sameer Ajmani <sameer@golang.org>
diff --git a/rate/rate.go b/rate/rate.go
index 8f6c7f4..93a798a 100644
--- a/rate/rate.go
+++ b/rate/rate.go
@@ -99,8 +99,9 @@
// bursts of at most b tokens.
func NewLimiter(r Limit, b int) *Limiter {
return &Limiter{
- limit: r,
- burst: b,
+ limit: r,
+ burst: b,
+ tokens: float64(b),
}
}
@@ -344,18 +345,6 @@
tokens: n,
timeToAct: t,
}
- } else if lim.limit == 0 {
- var ok bool
- if lim.burst >= n {
- ok = true
- lim.burst -= n
- }
- return Reservation{
- ok: ok,
- lim: lim,
- tokens: lim.burst,
- timeToAct: t,
- }
}
t, tokens := lim.advance(t)
diff --git a/rate/rate_test.go b/rate/rate_test.go
index 48ebaef..c91fc8f 100644
--- a/rate/rate_test.go
+++ b/rate/rate_test.go
@@ -618,3 +618,23 @@
t.Errorf("Limit(0, 1) want false when already used")
}
}
+
+func TestSetAfterZeroLimit(t *testing.T) {
+ lim := NewLimiter(0, 1)
+ // The limiter should start off full, so even though our rate limit is 0, our first request
+ // should be allowed…
+ if !lim.Allow() {
+ t.Errorf("Limit(0, 1) want true when first used")
+ }
+ // …the token bucket is not being replenished though, so the second request should not succeed
+ if lim.Allow() {
+ t.Errorf("Limit(0, 1) want false when already used")
+ }
+
+ lim.SetLimit(10)
+
+ tt := makeTestTime(t)
+
+ // We set the limit to 10/s so expect to get another token in 100ms
+ runWait(t, tt, lim, wait{"wait-after-set-nonzero-after-zero", context.Background(), 1, 1, true})
+}
