basic rate limit

src/rate_limit.rs

@@ -0,0 +1,67 @@
+use std::iter;
+use std::num::NonZeroUsize;
+use std::sync::Mutex;
+use std::time::SystemTime;
+
+use lru::LruCache;
+
+pub struct Limiter {
+    record: Mutex<LruCache<i32, Vec<u64>>>,
+    amount: u64,
+    interval: u64,
+}
+
+impl Limiter {
+    pub fn init(amount: u64, interval: u64) -> Self {
+        Self {
+            record: Mutex::new(LruCache::new(NonZeroUsize::new(2000).unwrap())),
+            amount,
+            interval,
+        }
+    }
+
+    pub fn check(&self, uid: i32) -> bool {
+        let t = SystemTime::now()
+            .duration_since(SystemTime::UNIX_EPOCH)
+            .unwrap()
+            .as_secs();
+        let mut r = self.record.lock().unwrap();
+        if let Some(ts) = r.pop(&uid) {
+            let new_ts: Vec<u64> = ts
+                .into_iter()
+                .chain(iter::once(t))
+                .filter(|&tt| tt + self.interval > t)
+                .collect();
+            let len = new_ts.len() as u64;
+            r.put(uid, new_ts);
+            len < self.amount
+        } else {
+            r.put(uid, vec![t]);
+            true
+        }
+    }
+}
+
+pub struct MainLimiters {
+    post_min: Limiter,
+    post_hour: Limiter,
+    get_hour: Limiter,
+}
+
+impl MainLimiters {
+    pub fn init() -> Self {
+        Self {
+            post_min: Limiter::init(6, 60),
+            post_hour: Limiter::init(50, 3600),
+            get_hour: Limiter::init(1000, 3600),
+        }
+    }
+
+    pub fn check(&self, is_post: bool, uid: i32) -> bool {
+        if is_post {
+            self.post_hour.check(uid) && self.post_min.check(uid)
+        } else {
+            self.get_hour.check(uid)
+        }
+    }
+}