// Rate limiting for wschat (token bucket algorithm)
use std::time
use std::sync
pub struct RateLimiter {
capacity: int,
tokens: Arc<Mutex<int>>,
refill_rate: int, // tokens per second,
last_refill: Arc<Mutex<int>>, // timestamp,
}
impl RateLimiter {
pub fn new(messages_per_second: int) -> Self {
RateLimiter {
capacity: messages_per_second * 10, // Allow bursts,
tokens: Arc::new(Mutex::new(messages_per_second * 10)),
refill_rate: messages_per_second,
last_refill: Arc::new(Mutex::new(time.now_unix())),
}
}
pub async fn check(self) -> bool {
self.refill().await;
let mut tokens = self.tokens.lock().await;
if *tokens > 0 {
*tokens -= 1
true
} else {
false
}
}
async fn refill(self) {
let now = time.now_unix();
let mut last_refill = self.last_refill.lock().await;
let elapsed = now - *last_refill;
if elapsed > 0 {
let new_tokens = elapsed * self.refill_rate;
let mut tokens = self.tokens.lock().await;
*tokens = (*tokens + new_tokens).min(self.capacity);
*last_refill = now;
}
}
}