feat: init

immediate/rate_limit.rs

@@ -0,0 +1,102 @@
+use std::time::Instant;
+
+use uuid::Uuid;
+
+use crate::cache::redis::AppRedis;
+use crate::error::AppResult;
+use ::redis::Cmd;
+
+use super::redis_keys::*;
+
+pub struct RateLimiter {
+    redis: AppRedis,
+    max_per_sec: u32,
+}
+
+impl RateLimiter {
+    pub fn new(redis: AppRedis) -> Self {
+        Self {
+            redis,
+            max_per_sec: WS_MAX_MESSAGES_PER_SEC,
+        }
+    }
+
+    pub fn with_limit(redis: AppRedis, max_per_sec: u32) -> Self {
+        Self { redis, max_per_sec }
+    }
+
+    pub fn check(&self, connection_id: Uuid) -> AppResult<bool> {
+        let key = format!("{WS_RATE_PREFIX}{connection_id}");
+        let mut conn = self.redis.get_connection()?;
+        let count: i64 = Cmd::new()
+            .arg("INCR")
+            .arg(&key)
+            .query(&mut *conn.inner_mut())
+            .map_err(crate::error::AppError::Redis)?;
+        if count == 1 {
+            let _ = Cmd::new()
+                .arg("EXPIRE")
+                .arg(&key)
+                .arg(1_u64)
+                .query::<()>(&mut *conn.inner_mut());
+        }
+        Ok(count <= self.max_per_sec as i64)
+    }
+
+    pub fn check_sliding(&self, connection_id: Uuid) -> AppResult<bool> {
+        let key = format!("{WS_RATE_PREFIX}{connection_id}");
+        let mut conn = self.redis.get_connection()?;
+        let count: i64 = Cmd::new()
+            .arg("INCR")
+            .arg(&key)
+            .query(&mut *conn.inner_mut())
+            .map_err(crate::error::AppError::Redis)?;
+        if count == 1 {
+            let _ = Cmd::new()
+                .arg("EXPIRE")
+                .arg(&key)
+                .arg(2_u64)
+                .query::<()>(&mut *conn.inner_mut());
+        }
+        Ok(count <= self.max_per_sec as i64)
+    }
+
+    pub fn remaining(&self, connection_id: Uuid) -> AppResult<u32> {
+        let key = format!("{WS_RATE_PREFIX}{connection_id}");
+        let mut conn = self.redis.get_connection()?;
+        let count: Option<i64> = Cmd::new()
+            .arg("GET")
+            .arg(&key)
+            .query(&mut *conn.inner_mut())
+            .map_err(crate::error::AppError::Redis)?;
+        Ok(self.max_per_sec.saturating_sub(count.unwrap_or(0) as u32))
+    }
+}
+
+pub struct LocalRateLimiter {
+    count: std::sync::atomic::AtomicU32,
+    start: std::sync::Mutex<Instant>,
+    max_per_sec: u32,
+}
+
+impl LocalRateLimiter {
+    pub fn new(max_per_sec: u32) -> Self {
+        Self {
+            count: std::sync::atomic::AtomicU32::new(0),
+            start: std::sync::Mutex::new(Instant::now()),
+            max_per_sec,
+        }
+    }
+
+    pub fn check(&self) -> bool {
+        let mut start = self.start.lock().unwrap();
+        if start.elapsed().as_secs() >= 1 {
+            self.count.store(0, std::sync::atomic::Ordering::Relaxed);
+            *start = Instant::now();
+        }
+        drop(start);
+        self.count
+            .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
+            < self.max_per_sec
+    }
+}