Skip to content

Commit

Permalink
initial fixed window attempt
Browse files Browse the repository at this point in the history
  • Loading branch information
Tevinthuku committed Apr 18, 2024
1 parent 4f871b6 commit bc4e2a2
Show file tree
Hide file tree
Showing 4 changed files with 103 additions and 19 deletions.
2 changes: 1 addition & 1 deletion ratelimiter/src/main.rs
Original file line number Diff line number Diff line change
@@ -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 {
Expand Down
80 changes: 80 additions & 0 deletions ratelimiter/src/rate_limiters/fixed_window_counter.rs
Original file line number Diff line number Diff line change
@@ -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
}
}
}
18 changes: 1 addition & 17 deletions ratelimiter/src/rate_limiters/mod.rs
Original file line number Diff line number Diff line change
@@ -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()
}
}
22 changes: 21 additions & 1 deletion ratelimiter/src/rate_limiters/token_bucket.rs
Original file line number Diff line number Diff line change
@@ -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 {
Expand Down

0 comments on commit bc4e2a2

Please sign in to comment.