diff --git a/ratelimiter/Cargo.toml b/ratelimiter/Cargo.toml index 1de26d6..696c749 100644 --- a/ratelimiter/Cargo.toml +++ b/ratelimiter/Cargo.toml @@ -7,4 +7,5 @@ edition = "2021" [dependencies] actix-web = "4" +chrono = "0.4.38" tokio = { version = "1.37.0", features = ["full"] } diff --git a/ratelimiter/src/rate_limiters/fixed_window_counter.rs b/ratelimiter/src/rate_limiters/fixed_window_counter.rs index 54f4155..8bda693 100644 --- a/ratelimiter/src/rate_limiters/fixed_window_counter.rs +++ b/ratelimiter/src/rate_limiters/fixed_window_counter.rs @@ -3,7 +3,7 @@ use std::{ sync::{Arc, Mutex}, }; -use tokio::time::Instant; +use chrono::{DateTime, Timelike, Utc}; use super::Ip; @@ -17,35 +17,29 @@ impl IpRateLimiter { let bucket = self .buckets .entry(ip) - .or_insert_with(|| FixedWindowCounter::new(std::time::Duration::from_secs(60), 60)); + .or_insert_with(|| FixedWindowCounter::new(60)); bucket.consume_token() } } #[derive(Debug, Clone)] pub struct FixedWindowCounter { - window_size: std::time::Duration, max_requests: usize, current_window: Arc>, } impl FixedWindowCounter { - pub fn new(window_size: std::time::Duration, max_requests: usize) -> Self { + pub fn new(max_requests: usize) -> Self { FixedWindowCounter { - window_size, max_requests, - current_window: Arc::new(Mutex::new(Window::new_starting_now( - window_size, - max_requests, - ))), + current_window: Arc::new(Mutex::new(Window::new(max_requests))), } } pub fn consume_token(&mut self) -> bool { let mut current_window = self.current_window.lock().unwrap(); if current_window.is_expired() { - *current_window = Window::new_starting_now(self.window_size, self.max_requests); - return current_window.consume_token(); + *current_window = Window::new(self.max_requests); } current_window.consume_token() } @@ -53,20 +47,24 @@ impl FixedWindowCounter { #[derive(Debug, Clone, Copy)] struct Window { - end: Instant, + end: DateTime, remaining_requests: usize, } impl Window { - fn new_starting_now(window_size: std::time::Duration, max_requests: usize) -> Self { - let now = Instant::now(); + fn new(max_requests: usize) -> Self { + let date_time_now = Utc::now(); + // safe to unwrap as we are using a valid seconds value of 59. + // this window ends at the 59th second of the current minute. + let end = date_time_now.with_second(59).unwrap(); + Window { - end: now + window_size, + end, remaining_requests: max_requests, } } fn is_expired(&self) -> bool { - Instant::now() >= self.end + Utc::now() >= self.end } fn consume_token(&mut self) -> bool {