snerd_rust/
rate_limiter.rs

1use chrono::{DateTime, Utc};
2use fs3::FileExt;
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5use std::fs::OpenOptions;
6use std::io::{BufRead, BufReader, Seek, SeekFrom, Write};
7use std::path::PathBuf;
8
9#[derive(Serialize, Deserialize, Debug, Clone)]
10pub struct RateWindow {
11    pub count: i32,
12    pub window_start: DateTime<Utc>,
13}
14
15#[derive(Clone)]
16pub struct RateLimiter {
17    file_path: PathBuf,
18}
19
20impl RateLimiter {
21    pub fn new(tasks_path: &PathBuf) -> Self {
22        let mut path = tasks_path.clone();
23        path.set_file_name("rate_limits.json");
24        Self { file_path: path }
25    }
26
27    pub fn check_and_increment(&self, group: &str, limit: i32) -> std::io::Result<bool> {
28        if let Some(parent) = self.file_path.parent() {
29            std::fs::create_dir_all(parent)?;
30        }
31
32        let mut file = OpenOptions::new()
33            .create(true)
34            .read(true)
35            .write(true)
36            .open(&self.file_path)?;
37
38        file.lock_exclusive()?;
39
40        let mut limits: HashMap<String, RateWindow> = HashMap::new();
41        {
42            let mut reader = BufReader::new(&file);
43            let mut line = String::new();
44            if let Ok(bytes_read) = reader.read_line(&mut line) {
45                if bytes_read > 0 {
46                    if let Ok(parsed) = serde_json::from_str(&line) {
47                        limits = parsed;
48                    }
49                }
50            }
51        }
52
53        let now = Utc::now();
54        let mut allow = false;
55
56        let window = limits.entry(group.to_string()).or_insert(RateWindow {
57            count: 0,
58            window_start: now,
59        });
60
61        if (now - window.window_start).num_seconds() >= 60 {
62            window.count = 0;
63            window.window_start = now;
64        }
65
66        if window.count < limit {
67            window.count += 1;
68            allow = true;
69        }
70
71        file.set_len(0)?;
72        file.seek(SeekFrom::Start(0))?;
73        let out = serde_json::to_string(&limits)?;
74        writeln!(file, "{}", out)?;
75        file.sync_all()?;
76
77        file.unlock()?;
78        Ok(allow)
79    }
80}
snerd_rust/rate_limiter.rs

snerd_rust/
rate_limiter.rs
