sofie_http/
throttle.rs

1use std::collections::HashMap;
2use std::net::IpAddr;
3use std::sync::Mutex;
4use std::time::Duration;
5
6use sofie_core::{Timestamp, Uuid, unpoisoned};
7
8const SWEEP_FLOOR: usize = 1024;
9/// A few megabytes: past it the requesters closest to refilled are forgotten, so a flood can
10/// neither exhaust memory nor lock newcomers out.
11const MAX_REQUESTERS: usize = 65_536;
12const EVICTED_WHEN_FULL: usize = MAX_REQUESTERS / 16;
13
14/// A token bucket: all at once, then one more each `period / count`.
15#[derive(Clone, Copy, Debug, PartialEq, Eq)]
16pub struct Rate {
17    count: u32,
18    period: Duration,
19}
20
21impl Rate {
22    #[must_use]
23    pub const fn per_second(count: u32) -> Self {
24        Self::new(count, Duration::from_secs(1))
25    }
26
27    #[must_use]
28    pub const fn per_minute(count: u32) -> Self {
29        Self::new(count, Duration::from_secs(60))
30    }
31
32    #[must_use]
33    pub const fn per_hour(count: u32) -> Self {
34        Self::new(count, Duration::from_hours(1))
35    }
36
37    const fn new(count: u32, period: Duration) -> Self {
38        assert!(count > 0, "a rate allows at least one request");
39        Self { count, period }
40    }
41
42    fn interval_micros(self) -> i64 {
43        let period = i64::try_from(self.period.as_micros()).unwrap_or(i64::MAX);
44        (period / i64::from(self.count)).max(1)
45    }
46}
47
48/// An IPv6 client is counted by its /64, which one subscriber usually holds whole.
49#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
50pub(crate) enum Requester {
51    Principal(Uuid),
52    Address(IpAddr),
53}
54
55impl Requester {
56    pub(crate) fn of(principal: Option<Uuid>, address: IpAddr) -> Self {
57        if let Some(principal) = principal {
58            return Self::Principal(principal);
59        }
60        match address.to_canonical() {
61            IpAddr::V4(v4) => Self::Address(v4.into()),
62            IpAddr::V6(v6) => {
63                Self::Address(IpAddr::V6((v6.to_bits() & !u128::from(u64::MAX)).into()))
64            }
65        }
66    }
67}
68
69/// Per process: behind a load balancer the effective limit is the rate times the instance count.
70pub(crate) struct Limiter {
71    rate: Rate,
72    requesters: Mutex<Requesters>,
73}
74
75/// The generic cell rate algorithm's one number per caller, in microseconds.
76struct Requesters {
77    due: HashMap<Requester, i64>,
78    sweep_at: usize,
79}
80
81impl Limiter {
82    pub(crate) fn new(rate: Rate) -> Self {
83        Self {
84            rate,
85            requesters: Mutex::new(Requesters {
86                due: HashMap::new(),
87                sweep_at: SWEEP_FLOOR,
88            }),
89        }
90    }
91
92    pub(crate) fn admit(&self, requester: Requester, now: Timestamp) -> Result<(), Duration> {
93        let now = now.as_microsecond();
94        let mut requesters = unpoisoned::lock(&self.requesters);
95        requesters.make_room_for(requester, now);
96        let due = self.due(&requesters, requester, now)?;
97        requesters
98            .due
99            .insert(requester, due + self.rate.interval_micros());
100        requesters.sweep(now);
101        Ok(())
102    }
103
104    /// As `admit`, without counting this request.
105    #[cfg(any(test, feature = "tokens"))]
106    pub(crate) fn check(&self, requester: Requester, now: Timestamp) -> Result<(), Duration> {
107        let requesters = unpoisoned::lock(&self.requesters);
108        self.due(&requesters, requester, now.as_microsecond())
109            .map(drop)
110    }
111
112    fn due(
113        &self,
114        requesters: &Requesters,
115        requester: Requester,
116        now: i64,
117    ) -> Result<i64, Duration> {
118        let tolerance = self.rate.interval_micros() * i64::from(self.rate.count - 1);
119        let due = requesters
120            .due
121            .get(&requester)
122            .map_or(now, |due| (*due).max(now));
123        let early = due - now;
124        if early > tolerance {
125            let wait = u64::try_from(early - tolerance).unwrap_or(u64::MAX);
126            return Err(Duration::from_micros(wait));
127        }
128        Ok(due)
129    }
130}
131
132impl Requesters {
133    fn make_room_for(&mut self, requester: Requester, now: i64) {
134        if self.due.len() < MAX_REQUESTERS || self.due.contains_key(&requester) {
135            return;
136        }
137        self.due.retain(|_, due| *due > now);
138        if self.due.len() < MAX_REQUESTERS {
139            return;
140        }
141        let mut dues: Vec<i64> = self.due.values().copied().collect();
142        let (_, cutoff, _) = dues.select_nth_unstable(EVICTED_WHEN_FULL);
143        let cutoff = *cutoff;
144        self.due.retain(|_, due| *due > cutoff);
145    }
146
147    /// Amortized: sweeps only once the map doubled since the last sweep.
148    fn sweep(&mut self, now: i64) {
149        if self.due.len() < self.sweep_at {
150            return;
151        }
152        self.due.retain(|_, due| *due > now);
153        self.sweep_at = (self.due.len() * 2).max(SWEEP_FLOOR);
154    }
155}
156
157#[cfg(test)]
158mod tests {
159    use super::*;
160
161    fn at(micros: i64) -> Timestamp {
162        Timestamp::from_microsecond(1_700_000_000_000_000 + micros).unwrap()
163    }
164
165    fn requester(last: u8) -> Requester {
166        Requester::of(None, IpAddr::from([192, 0, 2, last]))
167    }
168
169    #[test]
170    fn a_burst_of_the_whole_rate_passes_then_one_per_interval() {
171        let limiter = Limiter::new(Rate::per_second(4));
172
173        for _ in 0..4 {
174            assert!(limiter.admit(requester(1), at(0)).is_ok());
175        }
176        assert_eq!(
177            limiter.admit(requester(1), at(0)),
178            Err(Duration::from_millis(250))
179        );
180        assert!(
181            limiter.admit(requester(2), at(0)).is_ok(),
182            "requesters count apart"
183        );
184        assert!(limiter.admit(requester(1), at(250_000)).is_ok());
185        assert!(limiter.admit(requester(1), at(250_000)).is_err());
186    }
187
188    #[test]
189    fn a_check_tells_a_spent_budget_without_spending_any() {
190        let limiter = Limiter::new(Rate::per_second(2));
191
192        assert!(limiter.check(requester(1), at(0)).is_ok());
193        assert!(limiter.check(requester(1), at(0)).is_ok());
194        limiter.admit(requester(1), at(0)).unwrap();
195        limiter.admit(requester(1), at(0)).unwrap();
196
197        assert_eq!(
198            limiter.check(requester(1), at(0)),
199            Err(Duration::from_millis(500))
200        );
201    }
202
203    #[test]
204    fn a_principal_is_counted_over_its_address_and_ipv6_by_its_64() {
205        let principal = Uuid::now_v7().unwrap();
206        let first: IpAddr = "2001:db8::1".parse().unwrap();
207        let second: IpAddr = "2001:db8::ffff:2".parse().unwrap();
208
209        assert_eq!(Requester::of(None, first), Requester::of(None, second));
210        assert_eq!(
211            Requester::of(Some(principal), first),
212            Requester::Principal(principal)
213        );
214    }
215
216    #[test]
217    fn an_ipv4_mapped_address_is_counted_as_its_ipv4_address() {
218        let mapped = |text: &str| Requester::of(None, text.parse().unwrap());
219
220        assert_eq!(mapped("::ffff:192.0.2.1"), mapped("192.0.2.1"));
221        assert_ne!(mapped("::ffff:192.0.2.1"), mapped("::ffff:198.51.100.1"));
222    }
223
224    #[test]
225    fn requesters_whose_budget_refilled_are_forgotten() {
226        let limiter = Limiter::new(Rate::per_second(1));
227        let wave = |first: u8, now| {
228            for third in 0..4 {
229                for last in 0..=u8::MAX {
230                    let address = IpAddr::from([10, first, third, last]);
231                    limiter.admit(Requester::of(None, address), now).unwrap();
232                }
233            }
234        };
235
236        wave(0, at(0));
237        wave(1, at(2_000_000));
238
239        let requesters = limiter.requesters.lock().unwrap();
240        assert_eq!(requesters.due.len(), 1024);
241    }
242
243    #[test]
244    fn a_full_table_forgets_the_requesters_closest_to_refilled() {
245        let limiter = Limiter::new(Rate::per_second(1));
246        let address = |index: u32| Requester::of(None, IpAddr::from(index.to_be_bytes()));
247        let table_size = u32::try_from(MAX_REQUESTERS).unwrap();
248        {
249            let mut requesters = limiter.requesters.lock().unwrap();
250            for index in 0..table_size {
251                let due = at(2_000_000 + i64::from(index)).as_microsecond();
252                requesters.due.insert(address(index), due);
253            }
254            requesters.sweep_at = usize::MAX;
255        }
256
257        assert!(limiter.admit(address(u32::MAX), at(0)).is_ok());
258        let requesters = limiter.requesters.lock().unwrap();
259        assert!(requesters.due.len() < MAX_REQUESTERS);
260        assert!(!requesters.due.contains_key(&address(0)));
261        assert!(
262            requesters.due.contains_key(&address(table_size - 1)),
263            "the requester furthest from refilled is still counted"
264        );
265    }
266}
sofie_http/throttle.rs

sofie_http/
throttle.rs
