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