From bc4e2a2fb32a973a759795f420c72077577b3174 Mon Sep 17 00:00:00 2001
From: Tevin Thuku <tevinthuku@gmail.com>
Date: Thu, 18 Apr 2024 07:45:46 +0300
Subject: [PATCH] initial fixed window attempt

---
 ratelimiter/src/main.rs                       |  2 +-
 .../src/rate_limiters/fixed_window_counter.rs | 80 +++++++++++++++++++
 ratelimiter/src/rate_limiters/mod.rs          | 18 +----
 ratelimiter/src/rate_limiters/token_bucket.rs | 22 ++++-
 4 files changed, 103 insertions(+), 19 deletions(-)
 create mode 100644 ratelimiter/src/rate_limiters/fixed_window_counter.rs

diff --git a/ratelimiter/src/main.rs b/ratelimiter/src/main.rs
index 1f9fe32..7986387 100644
--- a/ratelimiter/src/main.rs
+++ b/ratelimiter/src/main.rs
@@ -1,7 +1,7 @@
 use std::sync::Mutex;
 
 use actix_web::{get, web, App, HttpRequest, HttpResponse, HttpServer, Responder};
-use ratelimiter::rate_limiters::IpRateLimiter;
+use ratelimiter::rate_limiters::fixed_window_counter::IpRateLimiter;
 
 #[get("/limited")]
 async fn limited(data: web::Data<AppStateWithIpRateLimiter>, req: HttpRequest) -> impl Responder {
diff --git a/ratelimiter/src/rate_limiters/fixed_window_counter.rs b/ratelimiter/src/rate_limiters/fixed_window_counter.rs
new file mode 100644
index 0000000..54f4155
--- /dev/null
+++ b/ratelimiter/src/rate_limiters/fixed_window_counter.rs
@@ -0,0 +1,80 @@
+use std::{
+    collections::HashMap,
+    sync::{Arc, Mutex},
+};
+
+use tokio::time::Instant;
+
+use super::Ip;
+
+#[derive(Default, Clone)]
+pub struct IpRateLimiter {
+    buckets: HashMap<Ip, FixedWindowCounter>,
+}
+
+impl IpRateLimiter {
+    pub fn consume_token(&mut self, ip: Ip) -> bool {
+        let bucket = self
+            .buckets
+            .entry(ip)
+            .or_insert_with(|| FixedWindowCounter::new(std::time::Duration::from_secs(60), 60));
+        bucket.consume_token()
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct FixedWindowCounter {
+    window_size: std::time::Duration,
+    max_requests: usize,
+    current_window: Arc<Mutex<Window>>,
+}
+
+impl FixedWindowCounter {
+    pub fn new(window_size: std::time::Duration, max_requests: usize) -> Self {
+        FixedWindowCounter {
+            window_size,
+            max_requests,
+            current_window: Arc::new(Mutex::new(Window::new_starting_now(
+                window_size,
+                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.consume_token()
+    }
+}
+
+#[derive(Debug, Clone, Copy)]
+struct Window {
+    end: Instant,
+    remaining_requests: usize,
+}
+
+impl Window {
+    fn new_starting_now(window_size: std::time::Duration, max_requests: usize) -> Self {
+        let now = Instant::now();
+        Window {
+            end: now + window_size,
+            remaining_requests: max_requests,
+        }
+    }
+    fn is_expired(&self) -> bool {
+        Instant::now() >= self.end
+    }
+
+    fn consume_token(&mut self) -> bool {
+        if self.remaining_requests > 0 {
+            self.remaining_requests -= 1;
+            true
+        } else {
+            false
+        }
+    }
+}
diff --git a/ratelimiter/src/rate_limiters/mod.rs b/ratelimiter/src/rate_limiters/mod.rs
index c3da4fe..44dbecc 100644
--- a/ratelimiter/src/rate_limiters/mod.rs
+++ b/ratelimiter/src/rate_limiters/mod.rs
@@ -1,20 +1,4 @@
-use std::collections::HashMap;
-
+pub mod fixed_window_counter;
 pub mod token_bucket;
 
 pub type Ip = String;
-
-#[derive(Debug, Clone, Default)]
-pub struct IpRateLimiter {
-    buckets: HashMap<Ip, token_bucket::TokenBucket>,
-}
-
-impl IpRateLimiter {
-    pub fn consume_token(&mut self, ip: Ip) -> bool {
-        let bucket = self
-            .buckets
-            .entry(ip)
-            .or_insert_with(|| token_bucket::TokenBucket::new(10, 1));
-        bucket.consume_token()
-    }
-}
diff --git a/ratelimiter/src/rate_limiters/token_bucket.rs b/ratelimiter/src/rate_limiters/token_bucket.rs
index 4657387..cad5561 100644
--- a/ratelimiter/src/rate_limiters/token_bucket.rs
+++ b/ratelimiter/src/rate_limiters/token_bucket.rs
@@ -1,4 +1,24 @@
-use std::sync::{Arc, Mutex};
+use std::{
+    collections::HashMap,
+    sync::{Arc, Mutex},
+};
+
+use super::Ip;
+
+#[derive(Debug, Clone, Default)]
+pub struct IpRateLimiter {
+    buckets: HashMap<Ip, TokenBucket>,
+}
+
+impl IpRateLimiter {
+    pub fn consume_token(&mut self, ip: Ip) -> bool {
+        let bucket = self
+            .buckets
+            .entry(ip)
+            .or_insert_with(|| TokenBucket::new(10, 1));
+        bucket.consume_token()
+    }
+}
 
 #[derive(Debug, Clone)]
 pub struct TokenBucket {