async_throttle/
single.rs

1use std::time::Duration;
2use tokio::sync::Mutex;
3use tokio::time::{interval, Interval};
4
5/// [`RateLimiter`] is a tool which can control the rate at which processing happens.
6///
7/// # Examples
8///
9/// ```
10/// use async_throttle::RateLimiter;
11///
12/// #[tokio::main]
13/// async fn main() {
14///    let period = std::time::Duration::from_millis(10);
15///    let rate_limiter = RateLimiter::new(period);
16///    
17///    // Takes 90ms to complete, the first iteration is instant, the next 9 iterations take 100ms
18///    for _ in 0..10 {
19///       rate_limiter.throttle(|| async { /* work */ }).await;
20///    }
21/// }
22pub struct RateLimiter {
23    /// The mutex that will be locked when the rate limiter is waiting for the interval to tick.
24    ///
25    /// It's important to use a tokio::sync::Mutex here instead of a std::sync::Mutex. The reason is
26    /// that the tokio::sync::Mutex does not block & the MutexGuard is held across await points.
27    ///
28    /// If you tried to use std::sync::Mutex instead, you would get a compiler error when
29    /// spawning tokio tasks because the MutexGuard would not be Send.
30    interval: Mutex<Interval>,
31}
32
33impl RateLimiter {
34    /// Creates a new rate limiter.
35    ///
36    /// # Examples
37    ///
38    /// ```
39    /// use tokio::sync::Mutex;
40    /// use anyhow::Result;
41    /// use std::time::Duration;
42    /// use async_throttle::RateLimiter;
43    ///
44    /// #[tokio::main]
45    /// async fn main() -> Result<()> {
46    ///     RateLimiter::new(Duration::from_millis(10));
47    ///     Ok(())
48    /// }
49    /// ```
50    pub fn new(period: Duration) -> Self {
51        Self {
52            interval: Mutex::new(interval(period)),
53        }
54    }
55
56    /// Throttles the execution of a function.
57    ///
58    /// # Examples
59    ///
60    /// ```
61    /// use async_throttle::RateLimiter;
62    /// use anyhow::Result;
63    /// use std::sync::Arc;
64    ///
65    /// async fn do_work() { /* some computation */ }
66    ///
67    /// async fn do_throttle(limiter: Arc<RateLimiter>) {
68    ///    limiter.throttle(|| do_work()).await
69    /// }
70    pub async fn throttle<Fut, F, T>(&self, f: F) -> T
71    where
72        Fut: std::future::Future<Output = T>,
73        F: FnOnce() -> Fut,
74    {
75        self.wait().await;
76        f().await
77    }
78
79    /// Waits for the interval to tick.
80    ///
81    /// This is the building block for the throttle function. It works by allowing only one task to
82    /// access the interval at a time. The first task to access the interval will tick it and then
83    /// release the lock. The next task to access the interval will tick it and then release the
84    /// lock.
85    async fn wait(&self) {
86        let mut interval = self.interval.lock().await;
87        interval.tick().await;
88    }
89}
async_throttle/single.rs

async_throttle/
single.rs
