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 <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>
Reviewed-by: Michael Knyszek <[email protected]>
Auto-Submit: Sameer Ajmani <[email protected]>

Author: leesio

rate/rate.go

@@ -99,8 +99,9 @@ func (lim *Limiter) Tokens() float64 {
 // 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 @@ func (lim *Limiter) reserveN(t time.Time, n int, maxFutureReserve time.Duration)
 			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)

rate/rate_test.go

@@ -618,3 +618,23 @@ func TestZeroLimit(t *testing.T) {
 		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})
+}