robinpath_modules/modules/
ratelimit_mod.rs

1use robinpath::{RobinPath, Value};
2use std::sync::{LazyLock, Mutex};
3use std::collections::HashMap;
4
5enum Limiter {
6    TokenBucket {
7        tokens: f64,
8        max_tokens: f64,
9        refill_rate: f64,
10        last_refill: u64,
11    },
12    SlidingWindow {
13        requests: Vec<u64>,
14        window_ms: u64,
15        max_requests: u64,
16    },
17    FixedWindow {
18        count: u64,
19        window_ms: u64,
20        max_requests: u64,
21        window_start: u64,
22    },
23}
24
25static LIMITERS: LazyLock<Mutex<HashMap<String, Limiter>>> =
26    LazyLock::new(|| Mutex::new(HashMap::new()));
27
28fn now_ms() -> u64 {
29    std::time::SystemTime::now()
30        .duration_since(std::time::UNIX_EPOCH)
31        .unwrap_or_default()
32        .as_millis() as u64
33}
34
35pub fn register(rp: &mut RobinPath) {
36    // ratelimit.create name type options → config
37    rp.register_builtin("ratelimit.create", |args, _| {
38        let name = args.first().map(|v| v.to_display_string()).unwrap_or_else(|| "default".to_string());
39        let limiter_type = args.get(1).map(|v| v.to_display_string()).unwrap_or_else(|| "token-bucket".to_string());
40        let opts = args.get(2).cloned().unwrap_or(Value::Null);
41        let now = now_ms();
42        let limiter = match limiter_type.as_str() {
43            "sliding-window" => {
44                let window_ms = get_opt_num(&opts, "windowMs", 60000.0) as u64;
45                let max_requests = get_opt_num(&opts, "maxRequests", 100.0) as u64;
46                Limiter::SlidingWindow { requests: Vec::new(), window_ms, max_requests }
47            }
48            "fixed-window" => {
49                let window_ms = get_opt_num(&opts, "windowMs", 60000.0) as u64;
50                let max_requests = get_opt_num(&opts, "maxRequests", 100.0) as u64;
51                Limiter::FixedWindow { count: 0, window_ms, max_requests, window_start: now }
52            }
53            _ => {
54                let max_tokens = get_opt_num(&opts, "maxTokens", 10.0);
55                let refill_rate = get_opt_num(&opts, "refillRate", 1.0);
56                Limiter::TokenBucket { tokens: max_tokens, max_tokens, refill_rate, last_refill: now }
57            }
58        };
59        LIMITERS.lock().unwrap().insert(name.clone(), limiter);
60        let mut obj = indexmap::IndexMap::new();
61        obj.insert("name".to_string(), Value::String(name));
62        obj.insert("type".to_string(), Value::String(limiter_type));
63        Ok(Value::Object(obj))
64    });
65
66    // ratelimit.acquire name count? → bool
67    rp.register_builtin("ratelimit.acquire", |args, _| {
68        let name = args.first().map(|v| v.to_display_string()).unwrap_or_else(|| "default".to_string());
69        let count = args.get(1).map(|v| v.to_number()).unwrap_or(1.0);
70        let mut limiters = LIMITERS.lock().unwrap();
71        let limiter = limiters.get_mut(&name)
72            .ok_or_else(|| format!("Rate limiter \"{}\" not found", name))?;
73        let now = now_ms();
74        let allowed = match limiter {
75            Limiter::TokenBucket { tokens, max_tokens, refill_rate, last_refill } => {
76                let elapsed = (now - *last_refill) as f64 / 1000.0;
77                *tokens = (*tokens + elapsed * *refill_rate).min(*max_tokens);
78                *last_refill = now;
79                if *tokens >= count {
80                    *tokens -= count;
81                    true
82                } else {
83                    false
84                }
85            }
86            Limiter::SlidingWindow { requests, window_ms, max_requests } => {
87                requests.retain(|&t| now - t < *window_ms);
88                if (requests.len() as u64) < *max_requests {
89                    for _ in 0..count as u64 {
90                        requests.push(now);
91                    }
92                    true
93                } else {
94                    false
95                }
96            }
97            Limiter::FixedWindow { count: c, window_ms, max_requests, window_start } => {
98                if now - *window_start >= *window_ms {
99                    *c = 0;
100                    *window_start = now;
101                }
102                if *c + count as u64 <= *max_requests {
103                    *c += count as u64;
104                    true
105                } else {
106                    false
107                }
108            }
109        };
110        Ok(Value::Bool(allowed))
111    });
112
113    // ratelimit.check name → bool
114    rp.register_builtin("ratelimit.check", |args, _| {
115        let name = args.first().map(|v| v.to_display_string()).unwrap_or_else(|| "default".to_string());
116        let mut limiters = LIMITERS.lock().unwrap();
117        let limiter = limiters.get_mut(&name)
118            .ok_or_else(|| format!("Rate limiter \"{}\" not found", name))?;
119        let now = now_ms();
120        let allowed = match limiter {
121            Limiter::TokenBucket { tokens, max_tokens, refill_rate, last_refill } => {
122                let elapsed = (now - *last_refill) as f64 / 1000.0;
123                let current = (*tokens + elapsed * *refill_rate).min(*max_tokens);
124                current >= 1.0
125            }
126            Limiter::SlidingWindow { requests, window_ms, max_requests } => {
127                let active = requests.iter().filter(|&&t| now - t < *window_ms).count() as u64;
128                active < *max_requests
129            }
130            Limiter::FixedWindow { count, window_ms, max_requests, window_start } => {
131                (if now - *window_start >= *window_ms { 0 } else { *count }) < *max_requests
132            }
133        };
134        Ok(Value::Bool(allowed))
135    });
136
137    // ratelimit.remaining name → number
138    rp.register_builtin("ratelimit.remaining", |args, _| {
139        let name = args.first().map(|v| v.to_display_string()).unwrap_or_else(|| "default".to_string());
140        let mut limiters = LIMITERS.lock().unwrap();
141        let limiter = limiters.get_mut(&name)
142            .ok_or_else(|| format!("Rate limiter \"{}\" not found", name))?;
143        let now = now_ms();
144        let remaining = match limiter {
145            Limiter::TokenBucket { tokens, max_tokens, refill_rate, last_refill } => {
146                let elapsed = (now - *last_refill) as f64 / 1000.0;
147                (*tokens + elapsed * *refill_rate).min(*max_tokens).floor()
148            }
149            Limiter::SlidingWindow { requests, window_ms, max_requests } => {
150                let active = requests.iter().filter(|&&t| now - t < *window_ms).count() as f64;
151                (*max_requests as f64 - active).max(0.0)
152            }
153            Limiter::FixedWindow { count, window_ms, max_requests, window_start } => {
154                let c = if now - *window_start >= *window_ms { 0 } else { *count };
155                (*max_requests - c) as f64
156            }
157        };
158        Ok(Value::Number(remaining))
159    });
160
161    // ratelimit.reset name → bool
162    rp.register_builtin("ratelimit.reset", |args, _| {
163        let name = args.first().map(|v| v.to_display_string()).unwrap_or_else(|| "default".to_string());
164        let mut limiters = LIMITERS.lock().unwrap();
165        let limiter = limiters.get_mut(&name)
166            .ok_or_else(|| format!("Rate limiter \"{}\" not found", name))?;
167        let now = now_ms();
168        match limiter {
169            Limiter::TokenBucket { tokens, max_tokens, last_refill, .. } => {
170                *tokens = *max_tokens;
171                *last_refill = now;
172            }
173            Limiter::SlidingWindow { requests, .. } => { requests.clear(); }
174            Limiter::FixedWindow { count, window_start, .. } => {
175                *count = 0;
176                *window_start = now;
177            }
178        }
179        Ok(Value::Bool(true))
180    });
181
182    // ratelimit.status name → {name, remaining, ...}
183    rp.register_builtin("ratelimit.status", |args, _| {
184        let name = args.first().map(|v| v.to_display_string()).unwrap_or_else(|| "default".to_string());
185        let mut limiters = LIMITERS.lock().unwrap();
186        let limiter = limiters.get_mut(&name)
187            .ok_or_else(|| format!("Rate limiter \"{}\" not found", name))?;
188        let now = now_ms();
189        let mut obj = indexmap::IndexMap::new();
190        obj.insert("name".to_string(), Value::String(name));
191        match limiter {
192            Limiter::TokenBucket { tokens, max_tokens, refill_rate, last_refill } => {
193                let elapsed = (now - *last_refill) as f64 / 1000.0;
194                let current = (*tokens + elapsed * *refill_rate).min(*max_tokens);
195                obj.insert("type".to_string(), Value::String("token-bucket".to_string()));
196                obj.insert("tokens".to_string(), Value::Number(current.floor()));
197                obj.insert("maxTokens".to_string(), Value::Number(*max_tokens));
198                obj.insert("refillRate".to_string(), Value::Number(*refill_rate));
199            }
200            Limiter::SlidingWindow { requests, window_ms, max_requests } => {
201                let active = requests.iter().filter(|&&t| now - t < *window_ms).count() as f64;
202                obj.insert("type".to_string(), Value::String("sliding-window".to_string()));
203                obj.insert("used".to_string(), Value::Number(active));
204                obj.insert("remaining".to_string(), Value::Number((*max_requests as f64 - active).max(0.0)));
205                obj.insert("maxRequests".to_string(), Value::Number(*max_requests as f64));
206            }
207            Limiter::FixedWindow { count, window_ms, max_requests, window_start } => {
208                let c = if now - *window_start >= *window_ms { 0 } else { *count };
209                obj.insert("type".to_string(), Value::String("fixed-window".to_string()));
210                obj.insert("used".to_string(), Value::Number(c as f64));
211                obj.insert("remaining".to_string(), Value::Number((*max_requests - c) as f64));
212                obj.insert("maxRequests".to_string(), Value::Number(*max_requests as f64));
213            }
214        }
215        Ok(Value::Object(obj))
216    });
217
218    // ratelimit.destroy name → bool
219    rp.register_builtin("ratelimit.destroy", |args, _| {
220        let name = args.first().map(|v| v.to_display_string()).unwrap_or_else(|| "default".to_string());
221        Ok(Value::Bool(LIMITERS.lock().unwrap().remove(&name).is_some()))
222    });
223}
224
225fn get_opt_num(opts: &Value, key: &str, default: f64) -> f64 {
226    if let Value::Object(obj) = opts {
227        obj.get(key).map(|v| v.to_number()).unwrap_or(default)
228    } else {
229        default
230    }
231}
robinpath_modules/modules/ratelimit_mod.rs

robinpath_modules/modules/
ratelimit_mod.rs
