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