fedimint_gateway_server/
rate_limit.rs

1use std::sync::Mutex;
2use std::time::SystemTime;
3
4use fedimint_core::time::now;
5
6/// A global token bucket limiting the rate of unauthenticated requests that
7/// create state on the gateway or its Lightning node.
8///
9/// The bucket holds at most `burst` tokens and refills at `refill_per_second`;
10/// each request takes one token and is rejected if none is available. The
11/// limiter is transport-agnostic since it guards the request handler itself
12/// rather than the HTTP layer, so requests arriving over Iroh are limited as
13/// well.
14#[derive(Debug)]
15pub struct TokenBucketRateLimiter {
16    burst: f64,
17    refill_per_second: f64,
18    state: Mutex<TokenBucketState>,
19}
20
21#[derive(Debug)]
22struct TokenBucketState {
23    tokens: f64,
24    last_refill: SystemTime,
25}
26
27impl TokenBucketRateLimiter {
28    pub fn new(burst: u32, refill_per_second: u32) -> Self {
29        Self {
30            burst: f64::from(burst),
31            refill_per_second: f64::from(refill_per_second),
32            state: Mutex::new(TokenBucketState {
33                tokens: f64::from(burst),
34                last_refill: now(),
35            }),
36        }
37    }
38
39    /// Takes a token from the bucket if one is available, returning whether
40    /// the request may proceed.
41    pub fn try_acquire(&self) -> bool {
42        self.try_acquire_at(now())
43    }
44
45    fn try_acquire_at(&self, now: SystemTime) -> bool {
46        let mut state = self
47            .state
48            .lock()
49            .expect("No code holding the lock can panic");
50
51        // `SystemTime` is not monotonic; if the clock moved backwards, refill
52        // nothing and restart the refill measurement from the earlier time.
53        let elapsed = now
54            .duration_since(state.last_refill)
55            .unwrap_or_default()
56            .as_secs_f64();
57        state.tokens = (state.tokens + elapsed * self.refill_per_second).min(self.burst);
58        state.last_refill = now;
59
60        if state.tokens >= 1.0 {
61            state.tokens -= 1.0;
62            true
63        } else {
64            false
65        }
66    }
67}
68
69#[cfg(test)]
70mod tests;
fedimint_gateway_server/rate_limit.rs

fedimint_gateway_server/
rate_limit.rs
