224 lines
6.7 KiB
Go
224 lines
6.7 KiB
Go
// Copyright 2023 The Cockroach Authors.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
|
// implied. See the License for the specific language governing
|
|
// permissions and limitations under the License.
|
|
|
|
package tokenbucket
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
)
|
|
|
|
// Tokens are abstract units (usually units of work).
|
|
type Tokens float64
|
|
|
|
// TokensPerSecond is the rate of token replenishment.
|
|
type TokensPerSecond float64
|
|
|
|
// TokenBucket implements the basic accounting for a token bucket.
|
|
//
|
|
// A token bucket has a rate of replenishment and a burst limit. Tokens are
|
|
// replenished over time, up to the burst limit.
|
|
//
|
|
// The token bucket keeps track of the current amount and updates it as time
|
|
// passes. The bucket can go into debt (i.e. negative current amount).
|
|
type TokenBucket struct {
|
|
rate TokensPerSecond
|
|
burst Tokens
|
|
nowFn func() time.Time
|
|
|
|
current Tokens
|
|
lastUpdated time.Time
|
|
|
|
exhaustedStart time.Time
|
|
exhaustedMicros int64
|
|
}
|
|
|
|
// Init the token bucket.
|
|
func (tb *TokenBucket) Init(rate TokensPerSecond, burst Tokens) {
|
|
tb.InitWithNowFn(rate, burst, func() time.Time {
|
|
return time.Now()
|
|
})
|
|
}
|
|
|
|
// Init the token bucket with a custom "Now" function.
|
|
// Note that Wait/WaitCtx cannot be used with a custom time function.
|
|
func (tb *TokenBucket) InitWithNowFn(rate TokensPerSecond, burst Tokens, nowFn func() time.Time) {
|
|
*tb = TokenBucket{
|
|
rate: rate,
|
|
burst: burst,
|
|
nowFn: nowFn,
|
|
current: burst,
|
|
lastUpdated: nowFn(),
|
|
}
|
|
}
|
|
|
|
// Update moves the time forward, accounting for the replenishment since the
|
|
// last update.
|
|
func (tb *TokenBucket) Update() {
|
|
now := tb.nowFn()
|
|
if since := now.Sub(tb.lastUpdated); since > 0 {
|
|
tb.current += Tokens(float64(tb.rate) * since.Seconds())
|
|
|
|
if tb.current > tb.burst {
|
|
tb.current = tb.burst
|
|
}
|
|
tb.lastUpdated = now
|
|
tb.updateExhaustedMicros()
|
|
}
|
|
}
|
|
|
|
// UpdateConfig updates the rate and burst limits. The change in burst will be
|
|
// applied to the current token quantity. For example, if the RateLimiter
|
|
// currently had 5 available tokens and the burst is updated from 10 to 20, the
|
|
// amount will increase to 15. Similarly, if the burst is decreased by 10, the
|
|
// current quota will decrease accordingly, potentially putting the limiter into
|
|
// debt.
|
|
func (tb *TokenBucket) UpdateConfig(rate TokensPerSecond, burst Tokens) {
|
|
tb.Update()
|
|
|
|
burstDelta := burst - tb.burst
|
|
tb.rate = rate
|
|
tb.burst = burst
|
|
|
|
tb.current += burstDelta
|
|
tb.updateExhaustedMicros()
|
|
}
|
|
|
|
// Reset resets the current tokens to whatever the burst is.
|
|
func (tb *TokenBucket) Reset() {
|
|
tb.current = tb.burst
|
|
tb.updateExhaustedMicros()
|
|
}
|
|
|
|
// Adjust returns tokens to the bucket (positive delta) or accounts for a debt
|
|
// of tokens (negative delta).
|
|
func (tb *TokenBucket) Adjust(delta Tokens) {
|
|
tb.Update()
|
|
tb.current += delta
|
|
if tb.current > tb.burst {
|
|
tb.current = tb.burst
|
|
}
|
|
tb.updateExhaustedMicros()
|
|
}
|
|
|
|
// TryToFulfill either removes the given amount if is available, or returns a
|
|
// time after which the request should be retried.
|
|
func (tb *TokenBucket) TryToFulfill(amount Tokens) (fulfilled bool, tryAgainAfter time.Duration) {
|
|
tb.Update()
|
|
|
|
// Deal with the case where the request is larger than the burst size. In
|
|
// this case we'll allow the acquisition to complete if and when the current
|
|
// value is equal to the burst. If the acquisition succeeds, it will put the
|
|
// limiter into debt.
|
|
want := amount
|
|
if want > tb.burst {
|
|
want = tb.burst
|
|
}
|
|
if delta := want - tb.current; delta > 0 {
|
|
// Compute the time it will take to get to the needed capacity.
|
|
timeDelta := time.Duration((float64(delta) * float64(time.Second)) / float64(tb.rate))
|
|
if timeDelta < time.Nanosecond {
|
|
timeDelta = time.Nanosecond
|
|
}
|
|
return false, timeDelta
|
|
}
|
|
|
|
tb.current -= amount
|
|
tb.updateExhaustedMicros()
|
|
return true, 0
|
|
}
|
|
|
|
// Wait removes the given amount, waiting as long as necessary.
|
|
func (tb *TokenBucket) Wait(amount Tokens) {
|
|
_ = tb.WaitCtx(context.Background(), amount)
|
|
}
|
|
|
|
// WaitCtx removes the given amount, waiting as long as necessary or until the
|
|
// context is canceled.
|
|
func (tb *TokenBucket) WaitCtx(ctx context.Context, amount Tokens) error {
|
|
// We want to check for context cancelation even if we don't need to wait.
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
default:
|
|
}
|
|
for {
|
|
fulfilled, tryAgainAfter := tb.TryToFulfill(amount)
|
|
if fulfilled {
|
|
return nil
|
|
}
|
|
select {
|
|
case <-time.After(tryAgainAfter):
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
}
|
|
}
|
|
}
|
|
|
|
// Exhausted returns the cumulative duration over which this token bucket was
|
|
// exhausted. Exported only for metrics.
|
|
func (tb *TokenBucket) Exhausted() time.Duration {
|
|
var ongoingExhaustionMicros int64
|
|
if !tb.exhaustedStart.IsZero() {
|
|
ongoingExhaustionMicros = tb.nowFn().Sub(tb.exhaustedStart).Microseconds()
|
|
}
|
|
return time.Duration(tb.exhaustedMicros+ongoingExhaustionMicros) * time.Microsecond
|
|
}
|
|
|
|
// EnsureLowerBound ensures that the current number of tokens is at least
|
|
// minTokens. The value of minTokens can be negative. This is useful in
|
|
// scenarios where the use of Adjust can be used to subtract a large amount of
|
|
// tokens, resulting in a token bucket with a very low negative number of
|
|
// tokens which would take a long time to recover to zero. In such scenarios,
|
|
// a periodic call to EnsureLowerBound is useful to reset to a better state.
|
|
func (tb *TokenBucket) EnsureLowerBound(minTokens Tokens) {
|
|
tb.Update()
|
|
if tb.current >= minTokens {
|
|
return
|
|
}
|
|
// tb.current < minTokens.
|
|
tb.current = minTokens
|
|
if tb.current > tb.burst {
|
|
tb.current = tb.burst
|
|
}
|
|
if tb.current > 0 {
|
|
tb.updateExhaustedMicros()
|
|
}
|
|
}
|
|
|
|
// Available returns the currently available tokens (can be -ve). Exported only
|
|
// for metrics.
|
|
func (tb *TokenBucket) Available() Tokens {
|
|
return tb.current
|
|
}
|
|
|
|
func (tb *TokenBucket) updateExhaustedMicros() {
|
|
now := tb.nowFn()
|
|
if tb.current <= 0 && tb.exhaustedStart.IsZero() {
|
|
tb.exhaustedStart = now
|
|
}
|
|
if tb.current > 0 && !tb.exhaustedStart.IsZero() {
|
|
tb.exhaustedMicros += now.Sub(tb.exhaustedStart).Microseconds()
|
|
tb.exhaustedStart = time.Time{}
|
|
}
|
|
}
|
|
|
|
// TestingInternalParameters returns the refill rate (configured), burst tokens
|
|
// (configured), and number of available tokens where available <= burst. It's
|
|
// used in tests.
|
|
func (tb *TokenBucket) TestingInternalParameters() (rate TokensPerSecond, burst, available Tokens) {
|
|
tb.Update()
|
|
return tb.rate, tb.burst, tb.current
|
|
}
|