ceremonyclient/pebble/internal/rate/rate.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
}