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