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>
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 @@
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 @@
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)
