relay_server/services/buffer/
throttle.rs

1use std::num::NonZeroU32;
2use std::sync::{Mutex, PoisonError};
3use std::time::Duration;
4
5use tokio::time::Instant;
6
7/// Paces how quickly something happens.
8#[derive(Debug)]
9pub struct Throttle {
10    state: Mutex<ThrottleState>,
11}
12
13#[derive(Debug)]
14struct ThrottleState {
15    rate: Option<f64>,
16    budget: f64,
17    last_refill: Instant,
18}
19
20impl Throttle {
21    /// Creates a new [`Throttle`] with the given rate.
22    ///
23    /// `rate` is per second. A rate of `None` disables the throttle.
24    pub fn new(rate: Option<NonZeroU32>) -> Self {
25        let rate = rate.map(|rate| f64::from(rate.get()));
26        Self {
27            state: Mutex::new(ThrottleState {
28                rate,
29                budget: rate.unwrap_or_default(),
30                last_refill: Instant::now(),
31            }),
32        }
33    }
34
35    /// Updates the rate of the throttle.
36    ///
37    /// A rate of `None` disables the throttle.
38    pub fn set_rate(&self, rate: Option<NonZeroU32>) {
39        let rate = rate.map(|rate| f64::from(rate.get()));
40
41        let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
42        if state.rate != rate {
43            state.rate = rate;
44            state.budget = rate.unwrap_or_default();
45            state.last_refill = Instant::now();
46        }
47    }
48
49    /// Accounts units against the throttle and waits until the rate is back within
50    /// the configured limit.
51    ///
52    /// The units are deducted immediately and any debt is paid off by waiting.
53    pub async fn acquire(&self, units: usize) {
54        let wait = {
55            let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
56
57            match state.rate {
58                None => None,
59                Some(rate) => {
60                    let now = Instant::now();
61                    let elapsed = now.duration_since(state.last_refill).as_secs_f64();
62                    state.last_refill = now;
63
64                    state.budget += elapsed * rate;
65                    state.budget = state.budget.min(rate);
66                    state.budget -= units as f64;
67
68                    (state.budget < 0.0).then(|| Duration::from_secs_f64(-state.budget / rate))
69                }
70            }
71        };
72
73        if let Some(wait) = wait {
74            tokio::time::sleep(wait).await;
75        }
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82
83    fn throttle(rate: u32) -> Throttle {
84        Throttle::new(NonZeroU32::new(rate))
85    }
86
87    #[tokio::test(start_paused = true)]
88    async fn test_within_rate_does_not_wait() {
89        let throttle = throttle(100);
90
91        let start = Instant::now();
92        throttle.acquire(100).await;
93
94        assert_eq!(start.elapsed(), Duration::ZERO);
95    }
96
97    #[tokio::test(start_paused = true)]
98    async fn test_debt_is_paid_off_by_waiting() {
99        let throttle = throttle(100);
100
101        let start = Instant::now();
102        throttle.acquire(100).await;
103        throttle.acquire(50).await;
104
105        assert_eq!(start.elapsed(), Duration::from_millis(500));
106    }
107
108    #[tokio::test(start_paused = true)]
109    async fn test_large_batch_is_delayed_not_starved() {
110        let throttle = throttle(10);
111
112        let start = Instant::now();
113        throttle.acquire(30).await;
114
115        assert_eq!(start.elapsed(), Duration::from_secs(2));
116    }
117
118    #[tokio::test(start_paused = true)]
119    async fn test_budget_units_refill_over_time() {
120        let throttle = throttle(100);
121        throttle.acquire(100).await;
122        tokio::time::advance(Duration::from_secs(1)).await;
123
124        let start = Instant::now();
125        throttle.acquire(100).await;
126
127        assert_eq!(start.elapsed(), Duration::ZERO);
128    }
129
130    #[tokio::test(start_paused = true)]
131    async fn test_refill_is_capped_at_rate() {
132        let throttle = throttle(100);
133        tokio::time::advance(Duration::from_secs(60)).await;
134
135        let start = Instant::now();
136        throttle.acquire(200).await;
137
138        assert_eq!(start.elapsed(), Duration::from_secs(1));
139    }
140
141    #[tokio::test(start_paused = true)]
142    async fn test_disabled_does_not_wait() {
143        let throttle = Throttle::new(None);
144
145        let start = Instant::now();
146        throttle.acquire(1_000_000).await;
147
148        assert_eq!(start.elapsed(), Duration::ZERO);
149    }
150
151    #[tokio::test(start_paused = true)]
152    async fn test_set_rate_enables_throttling() {
153        let throttle = Throttle::new(None);
154        throttle.set_rate(NonZeroU32::new(100));
155
156        let start = Instant::now();
157        throttle.acquire(100).await;
158        throttle.acquire(50).await;
159
160        assert_eq!(start.elapsed(), Duration::from_millis(500));
161    }
162
163    #[tokio::test(start_paused = true)]
164    async fn test_set_rate_disables_throttling() {
165        let throttle = throttle(10);
166        throttle.acquire(30).await;
167        throttle.set_rate(None);
168
169        let start = Instant::now();
170        throttle.acquire(1_000_000).await;
171
172        assert_eq!(start.elapsed(), Duration::ZERO);
173    }
174}
relay_server/services/buffer/throttle.rs

relay_server/services/buffer/
throttle.rs
