mirror of
https://source.quilibrium.com/quilibrium/ceremonyclient.git
synced 2024-11-11 02:35:18 +00:00
97 lines
2.6 KiB
Go
97 lines
2.6 KiB
Go
// Copyright 2023 The LevelDB-Go and Pebble Authors. All rights reserved. Use
|
|
// of this source code is governed by a BSD-style license that can be found in
|
|
// the LICENSE file.
|
|
|
|
// Package rate provides a rate limiter.
|
|
package rate // import "github.com/cockroachdb/pebble/internal/rate"
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/cockroachdb/tokenbucket"
|
|
)
|
|
|
|
// A Limiter controls how frequently events are allowed to happen.
|
|
// It implements a "token bucket" of size b, initially full and refilled
|
|
// at rate r tokens per second.
|
|
//
|
|
// Informally, in any large enough time interval, the Limiter limits the
|
|
// rate to r tokens per second, with a maximum burst size of b events.
|
|
//
|
|
// Limiter is thread-safe.
|
|
type Limiter struct {
|
|
mu struct {
|
|
sync.Mutex
|
|
tb tokenbucket.TokenBucket
|
|
rate float64
|
|
burst float64
|
|
}
|
|
sleepFn func(d time.Duration)
|
|
}
|
|
|
|
// NewLimiter returns a new Limiter that allows events up to rate r and permits
|
|
// bursts of at most b tokens.
|
|
func NewLimiter(r float64, b float64) *Limiter {
|
|
l := &Limiter{}
|
|
l.mu.tb.Init(tokenbucket.TokensPerSecond(r), tokenbucket.Tokens(b))
|
|
l.mu.rate = r
|
|
l.mu.burst = b
|
|
return l
|
|
}
|
|
|
|
// NewLimiterWithCustomTime returns a new Limiter that allows events up to rate
|
|
// r and permits bursts of at most b tokens. The limiter uses the given
|
|
// functions to retrieve the current time and to sleep (useful for testing).
|
|
func NewLimiterWithCustomTime(
|
|
r float64, b float64, nowFn func() time.Time, sleepFn func(d time.Duration),
|
|
) *Limiter {
|
|
l := &Limiter{}
|
|
l.mu.tb.InitWithNowFn(tokenbucket.TokensPerSecond(r), tokenbucket.Tokens(b), nowFn)
|
|
l.mu.rate = r
|
|
l.mu.burst = b
|
|
l.sleepFn = sleepFn
|
|
return l
|
|
}
|
|
|
|
// Wait sleeps until enough tokens are available. If n is more than the burst,
|
|
// the token bucket will go into debt, delaying future operations.
|
|
func (l *Limiter) Wait(n float64) {
|
|
for {
|
|
l.mu.Lock()
|
|
ok, d := l.mu.tb.TryToFulfill(tokenbucket.Tokens(n))
|
|
l.mu.Unlock()
|
|
if ok {
|
|
return
|
|
}
|
|
if l.sleepFn != nil {
|
|
l.sleepFn(d)
|
|
} else {
|
|
time.Sleep(d)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Remove removes tokens for an operation that bypassed any waiting; it can put
|
|
// the token bucket into debt, delaying future operations.
|
|
func (l *Limiter) Remove(n float64) {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
l.mu.tb.Adjust(-tokenbucket.Tokens(n))
|
|
}
|
|
|
|
// Rate returns the current rate limit.
|
|
func (l *Limiter) Rate() float64 {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
return l.mu.rate
|
|
}
|
|
|
|
// SetRate updates the rate limit.
|
|
func (l *Limiter) SetRate(r float64) {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
l.mu.tb.UpdateConfig(tokenbucket.TokensPerSecond(r), tokenbucket.Tokens(l.mu.burst))
|
|
l.mu.rate = r
|
|
}
|