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              :     dropped: u64,
       9              : }
      10              : 
      11              : pub struct RateLimitStats(u64);
      12              : 
      13              : impl std::fmt::Display for RateLimitStats {
      14            2 :     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
      15            2 :         write!(f, "{} dropped calls", self.0)
      16            2 :     }
      17              : }
      18              : 
      19              : impl RateLimit {
      20           39 :     pub fn new(interval: Duration) -> Self {
      21           39 :         Self {
      22           39 :             last: None,
      23           39 :             interval,
      24           39 :             dropped: 0,
      25           39 :         }
      26           39 :     }
      27              : 
      28              :     /// Call `f` if the rate limit allows.
      29              :     /// Don't call it otherwise.
      30           10 :     pub fn call<F: FnOnce()>(&mut self, f: F) {
      31           10 :         self.call2(|_| f())
      32           10 :     }
      33              : 
      34           18 :     pub fn call2<F: FnOnce(RateLimitStats)>(&mut self, f: F) {
      35           18 :         let now = Instant::now();
      36           13 :         match self.last {
      37           13 :             Some(last) if now - last <= self.interval => {
      38           11 :                 // ratelimit
      39           11 :                 self.dropped += 1;
      40           11 :             }
      41            7 :             _ => {
      42            7 :                 self.last = Some(now);
      43            7 :                 f(RateLimitStats(self.dropped));
      44            7 :                 self.dropped = 0;
      45            7 :             }
      46              :         }
      47           18 :     }
      48              : }
      49              : 
      50              : #[cfg(test)]
      51              : mod tests {
      52              :     use std::sync::atomic::AtomicUsize;
      53              : 
      54              :     #[test]
      55            1 :     fn basics() {
      56              :         use super::RateLimit;
      57              :         use std::sync::atomic::Ordering::Relaxed;
      58              :         use std::time::Duration;
      59              : 
      60            1 :         let called = AtomicUsize::new(0);
      61            1 :         let mut f = RateLimit::new(Duration::from_millis(100));
      62            1 : 
      63            3 :         let cl = || {
      64            3 :             called.fetch_add(1, Relaxed);
      65            3 :         };
      66              : 
      67            1 :         f.call(cl);
      68            1 :         assert_eq!(called.load(Relaxed), 1);
      69            1 :         f.call(cl);
      70            1 :         assert_eq!(called.load(Relaxed), 1);
      71            1 :         f.call(cl);
      72            1 :         assert_eq!(called.load(Relaxed), 1);
      73            1 :         std::thread::sleep(Duration::from_millis(100));
      74            1 :         f.call(cl);
      75            1 :         assert_eq!(called.load(Relaxed), 2);
      76            1 :         f.call(cl);
      77            1 :         assert_eq!(called.load(Relaxed), 2);
      78            1 :         std::thread::sleep(Duration::from_millis(100));
      79            1 :         f.call(cl);
      80            1 :         assert_eq!(called.load(Relaxed), 3);
      81            1 :     }
      82              : }