rate: Add SetBurst() to dynamically update burst size

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>

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.

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)