proxmox-http/src/client/rate_limiter.rs

   1 use std::time::{Duration, Instant};
   2 use std::convert::TryInto;
   3
   4 /// Token bucket based rate limiter
   5 pub struct RateLimiter {
   6     rate: u64, // tokens/second
   7     start_time: Instant,
   8     traffic: u64, // overall traffic
   9     bucket_size: u64,
  10     last_update: Instant,
  11     consumed_tokens: u64,
  12 }
  13
  14 impl RateLimiter {
  15
  16     const NO_DELAY: Duration = Duration::from_millis(0);
  17
  18     /// Creates a new instance, using [Instant::now] as start time.
  19     pub fn new(rate: u64, bucket_size: u64) -> Self {
  20         let start_time = Instant::now();
  21         Self::with_start_time(rate, bucket_size, start_time)
  22     }
  23
  24     /// Creates a new instance with specified `rate`, `bucket_size` and `start_time`.
  25     pub fn with_start_time(rate: u64, bucket_size: u64, start_time: Instant) -> Self {
  26         Self {
  27             rate,
  28             start_time,
  29             traffic: 0,
  30             bucket_size,
  31             last_update: start_time,
  32             // start with empty bucket (all tokens consumed)
  33             consumed_tokens: bucket_size,
  34         }
  35     }
  36
  37     /// Returns the average rate (since `start_time`)
  38     pub fn average_rate(&self, current_time: Instant) -> f64 {
  39         let time_diff = current_time.saturating_duration_since(self.start_time).as_secs_f64();
  40         if time_diff <= 0.0 {
  41             0.0
  42         } else {
  43             (self.traffic as f64) / time_diff
  44         }
  45     }
  46
  47     fn refill_bucket(&mut self, current_time: Instant) {
  48         let time_diff = match current_time.checked_duration_since(self.last_update) {
  49             Some(duration) => duration.as_nanos(),
  50             None => {
  51                 //log::error!("update_time: got negative time diff");
  52                 return;
  53             }
  54         };
  55
  56         if time_diff == 0 { return; }
  57
  58         self.last_update = current_time;
  59
  60         let allowed_traffic = ((time_diff.saturating_mul(self.rate as u128)) / 1_000_000_000)
  61             .try_into().unwrap_or(u64::MAX);
  62
  63         self.consumed_tokens = self.consumed_tokens.saturating_sub(allowed_traffic);
  64     }
  65
  66     /// Register traffic, returning a proposed delay to reach the expected rate.
  67     pub fn register_traffic(&mut self, current_time: Instant, data_len: u64) -> Duration {
  68         self.refill_bucket(current_time);
  69
  70         self.traffic += data_len;
  71         self.consumed_tokens += data_len;
  72
  73         if self.consumed_tokens <= self.bucket_size {
  74             return Self::NO_DELAY;
  75         }
  76         Duration::from_nanos((self.consumed_tokens - self.bucket_size).saturating_mul(1_000_000_000)/ self.rate)
  77     }
  78 }