sparrow/runtime/
ratelimit.rs1use std::collections::HashMap;
4use std::sync::Mutex;
5use std::time::Instant;
6
7#[derive(Debug)]
8pub struct RateLimiter {
9 window_secs: u64,
10 max_requests: u64,
11 burst: u64,
12 buckets: Mutex<HashMap<String, Bucket>>,
13}
14
15#[derive(Debug, Clone)]
16struct Bucket {
17 tokens: u64,
18 last_refill: Instant,
19}
20
21impl RateLimiter {
22 pub fn new(max_requests_per_minute: u64, burst: u64) -> Self {
23 Self {
24 window_secs: 60,
25 max_requests: max_requests_per_minute,
26 burst: burst.max(1),
27 buckets: Mutex::new(HashMap::new()),
28 }
29 }
30
31 pub fn check(&self, user_id: &str) -> (bool, u64, u64) {
34 let mut buckets = self.buckets.lock().unwrap();
35 let now = Instant::now();
36 let bucket = buckets.entry(user_id.to_string()).or_insert(Bucket {
37 tokens: self.burst,
38 last_refill: now,
39 });
40
41 let elapsed = now.duration_since(bucket.last_refill).as_secs();
43 if elapsed > 0 {
44 let refill =
45 (elapsed as f64 / self.window_secs as f64 * self.max_requests as f64) as u64;
46 bucket.tokens = (bucket.tokens + refill).min(self.burst);
47 bucket.last_refill = now;
48 }
49
50 if bucket.tokens > 0 {
51 bucket.tokens -= 1;
52 (true, bucket.tokens, self.window_secs)
53 } else {
54 (false, 0, self.window_secs)
55 }
56 }
57}
58
59impl Default for RateLimiter {
60 fn default() -> Self {
61 Self::new(60, 10)
62 }
63}
64
65use std::sync::Arc;
68use std::sync::atomic::{AtomicBool, Ordering};
69
70pub struct InterruptHandler {
71 interrupted: Arc<AtomicBool>,
72}
73
74impl InterruptHandler {
75 pub fn new() -> Self {
76 let handler = Self {
77 interrupted: Arc::new(AtomicBool::new(false)),
78 };
79 let flag = handler.interrupted.clone();
80
81 ctrlc::set_handler(move || {
83 flag.store(true, Ordering::SeqCst);
84 })
85 .ok();
86
87 handler
88 }
89
90 pub fn is_interrupted(&self) -> bool {
91 self.interrupted.load(Ordering::SeqCst)
92 }
93
94 pub fn reset(&self) {
95 self.interrupted.store(false, Ordering::SeqCst);
96 }
97}
98
99impl Default for InterruptHandler {
100 fn default() -> Self {
101 Self::new()
102 }
103}