46 lines
1.2 KiB
Rust
46 lines
1.2 KiB
Rust
|
//! App-level token bucket for incoming rendezvous accepts.
|
||
|
|
|
||
|
|
use std::time::{Duration, Instant};
|
||
|
|
|
||
|
|
pub struct TokenBucket {
|
||
|
|
rate_per_sec: f64,
|
||
|
|
burst: f64,
|
||
|
|
tokens: f64,
|
||
|
|
last: Instant,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl TokenBucket {
|
||
|
|
/// `count` tokens replenished over `window`, starting full up to `burst`.
|
||
|
|
pub fn new(count: u32, window: Duration, burst: u32) -> Self {
|
||
|
|
let secs = window.as_secs_f64().max(f64::EPSILON);
|
||
|
|
Self {
|
||
|
|
rate_per_sec: f64::from(count) / secs,
|
||
|
|
burst: f64::from(burst),
|
||
|
|
tokens: f64::from(burst),
|
||
|
|
last: Instant::now(),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn try_acquire(&mut self) -> bool {
|
||
|
|
self.try_acquire_at(Instant::now())
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn try_acquire_at(&mut self, now: Instant) -> bool {
|
||
|
|
let elapsed = now.saturating_duration_since(self.last).as_secs_f64();
|
||
|
|
self.last = now;
|
||
|
|
self.tokens = (self.tokens + elapsed * self.rate_per_sec).min(self.burst);
|
||
|
|
if self.tokens >= 1.0 {
|
||
|
|
self.tokens -= 1.0;
|
||
|
|
true
|
||
|
|
} else {
|
||
|
|
false
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
impl Default for TokenBucket {
|
||
|
|
fn default() -> Self {
|
||
|
|
Self::new(30, Duration::from_secs(60), 10)
|
||
|
|
}
|
||
|
|
}
|