simple_rate_limit/
lib.rs

1use ring_vec::RingVec;
2use std::time::{Duration, Instant};
3
4/// How many times an event may occur in a given [`Duration`].
5#[derive(Clone, Copy, Debug, Eq, PartialEq)]
6pub struct RateLimit {
7    pub count: usize,
8    pub period: Duration,
9}
10
11impl RateLimit {
12    /// Creates a new [`RateLimit`] with the provided values, but only
13    /// if the count is not zero (otherwise, what's the point?)
14    pub fn new(count: usize, period: Duration) -> Option<RateLimit> {
15        if count > 0 {
16            Some(RateLimit { count, period })
17        } else {
18            None
19        }
20    }
21}
22
23/// Enforces a [`RateLimit`] by maintaining a ring vector of
24/// timestamps capped at the [`RateLimit::count`].
25pub struct RateLimiter {
26    pub limit: RateLimit,
27    readings: RingVec<Instant>,
28}
29
30impl RateLimiter {
31    /// Create without preallocating storage. Ideal if it may go unused.
32    pub fn new(limit: RateLimit) -> RateLimiter {
33        RateLimiter { limit, readings: RingVec::new(limit.count) }
34    }
35
36    /// Create with preallocated storage. Ideal if you're likely to
37    /// use it a lot to avoid resizing during fill.
38    pub fn new_preallocated(limit: RateLimit) -> RateLimiter {
39        RateLimiter { limit, readings: RingVec::new_preallocated(limit.count) }
40    }
41
42    /// Checks whether we're able to perform the event at this
43    /// time. On success, logs the current time to count towards
44    /// future enforcement.
45    pub fn check(&mut self) -> bool { self.check_at(Instant::now()) }
46
47    /// Like [`RateLimiter::check`], but you can provide an arbitrary
48    /// timestamp (useful for tests!).
49    ///
50    /// Warning: do not go backwards in time, things will mess up.
51    pub fn check_at(&mut self, instant: Instant) -> bool {
52        if self.readings.push(instant).is_ok() { return true; }
53        let reclaimed = self.sweep(instant);
54        if reclaimed {
55            self.readings.push(instant).unwrap();
56        }
57        reclaimed
58    }
59
60    /// Removes all readings before the period of our [`RateLimit`],
61    /// relative to the provided [`Instant`].
62    pub fn sweep(&mut self, instant: Instant) -> bool {
63        let bench = instant - self.limit.period;
64        let mut reclaimed = false;
65        while let Some(x) = self.readings.peek() {
66            if *x < bench {
67                reclaimed = true;
68                self.readings.pop();
69            } else { break; }
70        }
71        reclaimed
72    }
73}
simple_rate_limit/lib.rs

simple_rate_limit/
lib.rs
