TLA  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 CBC           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                 : }