use std::time::Instant;
#[derive(Debug)]
pub(super) struct TokenBucket {
capacity: f64,
tokens: f64,
refill_per_sec: f64,
last_refill: Instant,
}
impl TokenBucket {
pub(super) fn new(max_bytes_per_sec: u64, burst_bytes: u64) -> Self {
Self {
capacity: burst_bytes as f64,
tokens: burst_bytes as f64,
refill_per_sec: max_bytes_per_sec as f64,
last_refill: Instant::now(),
}
}
pub(super) fn try_consume(&mut self, amount: u64) -> bool {
self.refill();
let amount_f = amount as f64;
if self.tokens < amount_f {
return false;
}
self.tokens -= amount_f;
true
}
fn refill(&mut self) {
let now = Instant::now();
let elapsed = now.duration_since(self.last_refill).as_secs_f64();
self.last_refill = now;
let added = elapsed * self.refill_per_sec;
self.tokens = (self.tokens + added).min(self.capacity);
}
}