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 <sameer@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
diff --git a/rate/rate.go b/rate/rate.go
index 26ed8b3..feab629 100644
--- a/rate/rate.go
+++ b/rate/rate.go
@@ -221,8 +221,9 @@
 // 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
diff --git a/rate/rate_test.go b/rate/rate_test.go
index a7c0285..cf45d92 100644
--- a/rate/rate_test.go
+++ b/rate/rate_test.go
@@ -396,7 +396,7 @@
 	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 @@
 	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()