Skip to content

Commit

Permalink
rate: Add SetBurst() to dynamically update burst size
Browse files Browse the repository at this point in the history
This commit adds the ability to dynamically update
burst size.

Fixes golang/go#23575

Signed-off-by: Simarpreet Singh <simar@linux.com>
Change-Id: I40da9ffdb108dee6ad15efb8700e3ae60d1169d9
Reviewed-on: https://go-review.googlesource.com/c/time/+/184082
Reviewed-by: Sameer Ajmani <sameer@golang.org>
Run-TryBot: Sameer Ajmani <sameer@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
  • Loading branch information
simar7 authored and Sajmani committed Sep 21, 2019
1 parent 9d24e82 commit c4c64ca
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 0 deletions.
17 changes: 17 additions & 0 deletions rate/rate.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
9 changes: 9 additions & 0 deletions rate/rate_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down

0 comments on commit c4c64ca

Please sign in to comment.