Line data    Source code

       1              : //! A helper to rate limit operations.
       2              : 
       3              : use std::time::{Duration, Instant};
       4              : 
       5              : pub struct RateLimit {
       6              :     last: Option<Instant>,
       7              :     interval: Duration,
       8              : }
       9              : 
      10              : impl RateLimit {
      11            2 :     pub fn new(interval: Duration) -> Self {
      12            2 :         Self {
      13            2 :             last: None,
      14            2 :             interval,
      15            2 :         }
      16            2 :     }
      17              : 
      18              :     /// Call `f` if the rate limit allows.
      19              :     /// Don't call it otherwise.
      20           12 :     pub fn call<F: FnOnce()>(&mut self, f: F) {
      21           12 :         let now = Instant::now();
      22           10 :         match self.last {
      23           10 :             Some(last) if now - last <= self.interval => {
      24            6 :                 // ratelimit
      25            6 :             }
      26            6 :             _ => {
      27            6 :                 self.last = Some(now);
      28            6 :                 f();
      29            6 :             }
      30              :         }
      31           12 :     }
      32              : }
      33              : 
      34              : #[cfg(test)]
      35              : mod tests {
      36              :     use std::sync::atomic::AtomicUsize;
      37              : 
      38              :     #[test]
      39            2 :     fn basics() {
      40            2 :         use super::RateLimit;
      41            2 :         use std::sync::atomic::Ordering::Relaxed;
      42            2 :         use std::time::Duration;
      43            2 : 
      44            2 :         let called = AtomicUsize::new(0);
      45            2 :         let mut f = RateLimit::new(Duration::from_millis(100));
      46            2 : 
      47            6 :         let cl = || {
      48            6 :             called.fetch_add(1, Relaxed);
      49            6 :         };
      50              : 
      51            2 :         f.call(cl);
      52            2 :         assert_eq!(called.load(Relaxed), 1);
      53            2 :         f.call(cl);
      54            2 :         assert_eq!(called.load(Relaxed), 1);
      55            2 :         f.call(cl);
      56            2 :         assert_eq!(called.load(Relaxed), 1);
      57            2 :         std::thread::sleep(Duration::from_millis(100));
      58            2 :         f.call(cl);
      59            2 :         assert_eq!(called.load(Relaxed), 2);
      60            2 :         f.call(cl);
      61            2 :         assert_eq!(called.load(Relaxed), 2);
      62            2 :         std::thread::sleep(Duration::from_millis(100));
      63            2 :         f.call(cl);
      64            2 :         assert_eq!(called.load(Relaxed), 3);
      65            2 :     }
      66              : }