use alloc::vec::Vec;
use core::sync::atomic::{AtomicBool, Ordering};
use core::time::Duration;
use portable_atomic::AtomicU64;
use spin::Mutex;
#[cfg_attr(
not(feature = "std"),
allow(
dead_code,
reason = "no_std-ready token-bucket core; awaits a clock-injecting no_std caller"
)
)]
const NANOS_PER_SEC: u128 = 1_000_000_000;
#[derive(Debug)]
#[cfg_attr(
not(feature = "std"),
allow(
dead_code,
reason = "no_std-ready token-bucket state; awaits a clock-injecting no_std caller"
)
)]
struct Bucket {
debited: u128,
credited: u128,
withdrawn: Vec<(u128, u128)>,
withdrawn_ahead_of_all: u128,
last_refill_nanos: u128,
changes: u64,
last_off: u64,
backoff: Option<Backoff>,
}
#[cfg_attr(
not(feature = "std"),
allow(
dead_code,
reason = "no_std-ready token-bucket state; awaits a clock-injecting no_std caller"
)
)]
impl Bucket {
fn full(rate: u64, now_nanos: u128) -> Self {
let mut bucket = Self {
debited: 0,
credited: 0,
withdrawn: Vec::new(),
withdrawn_ahead_of_all: 0,
last_refill_nanos: 0,
changes: 0,
last_off: 0,
backoff: None,
};
bucket.fill(rate, now_nanos);
bucket
}
fn outstanding(&self) -> u128 {
let withdrawn: u128 = self.withdrawn.iter().map(|&(_, bytes)| bytes).sum();
self.debited - self.withdrawn_ahead_of_all - withdrawn
}
fn standing(&self, position: u128) -> u128 {
let ahead: u128 = self
.withdrawn
.iter()
.filter(|&&(at, _)| at <= position)
.map(|&(_, bytes)| bytes)
.sum();
(position - ahead).saturating_sub(self.withdrawn_ahead_of_all)
}
fn fill(&mut self, rate: u64, now_nanos: u128) {
self.credited = self.outstanding() + u128::from(rate);
self.advance_clock(now_nanos);
self.settle_withdrawn();
}
fn advance_clock(&mut self, now_nanos: u128) {
self.last_refill_nanos = self.last_refill_nanos.max(now_nanos);
}
fn refill(&mut self, rate: u64, now_nanos: u128) {
let elapsed = now_nanos.saturating_sub(self.last_refill_nanos);
if elapsed == 0 || rate == 0 {
return;
}
let refilled = elapsed.saturating_mul(u128::from(rate)) / NANOS_PER_SEC;
if refilled == 0 {
return;
}
let ceiling = self.outstanding() + u128::from(rate);
if self.credited + refilled >= ceiling {
self.credited = ceiling;
self.last_refill_nanos = now_nanos;
} else {
self.credited += refilled;
let spent = (refilled * NANOS_PER_SEC).div_ceil(u128::from(rate));
self.last_refill_nanos += spent;
}
self.settle_withdrawn();
}
fn cap(&mut self, rate: u64) {
let ceiling = self.outstanding() + u128::from(rate);
if self.credited > ceiling {
self.credited = ceiling;
}
}
fn settle_withdrawn(&mut self) {
let mut i = 0;
while let Some(&(at, bytes)) = self.withdrawn.get(i) {
if self.credited >= self.standing(at) {
self.withdrawn.swap_remove(i);
self.withdrawn_ahead_of_all += bytes;
} else {
i += 1;
}
}
}
fn debit(&mut self, bytes: u64) -> u128 {
self.debited += u128::from(bytes);
self.debited
}
fn wait_for(&self, position: u128, rate: u64) -> Duration {
let owed = self.standing(position).saturating_sub(self.credited);
let rate = u128::from(rate);
#[expect(clippy::cast_possible_truncation, reason = "at most NANOS_PER_SEC")]
let nanos = ((owed % rate) * NANOS_PER_SEC).div_ceil(rate) as u64;
let secs = u64::try_from(owed / rate).unwrap_or(u64::MAX);
Duration::from_secs(secs).saturating_add(Duration::from_nanos(nanos))
}
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct LatencyBackoff {
ceiling: Duration,
hysteresis: f64,
floor: f64,
period: Duration,
}
impl LatencyBackoff {
pub const DEFAULT_HYSTERESIS: f64 = 0.2;
pub const DEFAULT_FLOOR: f64 = 0.05;
pub const DEFAULT_PERIOD: Duration = Duration::from_millis(100);
#[must_use]
pub const fn new(ceiling: Duration) -> Self {
Self {
ceiling,
hysteresis: Self::DEFAULT_HYSTERESIS,
floor: Self::DEFAULT_FLOOR,
period: Self::DEFAULT_PERIOD,
}
}
#[must_use]
pub fn with_hysteresis(mut self, fraction: f64) -> Self {
self.hysteresis = if fraction.is_nan() {
Self::DEFAULT_HYSTERESIS
} else {
fraction.clamp(0.0, 0.95)
};
self
}
#[must_use]
pub fn with_floor(mut self, fraction: f64) -> Self {
self.floor = if fraction.is_nan() {
Self::DEFAULT_FLOOR
} else {
fraction.clamp(0.001, 1.0)
};
self
}
#[must_use]
pub const fn with_period(mut self, period: Duration) -> Self {
self.period = period;
self
}
#[must_use]
pub const fn ceiling(&self) -> Duration {
self.ceiling
}
#[must_use]
pub const fn hysteresis(&self) -> f64 {
self.hysteresis
}
#[must_use]
pub const fn floor(&self) -> f64 {
self.floor
}
#[must_use]
pub const fn period(&self) -> Duration {
self.period
}
}
const LATENCY_SMOOTHING: f64 = 0.875;
const STEP_GAIN: f64 = 0.5;
const MAX_STEP_UP: f64 = 2.0;
#[derive(Debug)]
struct Backoff {
config: LatencyBackoff,
smoothed_nanos: Option<f64>,
fraction: f64,
last_step_nanos: u128,
}
impl Backoff {
fn new(config: LatencyBackoff, now_nanos: u128) -> Self {
Self {
config,
smoothed_nanos: None,
fraction: 1.0,
last_step_nanos: now_nanos,
}
}
#[expect(
clippy::cast_precision_loss,
reason = "latencies in nanoseconds as f64 lose precision only past 2^53 ns, about 104 days"
)]
#[expect(
clippy::suboptimal_flops,
reason = "f64::mul_add is not in core, and this module builds without std"
)]
fn observe(&mut self, latency: Duration, now_nanos: u128) {
let sample = latency.as_nanos() as f64;
let smoothed = match self.smoothed_nanos {
None => sample,
Some(previous) => (1.0 - LATENCY_SMOOTHING) * sample + LATENCY_SMOOTHING * previous,
};
self.smoothed_nanos = Some(smoothed);
if now_nanos.saturating_sub(self.last_step_nanos) < self.config.period.as_nanos() {
return;
}
self.last_step_nanos = now_nanos;
let ceiling = self.config.ceiling.as_nanos() as f64;
let climbs_below = ceiling * (1.0 - self.config.hysteresis);
if smoothed > ceiling {
let step = (1.0 - STEP_GAIN) + STEP_GAIN * (ceiling / smoothed);
self.fraction = (self.fraction * step).max(self.config.floor);
} else if smoothed < climbs_below {
let step = ((1.0 - STEP_GAIN) + STEP_GAIN * (ceiling / smoothed)).min(MAX_STEP_UP);
self.fraction = (self.fraction * step).min(1.0);
}
}
fn reconfigure(&mut self, config: LatencyBackoff) {
self.config = config;
self.fraction = self.fraction.max(config.floor);
}
}
fn effective_rate(configured: u64, backoff: Option<&Backoff>) -> u64 {
match backoff {
Some(backoff) if configured != 0 && backoff.fraction < 1.0 => {
#[expect(
clippy::cast_precision_loss,
clippy::cast_possible_truncation,
clippy::cast_sign_loss,
reason = "the product is in [0, configured], a fraction of a u64"
)]
let scaled = (configured as f64 * backoff.fraction) as u64;
scaled.clamp(1, configured)
}
_ => configured,
}
}
#[cfg(feature = "std")]
struct Ticket {
position: u128,
taken_at: u64,
}
#[cfg(feature = "std")]
#[derive(Debug, Default)]
pub(crate) struct Wakeup {
generation: parking_lot::Mutex<u64>,
rung: parking_lot::Condvar,
}
#[cfg(feature = "std")]
impl Wakeup {
fn generation(&self) -> u64 {
*self.generation.lock()
}
pub(crate) fn ring(&self) {
*self.generation.lock() += 1;
self.rung.notify_all();
}
fn sleep(&self, seen: u64, timeout: Duration) {
let mut generation = self.generation.lock();
if *generation == seen {
let _ = self.rung.wait_for(&mut generation, timeout);
}
}
}
#[derive(Debug)]
pub struct RateLimiter {
configured_bytes_per_sec: AtomicU64,
rate_bytes_per_sec: AtomicU64,
steers_by_latency: AtomicBool,
bucket: Mutex<Bucket>,
#[cfg(feature = "std")]
wakeup: alloc::sync::Arc<Wakeup>,
}
impl RateLimiter {
#[must_use]
pub fn new(rate_bytes_per_sec: u64) -> Self {
Self {
configured_bytes_per_sec: AtomicU64::new(rate_bytes_per_sec),
rate_bytes_per_sec: AtomicU64::new(rate_bytes_per_sec),
steers_by_latency: AtomicBool::new(false),
bucket: Mutex::new(Bucket::full(rate_bytes_per_sec, 0)),
#[cfg(feature = "std")]
wakeup: alloc::sync::Arc::default(),
}
}
#[cfg(feature = "std")]
pub(crate) fn wakeup(&self) -> alloc::sync::Arc<Wakeup> {
alloc::sync::Arc::clone(&self.wakeup)
}
#[cfg(feature = "std")]
pub fn wake_waiters(&self) {
self.wakeup.ring();
}
#[must_use]
pub fn rate(&self) -> u64 {
self.configured_bytes_per_sec.load(Ordering::Relaxed)
}
#[must_use]
pub fn effective_rate(&self) -> u64 {
self.rate_bytes_per_sec.load(Ordering::Relaxed)
}
fn switch_rate(&self, bucket: &mut Bucket, new: u64, now_nanos: u128) -> bool {
let old = self.rate_bytes_per_sec.load(Ordering::Relaxed);
if old == new {
return false;
}
if old == 0 {
bucket.fill(new, now_nanos);
} else {
bucket.refill(old, now_nanos);
bucket.advance_clock(now_nanos);
if new != 0 {
bucket.cap(new);
}
}
bucket.changes += 1;
if new == 0 {
bucket.last_off = bucket.changes;
}
self.rate_bytes_per_sec.store(new, Ordering::Relaxed);
true
}
fn rate_moved(&self) {
#[cfg(feature = "std")]
self.wakeup.ring();
}
pub fn set_latency_backoff_at(&self, backoff: Option<LatencyBackoff>, now: Duration) {
let now_nanos = now.as_nanos();
let mut bucket = self.bucket.lock();
match (backoff, bucket.backoff.as_mut()) {
(None, _) => bucket.backoff = None,
(Some(config), Some(current)) => current.reconfigure(config),
(Some(config), None) => bucket.backoff = Some(Backoff::new(config, now_nanos)),
}
self.steers_by_latency
.store(bucket.backoff.is_some(), Ordering::Relaxed);
let granted = effective_rate(self.rate(), bucket.backoff.as_ref());
let moved = self.switch_rate(&mut bucket, granted, now_nanos);
drop(bucket);
if moved {
self.rate_moved();
}
}
#[cfg(feature = "std")]
pub fn set_latency_backoff(&self, backoff: Option<LatencyBackoff>) {
self.set_latency_backoff_at(backoff, Self::std_now());
}
#[must_use]
pub(crate) fn steers_by_latency(&self) -> bool {
self.steers_by_latency.load(Ordering::Relaxed)
&& self.rate_bytes_per_sec.load(Ordering::Relaxed) != 0
}
#[must_use]
pub fn latency_backoff(&self) -> Option<LatencyBackoff> {
self.bucket.lock().backoff.as_ref().map(|b| b.config)
}
pub fn record_read_latency_at(&self, latency: Duration, now: Duration) {
if self.rate_bytes_per_sec.load(Ordering::Relaxed) == 0 {
return;
}
let now_nanos = now.as_nanos();
let mut bucket = self.bucket.lock();
let Some(backoff) = bucket.backoff.as_mut() else {
return;
};
backoff.observe(latency, now_nanos);
let granted = effective_rate(self.rate(), bucket.backoff.as_ref());
let moved = self.switch_rate(&mut bucket, granted, now_nanos);
drop(bucket);
if moved {
self.rate_moved();
}
}
#[cfg(feature = "std")]
pub fn record_read_latency(&self, latency: Duration) {
self.record_read_latency_at(latency, Self::std_now());
}
pub fn set_rate_at(&self, bytes_per_sec: u64, now: Duration) {
let now_nanos = now.as_nanos();
let mut bucket = self.bucket.lock();
self.configured_bytes_per_sec
.store(bytes_per_sec, Ordering::Relaxed);
let granted = effective_rate(bytes_per_sec, bucket.backoff.as_ref());
let moved = self.switch_rate(&mut bucket, granted, now_nanos);
drop(bucket);
if moved {
self.rate_moved();
}
}
#[cfg(feature = "std")]
pub fn set_rate(&self, bytes_per_sec: u64) {
self.set_rate_at(bytes_per_sec, Self::std_now());
}
#[must_use]
#[cfg_attr(
not(feature = "std"),
allow(
dead_code,
reason = "no_std-ready token-bucket decision; awaits a clock-injecting no_std caller"
)
)]
pub fn acquire_wait(&self, bytes: u64, now: Duration) -> Duration {
if self.rate_bytes_per_sec.load(Ordering::Relaxed) == 0 {
return Duration::ZERO;
}
let mut bucket = self.bucket.lock();
let rate = self.effective_rate();
if rate == 0 {
return Duration::ZERO;
}
bucket.refill(rate, now.as_nanos());
let position = bucket.debit(bytes);
bucket.wait_for(position, rate)
}
#[cfg(feature = "std")]
#[must_use]
pub fn reserve(&self, bytes: u64) -> Duration {
self.acquire_wait(bytes, Self::std_now())
}
#[cfg(feature = "std")]
pub fn request_interruptible(&self, bytes: u64, should_stop: impl Fn() -> bool) -> bool {
self.request(bytes, should_stop, false)
}
#[cfg(feature = "std")]
pub fn request_abortable(&self, bytes: u64, should_stop: impl Fn() -> bool) -> bool {
self.request(bytes, should_stop, true)
}
#[cfg(feature = "std")]
fn request(&self, bytes: u64, should_stop: impl Fn() -> bool, withdraw_on_stop: bool) -> bool {
if self.rate_bytes_per_sec.load(Ordering::Relaxed) == 0 {
return false;
}
if should_stop() {
return true;
}
let Some(ticket) = self.take_ticket(bytes, Self::std_now()) else {
return false;
};
loop {
let seen = self.wakeup.generation();
if should_stop() {
if withdraw_on_stop {
self.withdraw(&ticket, bytes);
}
return true;
}
let Some(wait) = self.remaining(&ticket, Self::std_now()) else {
return false;
};
self.wakeup.sleep(seen, wait);
}
}
#[cfg(feature = "std")]
fn take_ticket(&self, bytes: u64, now: Duration) -> Option<Ticket> {
let mut bucket = self.bucket.lock();
let rate = self.effective_rate();
if rate == 0 {
return None;
}
bucket.refill(rate, now.as_nanos());
Some(Ticket {
position: bucket.debit(bytes),
taken_at: bucket.changes,
})
}
#[cfg(feature = "std")]
fn remaining(&self, ticket: &Ticket, now: Duration) -> Option<Duration> {
let mut bucket = self.bucket.lock();
if bucket.last_off > ticket.taken_at {
return None;
}
let rate = self.effective_rate();
debug_assert_ne!(rate, 0);
bucket.refill(rate, now.as_nanos());
let wait = bucket.wait_for(ticket.position, rate);
(!wait.is_zero()).then_some(wait)
}
#[cfg(feature = "std")]
fn withdraw(&self, ticket: &Ticket, bytes: u64) {
let mut bucket = self.bucket.lock();
if bucket.last_off > ticket.taken_at {
return;
}
bucket.withdrawn.push((ticket.position, u128::from(bytes)));
bucket.cap(self.effective_rate());
bucket.settle_withdrawn();
drop(bucket);
self.wakeup.ring();
}
#[cfg(not(feature = "std"))]
pub fn request_interruptible(&self, _bytes: u64, should_stop: impl Fn() -> bool) -> bool {
should_stop()
}
#[cfg(not(feature = "std"))]
pub fn request_abortable(&self, _bytes: u64, should_stop: impl Fn() -> bool) -> bool {
should_stop()
}
#[cfg(feature = "std")]
fn std_now() -> Duration {
use std::sync::OnceLock;
static ORIGIN: OnceLock<std::time::Instant> = OnceLock::new();
ORIGIN.get_or_init(std::time::Instant::now).elapsed()
}
}
#[cfg(test)]
#[expect(clippy::unwrap_used, clippy::expect_used, reason = "test code")]
mod tests;