diff --git a/rate/rate.go b/rate/rate.go index ae93e24..85c18b5 100644 --- a/rate/rate.go +++ b/rate/rate.go @@ -281,6 +281,23 @@ func (lim *Limiter) SetLimitAt(now time.Time, newLimit Limit) { lim.limit = newLimit } +// SetBurst is shorthand for SetBurstAt(time.Now(), newBurst). +func (lim *Limiter) SetBurst(newBurst int) { + lim.SetBurstAt(time.Now(), newBurst) +} + +// SetBurstAt sets a new burst size for the limiter. +func (lim *Limiter) SetBurstAt(now time.Time, newBurst int) { + lim.mu.Lock() + defer lim.mu.Unlock() + + now, _, tokens := lim.advance(now) + + lim.last = now + lim.tokens = tokens + lim.burst = newBurst +} + // 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. diff --git a/rate/rate_test.go b/rate/rate_test.go index ec8c66d..448300d 100644 --- a/rate/rate_test.go +++ b/rate/rate_test.go @@ -352,6 +352,15 @@ func TestReserveSetLimit(t *testing.T) { runReserve(t, lim, request{t2, 1, t4, true}) // violates Limit and Burst } +func TestReserveSetBurst(t *testing.T) { + lim := NewLimiter(5, 2) + + runReserve(t, lim, request{t0, 2, t0, true}) + runReserve(t, lim, request{t0, 2, t4, true}) + lim.SetBurstAt(t3, 4) + runReserve(t, lim, request{t0, 4, t9, true}) // violates Limit and Burst +} + func TestReserveSetLimitCancel(t *testing.T) { lim := NewLimiter(5, 2)