Upgrade versions of vendored libs

vendor/github.com/cockroachdb/tokenbucket/token_bucket.go

@@ -0,0 +1,202 @@
+// 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
+}
+
+// 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
+}