use std::time::Instant;
use crate::bucket::RateLimiter;
#[derive(Clone)]
pub struct TokenBucket {
capacity: f64,
tokens: f64,
refill_rate: f64,
last_refill: Instant,
}
impl TokenBucket {
pub fn new(capacity: f64, refill_rate: f64) -> Self {
Self {
capacity,
refill_rate,
tokens: capacity,
last_refill: Instant::now(),
}
}
pub fn available_tokens(&mut self) -> f64 {
self.refill();
self.tokens
}
pub fn allow_n(&mut self, n: f64) -> bool {
self.refill();
if self.tokens >= n {
self.tokens -= n;
true
} else {
false
}
}
pub fn allow(&mut self) -> bool {
self.allow_n(1.0)
}
fn refill(&mut self) {
let now = Instant::now();
let elapsed = now.duration_since(self.last_refill).as_secs_f64();
let added = elapsed * self.refill_rate;
self.tokens = (self.tokens + added).min(self.capacity);
self.last_refill = now;
}
}
impl RateLimiter for TokenBucket {
fn allow(&mut self) -> bool {
TokenBucket::allow(self)
}
fn allow_n(&mut self, n: f64) -> bool {
TokenBucket::allow_n(self, n)
}
fn available_tokens(&mut self) -> f64 {
TokenBucket::available_tokens(self)
}
}