x/time/rate: provides a rate limiter.
The rate limiter works with golang.org/x/net/context's cancelation
mechanism: the Wait method blocks until the limiter permits the
operation to proceed or the context is canceled (in which case the
requested rate allocation is remitted for use by other operations).
Co-author: Arkadi Pyuro <arkadi@google.com>
Change-Id: I841db4e61fd169ac119fdc43105a1e16ca97e0a2
Reviewed-on: https://go-review.googlesource.com/16672
Reviewed-by: Russ Cox <rsc@golang.org>
diff --git a/rate/rate.go b/rate/rate.go
new file mode 100644
index 0000000..2131b92
--- /dev/null
+++ b/rate/rate.go
@@ -0,0 +1,368 @@
+// Copyright 2015 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package rate provides a rate limiter.
+package rate
+
+import (
+ "fmt"
+ "math"
+ "sync"
+ "time"
+
+ "golang.org/x/net/context"
+)
+
+// Limit defines the maximum frequency of some events.
+// Limit is represented as number of events per second.
+// A zero Limit allows no events.
+type Limit float64
+
+// Inf is the infinite rate limit; it allows all events (even if burst is zero).
+const Inf = Limit(math.MaxFloat64)
+
+// Every converts a minimum time interval between events to a Limit.
+func Every(interval time.Duration) Limit {
+ if interval <= 0 {
+ return Inf
+ }
+ return 1 / Limit(interval.Seconds())
+}
+
+// A Limiter controls how frequently events are allowed to happen.
+// It implements a "token bucket" of size b, initially full and refilled
+// at rate r tokens per second.
+// Informally, in any large enough time interval, the Limiter limits the
+// rate to r tokens per second, with a maximum burst size of b events.
+// As a special case, if r == Inf (the infinite rate), b is ignored.
+// See https://en.wikipedia.org/wiki/Token_bucket for more about token buckets.
+//
+// The zero value is a valid Limiter, but it will reject all events.
+// Use NewLimiter to create non-zero Limiters.
+//
+// Limiter has three main methods, Allow, Reserve, and Wait.
+// Most callers should use Wait.
+//
+// Each of the three methods consumes a single token.
+// They differ in their behavior when no token is available.
+// If no token is available, Allow returns false.
+// If no token is available, Reserve returns a reservation for a future token
+// and the amount of time the caller must wait before using it.
+// If no token is available, Wait blocks until one can be obtained
+// or its associated context.Context is canceled.
+//
+// The methods AllowN, ReserveN, and WaitN consume n tokens.
+type Limiter struct {
+ limit Limit
+ burst int
+
+ mu sync.Mutex
+ tokens float64
+ // last is the last time the limiter's tokens field was updated
+ last time.Time
+ // lastEvent is the latest time of a rate-limited event (past or future)
+ lastEvent time.Time
+}
+
+// Limit returns the maximum overall event rate.
+func (lim *Limiter) Limit() Limit {
+ lim.mu.Lock()
+ defer lim.mu.Unlock()
+ return lim.limit
+}
+
+// Burst returns the maximum burst size. Burst is the maximum number of tokens
+// that can be consumed in a single call to Allow, Reserve, or Wait, so higher
+// Burst values allow more events to happen at once.
+// A zero Burst allows no events, unless limit == Inf.
+func (lim *Limiter) Burst() int {
+ return lim.burst
+}
+
+// NewLimiter returns a new Limiter that allows events up to rate r and permits
+// bursts of at most b tokens.
+func NewLimiter(r Limit, b int) *Limiter {
+ return &Limiter{
+ limit: r,
+ burst: b,
+ }
+}
+
+// Allow is shorthand for AllowN(time.Now(), 1).
+func (lim *Limiter) Allow() bool {
+ return lim.AllowN(time.Now(), 1)
+}
+
+// AllowN reports whether n events may happen at time now.
+// Use this method if you intend to drop / skip events that exceed the rate limit.
+// Otherwise use Reserve or Wait.
+func (lim *Limiter) AllowN(now time.Time, n int) bool {
+ return lim.reserveN(now, n, 0).ok
+}
+
+// A Reservation holds information about events that are permitted by a Limiter to happen after a delay.
+// A Reservation may be canceled, which may enable the Limiter to permit additional events.
+type Reservation struct {
+ ok bool
+ lim *Limiter
+ tokens int
+ timeToAct time.Time
+ // This is the Limit at reservation time, it can change later.
+ limit Limit
+}
+
+// OK returns whether the limiter can provide the requested number of tokens
+// within the maximum wait time. If OK is false, Delay returns InfDuration, and
+// Cancel does nothing.
+func (r *Reservation) OK() bool {
+ return r.ok
+}
+
+// Delay is shorthand for DelayFrom(time.Now()).
+func (r *Reservation) Delay() time.Duration {
+ return r.DelayFrom(time.Now())
+}
+
+// InfDuration is the duration returned by Delay when a Reservation is not OK.
+const InfDuration = time.Duration(1<<63 - 1)
+
+// DelayFrom returns the duration for which the reservation holder must wait
+// before taking the reserved action. Zero duration means act immediately.
+// InfDuration means the limiter cannot grant the tokens requested in this
+// Reservation within the maximum wait time.
+func (r *Reservation) DelayFrom(now time.Time) time.Duration {
+ if !r.ok {
+ return InfDuration
+ }
+ delay := r.timeToAct.Sub(now)
+ if delay < 0 {
+ return 0
+ }
+ return delay
+}
+
+// Cancel is shorthand for CancelAt(time.Now()).
+func (r *Reservation) Cancel() {
+ r.CancelAt(time.Now())
+ return
+}
+
+// CancelAt indicates that the reservation holder will not perform the reserved action
+// and reverses the effects of this Reservation on the rate limit as much as possible,
+// considering that other reservations may have already been made.
+func (r *Reservation) CancelAt(now time.Time) {
+ if !r.ok {
+ return
+ }
+
+ r.lim.mu.Lock()
+ defer r.lim.mu.Unlock()
+
+ if r.lim.limit == Inf || r.tokens == 0 || r.timeToAct.Before(now) {
+ return
+ }
+
+ // calculate tokens to restore
+ // The duration between lim.lastEvent and r.timeToAct tells us how many tokens were reserved
+ // after r was obtained. These tokens should not be restored.
+ restoreTokens := float64(r.tokens) - r.limit.tokensFromDuration(r.lim.lastEvent.Sub(r.timeToAct))
+ if restoreTokens <= 0 {
+ return
+ }
+ // advance time to now
+ now, _, tokens := r.lim.advance(now)
+ // calculate new number of tokens
+ tokens += restoreTokens
+ if burst := float64(r.lim.burst); tokens > burst {
+ tokens = burst
+ }
+ // update state
+ r.lim.last = now
+ r.lim.tokens = tokens
+ if r.timeToAct == r.lim.lastEvent {
+ prevEvent := r.timeToAct.Add(r.limit.durationFromTokens(float64(-r.tokens)))
+ if !prevEvent.Before(now) {
+ r.lim.lastEvent = prevEvent
+ }
+ }
+
+ return
+}
+
+// Reserve is shorthand for ReserveN(time.Now(), 1).
+func (lim *Limiter) Reserve() *Reservation {
+ return lim.ReserveN(time.Now(), 1)
+}
+
+// ReserveN returns a Reservation that indicates how long the caller must wait before n events happen.
+// The Limiter takes this Reservation into account when allowing future events.
+// ReserveN returns false if n exceeds the Limiter's burst size.
+// Usage example:
+// r, ok := lim.ReserveN(time.Now(), 1)
+// if !ok {
+// // Not allowed to act! Did you remember to set lim.burst to be > 0 ?
+// }
+// time.Sleep(r.Delay())
+// Act()
+// Use this method if you wish to wait and slow down in accordance with the rate limit without dropping events.
+// If you need to respect a deadline or cancel the delay, use Wait instead.
+// To drop or skip events exceeding rate limit, use Allow instead.
+func (lim *Limiter) ReserveN(now time.Time, n int) *Reservation {
+ r := lim.reserveN(now, n, InfDuration)
+ return &r
+}
+
+// Wait is shorthand for WaitN(ctx, 1).
+func (lim *Limiter) Wait(ctx context.Context) (err error) {
+ return lim.WaitN(ctx, 1)
+}
+
+// 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.
+func (lim *Limiter) WaitN(ctx context.Context, n int) (err error) {
+ if n > lim.burst {
+ return fmt.Errorf("rate: Wait(n=%d) exceeds limiter's burst %d", n, lim.burst)
+ }
+ // Check if ctx is already cancelled
+ select {
+ case <-ctx.Done():
+ return ctx.Err()
+ default:
+ }
+ // Determine wait limit
+ now := time.Now()
+ waitLimit := InfDuration
+ if deadline, ok := ctx.Deadline(); ok {
+ waitLimit = deadline.Sub(now)
+ }
+ // Reserve
+ r := lim.reserveN(now, n, waitLimit)
+ if !r.ok {
+ return fmt.Errorf("rate: Wait(n=%d) would exceed context deadline", n)
+ }
+ // Wait
+ t := time.NewTimer(r.DelayFrom(now))
+ defer t.Stop()
+ select {
+ case <-t.C:
+ // We can proceed.
+ return nil
+ case <-ctx.Done():
+ // Context was canceled before we could proceed. Cancel the
+ // reservation, which may permit other events to proceed sooner.
+ r.Cancel()
+ return ctx.Err()
+ }
+}
+
+// SetLimit is shorthand for SetLimitAt(time.Now(), newLimit).
+func (lim *Limiter) SetLimit(newLimit Limit) {
+ lim.SetLimitAt(time.Now(), newLimit)
+}
+
+// SetLimitAt sets a new Limit for the limiter. The new Limit, and Burst, may be violated
+// or underutilized by those which reserved (using Reserve or Wait) but did not yet act
+// before SetLimitAt was called.
+func (lim *Limiter) SetLimitAt(now time.Time, newLimit Limit) {
+ lim.mu.Lock()
+ defer lim.mu.Unlock()
+
+ now, _, tokens := lim.advance(now)
+
+ lim.last = now
+ lim.tokens = tokens
+ lim.limit = newLimit
+}
+
+// reserveN is a helper method for AllowN, ReserveN, and WaitN.
+// maxFutureReserve specifies the maximum reservation wait duration allowed.
+// reserveN returns Reservation, not *Reservation, to avoid allocation in AllowN and WaitN.
+func (lim *Limiter) reserveN(now time.Time, n int, maxFutureReserve time.Duration) Reservation {
+ lim.mu.Lock()
+ defer lim.mu.Unlock()
+
+ if lim.limit == Inf {
+ return Reservation{
+ ok: true,
+ lim: lim,
+ tokens: n,
+ timeToAct: now,
+ }
+ }
+
+ now, last, tokens := lim.advance(now)
+
+ // Calculate the remaining number of tokens resulting from the request.
+ tokens -= float64(n)
+
+ // Calculate the wait duration
+ var waitDuration time.Duration
+ if tokens < 0 {
+ waitDuration = lim.limit.durationFromTokens(-tokens)
+ }
+
+ // Decide result
+ ok := n <= lim.burst && waitDuration <= maxFutureReserve
+
+ // Prepare reservation
+ r := Reservation{
+ ok: ok,
+ lim: lim,
+ limit: lim.limit,
+ }
+ if ok {
+ r.tokens = n
+ r.timeToAct = now.Add(waitDuration)
+ }
+
+ // Update state
+ if ok {
+ lim.last = now
+ lim.tokens = tokens
+ lim.lastEvent = r.timeToAct
+ } else {
+ lim.last = last
+ }
+
+ return r
+}
+
+// advance calculates and returns an updated state for lim resulting from the passage of time.
+// lim is not changed.
+func (lim *Limiter) advance(now time.Time) (newNow time.Time, newLast time.Time, newTokens float64) {
+ last := lim.last
+ if now.Before(last) {
+ last = now
+ }
+
+ // Avoid making delta overflow below when last is very old.
+ maxElapsed := lim.limit.durationFromTokens(float64(lim.burst) - lim.tokens)
+ elapsed := now.Sub(last)
+ if elapsed > maxElapsed {
+ elapsed = maxElapsed
+ }
+
+ // Calculate the new number of tokens, due to time that passed.
+ delta := lim.limit.tokensFromDuration(elapsed)
+ tokens := lim.tokens + delta
+ if burst := float64(lim.burst); tokens > burst {
+ tokens = burst
+ }
+
+ return now, last, tokens
+}
+
+// durationFromTokens is a unit conversion function from the number of tokens to the duration
+// of time it takes to accumulate them at a rate of limit tokens per second.
+func (limit Limit) durationFromTokens(tokens float64) time.Duration {
+ seconds := tokens / float64(limit)
+ return time.Nanosecond * time.Duration(1e9*seconds)
+}
+
+// tokensFromDuration is a unit conversion function from a time duration to the number of tokens
+// which could be accumulated during that duration at a rate of limit tokens per second.
+func (limit Limit) tokensFromDuration(d time.Duration) float64 {
+ return d.Seconds() * float64(limit)
+}
diff --git a/rate/rate_test.go b/rate/rate_test.go
new file mode 100644
index 0000000..28338be
--- /dev/null
+++ b/rate/rate_test.go
@@ -0,0 +1,421 @@
+// Copyright 2015 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package rate
+
+import (
+ "math"
+ "sync"
+ "sync/atomic"
+ "testing"
+ "time"
+
+ "golang.org/x/net/context"
+)
+
+func TestLimit(t *testing.T) {
+ if Limit(10) == Inf {
+ t.Errorf("Limit(10) == Inf should be false")
+ }
+}
+
+func closeEnough(a, b Limit) bool {
+ return (math.Abs(float64(a)/float64(b)) - 1.0) < 1e-9
+}
+
+func TestEvery(t *testing.T) {
+ cases := []struct {
+ interval time.Duration
+ lim Limit
+ }{
+ {0, Inf},
+ {-1, Inf},
+ {1 * time.Nanosecond, Limit(1e9)},
+ {1 * time.Microsecond, Limit(1e6)},
+ {1 * time.Millisecond, Limit(1e3)},
+ {10 * time.Millisecond, Limit(100)},
+ {100 * time.Millisecond, Limit(10)},
+ {1 * time.Second, Limit(1)},
+ {2 * time.Second, Limit(0.5)},
+ {time.Duration(2.5 * float64(time.Second)), Limit(0.4)},
+ {4 * time.Second, Limit(0.25)},
+ {10 * time.Second, Limit(0.1)},
+ {time.Duration(math.MaxInt64), Limit(1e9 / float64(math.MaxInt64))},
+ }
+ for _, tc := range cases {
+ lim := Every(tc.interval)
+ if !closeEnough(lim, tc.lim) {
+ t.Errorf("Every(%v) = %v want %v", tc.interval, lim, tc.lim)
+ }
+ }
+}
+
+const (
+ d = 100 * time.Millisecond
+)
+
+var (
+ t0 = time.Now()
+ t1 = t0.Add(time.Duration(1) * d)
+ t2 = t0.Add(time.Duration(2) * d)
+ t3 = t0.Add(time.Duration(3) * d)
+ t4 = t0.Add(time.Duration(4) * d)
+ t5 = t0.Add(time.Duration(5) * d)
+ t9 = t0.Add(time.Duration(9) * d)
+)
+
+type allow struct {
+ t time.Time
+ n int
+ ok bool
+}
+
+func run(t *testing.T, lim *Limiter, allows []allow) {
+ for i, allow := range allows {
+ ok := lim.AllowN(allow.t, allow.n)
+ if ok != allow.ok {
+ t.Errorf("step %d: lim.AllowN(%v, %v) = %v want %v",
+ i, allow.t, allow.n, ok, allow.ok)
+ }
+ }
+}
+
+func TestLimiterBurst1(t *testing.T) {
+ run(t, NewLimiter(10, 1), []allow{
+ {t0, 1, true},
+ {t0, 1, false},
+ {t0, 1, false},
+ {t1, 1, true},
+ {t1, 1, false},
+ {t1, 1, false},
+ {t2, 2, false}, // burst size is 1, so n=2 always fails
+ {t2, 1, true},
+ {t2, 1, false},
+ })
+}
+
+func TestLimiterBurst3(t *testing.T) {
+ run(t, NewLimiter(10, 3), []allow{
+ {t0, 2, true},
+ {t0, 2, false},
+ {t0, 1, true},
+ {t0, 1, false},
+ {t1, 4, false},
+ {t2, 1, true},
+ {t3, 1, true},
+ {t4, 1, true},
+ {t4, 1, true},
+ {t4, 1, false},
+ {t4, 1, false},
+ {t9, 3, true},
+ {t9, 0, true},
+ })
+}
+
+func TestLimiterJumpBackwards(t *testing.T) {
+ run(t, NewLimiter(10, 3), []allow{
+ {t1, 1, true}, // start at t1
+ {t0, 1, true}, // jump back to t0, two tokens remain
+ {t0, 1, true},
+ {t0, 1, false},
+ {t0, 1, false},
+ {t1, 1, true}, // got a token
+ {t1, 1, false},
+ {t1, 1, false},
+ {t2, 1, true}, // got another token
+ {t2, 1, false},
+ {t2, 1, false},
+ })
+}
+
+func TestSimultaneousRequests(t *testing.T) {
+ const (
+ limit = 1
+ burst = 5
+ numRequests = 15
+ )
+ var (
+ wg sync.WaitGroup
+ numOK = uint32(0)
+ )
+
+ // Very slow replenishing bucket.
+ lim := NewLimiter(limit, burst)
+
+ // Tries to take a token, atomically updates the counter and decreases the wait
+ // group counter.
+ f := func() {
+ defer wg.Done()
+ if ok := lim.Allow(); ok {
+ atomic.AddUint32(&numOK, 1)
+ }
+ }
+
+ wg.Add(numRequests)
+ for i := 0; i < numRequests; i++ {
+ go f()
+ }
+ wg.Wait()
+ if numOK != burst {
+ t.Errorf("numOK = %d, want %d", numOK, burst)
+ }
+}
+
+func TestLongRunningQPS(t *testing.T) {
+ // The test runs for a few seconds executing many requests and then checks
+ // that overall number of requests is reasonable.
+ const (
+ limit = 100
+ burst = 100
+ )
+ var numOK = int32(0)
+
+ lim := NewLimiter(limit, burst)
+
+ var wg sync.WaitGroup
+ f := func() {
+ if ok := lim.Allow(); ok {
+ atomic.AddInt32(&numOK, 1)
+ }
+ wg.Done()
+ }
+
+ start := time.Now()
+ end := start.Add(5 * time.Second)
+ for time.Now().Before(end) {
+ wg.Add(1)
+ go f()
+
+ // This will still offer ~500 requests per second, but won't consume
+ // outrageous amount of CPU.
+ time.Sleep(2 * time.Millisecond)
+ }
+ wg.Wait()
+ elapsed := time.Since(start)
+ ideal := burst + (limit * float64(elapsed) / float64(time.Second))
+
+ // We should never get more requests than allowed.
+ if want := int32(ideal + 1); numOK > want {
+ t.Errorf("numOK = %d, want %d (ideal %f)", numOK, want, ideal)
+ }
+ // We should get very close to the number of requests allowed.
+ if want := int32(0.999 * ideal); numOK < want {
+ t.Errorf("numOK = %d, want %d (ideal %f)", numOK, want, ideal)
+ }
+}
+
+type request struct {
+ t time.Time
+ n int
+ act time.Time
+ ok bool
+}
+
+// dFromDuration converts a duration to a multiple of the global constant d
+func dFromDuration(dur time.Duration) int {
+ // Adding a millisecond to be swallowed by the integer division
+ // because we don't care about small inaccuracies
+ return int((dur + time.Millisecond) / d)
+}
+
+// dSince returns multiples of d since t0
+func dSince(t time.Time) int {
+ return dFromDuration(t.Sub(t0))
+}
+
+func runReserve(t *testing.T, lim *Limiter, req request) *Reservation {
+ return runReserveMax(t, lim, req, InfDuration)
+}
+
+func runReserveMax(t *testing.T, lim *Limiter, req request, maxReserve time.Duration) *Reservation {
+ r := lim.reserveN(req.t, req.n, maxReserve)
+ if r.ok && (dSince(r.timeToAct) != dSince(req.act)) || r.ok != req.ok {
+ t.Errorf("lim.reserveN(t%d, %v, %v) = (t%d, %v) want (t%d, %v)",
+ dSince(req.t), req.n, maxReserve, dSince(r.timeToAct), r.ok, dSince(req.act), req.ok)
+ }
+ return &r
+}
+
+func TestSimpleReserve(t *testing.T) {
+ lim := NewLimiter(10, 2)
+
+ runReserve(t, lim, request{t0, 2, t0, true})
+ runReserve(t, lim, request{t0, 2, t2, true})
+ runReserve(t, lim, request{t3, 2, t4, true})
+}
+
+func TestMix(t *testing.T) {
+ lim := NewLimiter(10, 2)
+
+ runReserve(t, lim, request{t0, 3, t1, false}) // should return false because n > Burst
+ runReserve(t, lim, request{t0, 2, t0, true})
+ run(t, lim, []allow{{t1, 2, false}}) // not enought tokens - don't allow
+ runReserve(t, lim, request{t1, 2, t2, true})
+ run(t, lim, []allow{{t1, 1, false}}) // negative tokens - don't allow
+ run(t, lim, []allow{{t3, 1, true}})
+}
+
+func TestCancelInvalid(t *testing.T) {
+ lim := NewLimiter(10, 2)
+
+ runReserve(t, lim, request{t0, 2, t0, true})
+ r := runReserve(t, lim, request{t0, 3, t3, false})
+ r.CancelAt(t0) // should have no effect
+ runReserve(t, lim, request{t0, 2, t2, true}) // did not get extra tokens
+}
+
+func TestCancelLast(t *testing.T) {
+ lim := NewLimiter(10, 2)
+
+ runReserve(t, lim, request{t0, 2, t0, true})
+ r := runReserve(t, lim, request{t0, 2, t2, true})
+ r.CancelAt(t1) // got 2 tokens back
+ runReserve(t, lim, request{t1, 2, t2, true})
+}
+
+func TestCancelTooLate(t *testing.T) {
+ lim := NewLimiter(10, 2)
+
+ runReserve(t, lim, request{t0, 2, t0, true})
+ r := runReserve(t, lim, request{t0, 2, t2, true})
+ r.CancelAt(t3) // too late to cancel - should have no effect
+ runReserve(t, lim, request{t3, 2, t4, true})
+}
+
+func TestCancel0Tokens(t *testing.T) {
+ lim := NewLimiter(10, 2)
+
+ runReserve(t, lim, request{t0, 2, t0, true})
+ r := runReserve(t, lim, request{t0, 1, t1, true})
+ runReserve(t, lim, request{t0, 1, t2, true})
+ r.CancelAt(t0) // got 0 tokens back
+ runReserve(t, lim, request{t0, 1, t3, true})
+}
+
+func TestCancel1Token(t *testing.T) {
+ lim := NewLimiter(10, 2)
+
+ runReserve(t, lim, request{t0, 2, t0, true})
+ r := runReserve(t, lim, request{t0, 2, t2, true})
+ runReserve(t, lim, request{t0, 1, t3, true})
+ r.CancelAt(t2) // got 1 token back
+ runReserve(t, lim, request{t2, 2, t4, true})
+}
+
+func TestCancelMulti(t *testing.T) {
+ lim := NewLimiter(10, 4)
+
+ runReserve(t, lim, request{t0, 4, t0, true})
+ rA := runReserve(t, lim, request{t0, 3, t3, true})
+ runReserve(t, lim, request{t0, 1, t4, true})
+ rC := runReserve(t, lim, request{t0, 1, t5, true})
+ rC.CancelAt(t1) // get 1 token back
+ rA.CancelAt(t1) // get 2 tokens back, as if C was never reserved
+ runReserve(t, lim, request{t1, 3, t5, true})
+}
+
+func TestReserveJumpBack(t *testing.T) {
+ lim := NewLimiter(10, 2)
+
+ runReserve(t, lim, request{t1, 2, t1, true}) // start at t1
+ runReserve(t, lim, request{t0, 1, t1, true}) // should violate Limit,Burst
+ runReserve(t, lim, request{t2, 2, t3, true})
+}
+
+func TestReserveJumpBackCancel(t *testing.T) {
+ lim := NewLimiter(10, 2)
+
+ runReserve(t, lim, request{t1, 2, t1, true}) // start at t1
+ r := runReserve(t, lim, request{t1, 2, t3, true})
+ runReserve(t, lim, request{t1, 1, t4, true})
+ r.CancelAt(t0) // cancel at t0, get 1 token back
+ runReserve(t, lim, request{t1, 2, t4, true}) // should violate Limit,Burst
+}
+
+func TestReserveSetLimit(t *testing.T) {
+ lim := NewLimiter(5, 2)
+
+ runReserve(t, lim, request{t0, 2, t0, true})
+ runReserve(t, lim, request{t0, 2, t4, true})
+ lim.SetLimitAt(t2, 10)
+ runReserve(t, lim, request{t2, 1, t4, true}) // violates Limit and Burst
+}
+
+func TestReserveSetLimitCancel(t *testing.T) {
+ lim := NewLimiter(5, 2)
+
+ runReserve(t, lim, request{t0, 2, t0, true})
+ r := runReserve(t, lim, request{t0, 2, t4, true})
+ lim.SetLimitAt(t2, 10)
+ r.CancelAt(t2) // 2 tokens back
+ runReserve(t, lim, request{t2, 2, t3, true})
+}
+
+func TestReserveMax(t *testing.T) {
+ lim := NewLimiter(10, 2)
+ maxT := d
+
+ runReserveMax(t, lim, request{t0, 2, t0, true}, maxT)
+ runReserveMax(t, lim, request{t0, 1, t1, true}, maxT) // reserve for close future
+ runReserveMax(t, lim, request{t0, 1, t2, false}, maxT) // time to act too far in the future
+}
+
+type wait struct {
+ name string
+ ctx context.Context
+ n int
+ delay int // in multiples of d
+ nilErr bool
+}
+
+func runWait(t *testing.T, lim *Limiter, w wait) {
+ start := time.Now()
+ err := lim.WaitN(w.ctx, w.n)
+ delay := time.Now().Sub(start)
+ if (w.nilErr && err != nil) || (!w.nilErr && err == nil) || w.delay != dFromDuration(delay) {
+ errString := "<nil>"
+ if !w.nilErr {
+ errString = "<non-nil error>"
+ }
+ t.Errorf("lim.WaitN(%v, lim, %v) = %v with delay %v ; want %v with delay %v",
+ w.name, w.n, err, delay, errString, d*time.Duration(w.delay))
+ }
+}
+
+func TestWaitSimple(t *testing.T) {
+ lim := NewLimiter(10, 3)
+
+ ctx, cancel := context.WithCancel(context.Background())
+ 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{"act-now", context.Background(), 2, 0, true})
+ runWait(t, lim, wait{"act-later", context.Background(), 3, 2, true})
+}
+
+func TestWaitCancel(t *testing.T) {
+ lim := NewLimiter(10, 3)
+
+ ctx, cancel := context.WithCancel(context.Background())
+ runWait(t, lim, wait{"act-now", ctx, 2, 0, true}) // after this lim.tokens = 1
+ go func() {
+ time.Sleep(d)
+ cancel()
+ }()
+ runWait(t, lim, wait{"will-cancel", ctx, 3, 1, false})
+ // should get 3 tokens back, and have lim.tokens = 2
+ t.Logf("tokens:%v last:%v lastEvent:%v", lim.tokens, lim.last, lim.lastEvent)
+ runWait(t, lim, wait{"act-now-after-cancel", context.Background(), 2, 0, true})
+}
+
+func TestWaitTimeout(t *testing.T) {
+ lim := NewLimiter(10, 3)
+
+ ctx, cancel := context.WithTimeout(context.Background(), d)
+ defer cancel()
+ runWait(t, lim, wait{"act-now", ctx, 2, 0, true})
+ runWait(t, lim, wait{"w-timeout-err", ctx, 3, 0, false})
+}
