superfluid_daemon/
ratelimit.rs

1//! Token-bucket rate limiter keyed by client IP (`--rate-limit N`).
2
3use std::collections::HashMap;
4use std::net::IpAddr;
5use std::sync::Mutex;
6use std::time::Instant;
7
8const MAX_BUCKETS: usize = 100_000;
9const IDLE_EVICT_SECS: f64 = 300.0;
10const SWEEP_EVERY: u32 = 1000;
11
12struct Bucket {
13    tokens: f64,
14    last_refill: Instant,
15}
16
17impl Bucket {
18    fn take(&mut self, rpm: f64, now: Instant) -> bool {
19        let elapsed = now.duration_since(self.last_refill).as_secs_f64();
20        self.tokens = (self.tokens + elapsed * (rpm / 60.0)).min(rpm);
21        self.last_refill = now;
22        if self.tokens >= 1.0 {
23            self.tokens -= 1.0;
24            return true;
25        }
26        false
27    }
28}
29
30pub struct TokenBucket {
31    rpm: u32,
32    bucket: Mutex<Bucket>,
33}
34
35impl TokenBucket {
36    pub fn new(rpm: u32) -> TokenBucket {
37        TokenBucket { rpm, bucket: Mutex::new(Bucket { tokens: rpm as f64, last_refill: Instant::now() }) }
38    }
39
40    pub fn allow(&self) -> bool {
41        self.rpm == 0 || self.bucket.lock().expect("token bucket").take(self.rpm as f64, Instant::now())
42    }
43
44    pub fn next_token_in(&self) -> std::time::Duration {
45        if self.rpm == 0 {
46            return std::time::Duration::ZERO;
47        }
48        let b = self.bucket.lock().expect("token bucket");
49        let rpm = self.rpm as f64;
50        let now_tokens = (b.tokens + Instant::now().duration_since(b.last_refill).as_secs_f64() * rpm / 60.0).min(rpm);
51        std::time::Duration::from_secs_f64(((1.0 - now_tokens).max(0.0)) * 60.0 / rpm)
52    }
53}
54
55pub struct RateLimiter {
56    rpm: u32,
57    state: Mutex<State>,
58}
59
60struct State {
61    buckets: HashMap<IpAddr, Bucket>,
62    sweep_counter: u32,
63}
64
65impl RateLimiter {
66    pub fn new(rpm: u32) -> RateLimiter {
67        RateLimiter {
68            rpm,
69            state: Mutex::new(State { buckets: HashMap::new(), sweep_counter: 0 }),
70        }
71    }
72
73    pub fn allow(&self, client: IpAddr) -> bool {
74        if self.rpm == 0 {
75            return true;
76        }
77        let now = Instant::now();
78        let mut st = self.state.lock().expect("rate limiter");
79        st.sweep_counter += 1;
80        if st.sweep_counter >= SWEEP_EVERY {
81            st.sweep_counter = 0;
82            evict_idle(&mut st.buckets, now);
83        }
84        if st.buckets.len() >= MAX_BUCKETS && !st.buckets.contains_key(&client) {
85            evict_idle(&mut st.buckets, now);
86            if st.buckets.len() >= MAX_BUCKETS {
87                evict_least_recent(&mut st.buckets);
88            }
89        }
90        let rpm = self.rpm as f64;
91        st.buckets
92            .entry(client)
93            .or_insert(Bucket { tokens: rpm, last_refill: now })
94            .take(rpm, now)
95    }
96
97    pub fn bucket_count(&self) -> usize {
98        self.state.lock().expect("rate limiter").buckets.len()
99    }
100}
101
102fn evict_idle(buckets: &mut HashMap<IpAddr, Bucket>, now: Instant) {
103    buckets.retain(|_, b| now.duration_since(b.last_refill).as_secs_f64() <= IDLE_EVICT_SECS);
104}
105
106fn evict_least_recent(buckets: &mut HashMap<IpAddr, Bucket>) {
107    if buckets.is_empty() {
108        return;
109    }
110    let mut stamps: Vec<Instant> = buckets.values().map(|b| b.last_refill).collect();
111    let drop = (stamps.len() / 10 + 1).min(stamps.len() - 1);
112    stamps.sort_unstable();
113    let cutoff = stamps[drop];
114    buckets.retain(|_, b| b.last_refill >= cutoff);
115}
superfluid_daemon/ratelimit.rs

superfluid_daemon/
ratelimit.rs
