rate: ignore burst in WaitN when rate limit is Inf

WaitN fails when the number of tokens requested exceeds the limiter's
burst size, even when the rate limit is Inf.  The documented behavior is
that the burst size is ignored when the rate limit is Inf.  Change WaitN
to conform with the documentation, and add tests.

Fixes golang/go#16854

Change-Id: I3ca0a8dac47641c76bed4258cd7010f0a704ff8c
Reviewed-on: https://go-review.googlesource.com/28610
Run-TryBot: Sameer Ajmani <[email protected]>
TryBot-Result: Gobot Gobot <[email protected]>
Reviewed-by: Brad Fitzpatrick <[email protected]>

Author: Sajmani

rate/rate.go

@@ -221,8 +221,9 @@ func (lim *Limiter) Wait(ctx context.Context) (err error) {
 // WaitN blocks until lim permits n events to happen.
 // It returns an error if n exceeds the Limiter's burst size, the Context is
 // canceled, or the expected wait time exceeds the Context's Deadline.
+// The burst limit is ignored if the rate limit is Inf.
 func (lim *Limiter) WaitN(ctx context.Context, n int) (err error) {
-	if n > lim.burst {
+	if n > lim.burst && lim.limit != Inf {
 		return fmt.Errorf("rate: Wait(n=%d) exceeds limiter's burst %d", n, lim.burst)
 	}
 	// Check if ctx is already cancelled

rate/rate_test.go

@@ -396,7 +396,7 @@ func TestWaitSimple(t *testing.T) {
 	cancel()
 	runWait(t, lim, wait{"already-cancelled", ctx, 1, 0, false})
 
-	runWait(t, lim, wait{"n-gt-burst", context.Background(), 4, 0, false})
+	runWait(t, lim, wait{"exceed-burst-error", context.Background(), 4, 0, false})
 
 	runWait(t, lim, wait{"act-now", context.Background(), 2, 0, true})
 	runWait(t, lim, wait{"act-later", context.Background(), 3, 2, true})
@@ -426,6 +426,12 @@ func TestWaitTimeout(t *testing.T) {
 	runWait(t, lim, wait{"w-timeout-err", ctx, 3, 0, false})
 }
 
+func TestWaitInf(t *testing.T) {
+	lim := NewLimiter(Inf, 0)
+
+	runWait(t, lim, wait{"exceed-burst-no-error", context.Background(), 3, 0, true})
+}
+
 func BenchmarkAllowN(b *testing.B) {
 	lim := NewLimiter(Every(1*time.Second), 1)
 	now := time.Now()