backblaze_b2_client/
throttle.rs

1use std::{
2    ops::AddAssign,
3    time::{Duration, Instant},
4};
5
6use num::Unsigned;
7use tokio::time::sleep;
8
9#[derive(Debug)]
10pub struct Throttle<T: Unsigned + AddAssign + Copy + PartialOrd> {
11    max_per_period: T,
12    count_start: Instant,
13    period: Duration,
14    current_count: T,
15}
16
17impl<T: Unsigned + AddAssign + Copy + PartialOrd> Throttle<T> {
18    pub fn new(max_per_period: T, period: Duration) -> Self {
19        Self {
20            max_per_period,
21            period,
22            count_start: Instant::now(),
23            current_count: T::zero(),
24        }
25    }
26
27    /// Equivalent to
28    /// ```rust
29    /// Throttle::new(max_per_period, Duration::from_secs(1))
30    /// ```
31    pub fn per_second(max_per_period: T) -> Self {
32        Self::new(max_per_period, Duration::from_secs(1))
33    }
34
35    /// Equivalent to
36    /// ```rust
37    /// Throttle::new(max_per_period, Duration::from_secs(60))
38    /// ```
39    pub fn per_minute(max_per_period: T) -> Self {
40        Self::new(max_per_period, Duration::from_secs(60))
41    }
42
43    /// Advances the throttle by 1, waiting if the throttle has been exhausted
44    pub async fn advance(&mut self) -> T {
45        self.advance_by(T::one()).await
46    }
47
48    /// Advances the throttle by the given amount, waiting if the throttle has been exhausted
49    pub async fn advance_by(&mut self, by: T) -> T {
50        if self.count_start.elapsed() >= self.period {
51            self.current_count = T::zero();
52            self.count_start = Instant::now();
53        }
54
55        if self.current_count >= self.max_per_period {
56            sleep(self.period - self.count_start.elapsed()).await;
57            self.current_count = T::zero();
58            self.count_start = Instant::now();
59        }
60
61        self.current_count += by;
62
63        return if self.current_count > self.max_per_period {
64            T::zero()
65        } else {
66            self.max_per_period - self.current_count
67        };
68    }
69
70    /// If throttle period has been exhausted, waits for the period to end <br>
71    /// otherwise returns immediately
72    pub async fn wait_if_exhausted(&self) {
73        if self.count_start.elapsed() >= self.period {
74            return;
75        }
76
77        if self.current_count >= self.max_per_period {
78            sleep(self.period - self.count_start.elapsed()).await;
79        }
80    }
81
82    /// Returns the remaining count for the current period
83    pub fn remaining(&self) -> T {
84        if self.count_start.elapsed() >= self.period {
85            return self.max_per_period;
86        }
87
88        return if self.current_count > self.max_per_period {
89            T::zero()
90        } else {
91            self.max_per_period - self.current_count
92        };
93    }
94}
95
96impl<T: Unsigned + AddAssign + Copy + PartialOrd> Clone for Throttle<T> {
97    fn clone(&self) -> Self {
98        Self {
99            max_per_period: self.max_per_period,
100            period: self.period.clone(),
101            count_start: Instant::now(),
102            current_count: T::zero(),
103        }
104    }
105}
backblaze_b2_client/throttle.rs

backblaze_b2_client/
throttle.rs
