github_mcp/core/
rate_limiter.rs1use std::sync::Mutex;
4use std::time::{Duration, Instant};
5
6use super::errors::McpifyError;
7
8pub struct RateLimiter {
11 limit: usize,
12 window: Duration,
13 timestamps: Mutex<Vec<Instant>>,
14}
15
16impl RateLimiter {
17 pub fn new(limit: usize, window: Duration) -> Self {
18 Self {
19 limit,
20 window,
21 timestamps: Mutex::new(Vec::new()),
22 }
23 }
24
25 pub fn acquire(&self) -> Result<(), McpifyError> {
26 let now = Instant::now();
27 let mut timestamps = self.timestamps.lock().unwrap();
28 timestamps.retain(|&t| now.duration_since(t) < self.window);
29
30 if timestamps.len() >= self.limit {
31 return Err(McpifyError::RateLimitExceeded);
32 }
33 timestamps.push(now);
34 Ok(())
35 }
36}
37
38impl Default for RateLimiter {
39 fn default() -> Self {
40 Self::new(100, Duration::from_secs(1))
41 }
42}
43
44#[cfg(test)]
45mod tests {
46 use super::*;
47
48 #[test]
49 fn allows_calls_up_to_the_limit() {
50 let limiter = RateLimiter::new(3, Duration::from_secs(1));
51 for _ in 0..3 {
52 assert!(limiter.acquire().is_ok());
53 }
54 }
55
56 #[test]
57 fn rejects_calls_beyond_the_limit_within_the_window() {
58 let limiter = RateLimiter::new(1, Duration::from_secs(60));
59 limiter.acquire().unwrap();
60 let err = limiter.acquire().unwrap_err();
61 assert_eq!(err.code(), "RATE_LIMIT_EXCEEDED");
62 }
63}