mirror of
				https://source.quilibrium.com/quilibrium/ceremonyclient.git
				synced 2025-11-04 03:07:28 +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
 | 
						|
}
 |