Repository: tokio-rs/tokio
File: tokio/src/sync/semaphore.rs (Rust)

Code before:
/// ```
///
/// Implement a simple token bucket for rate limiting
///
/// Many applications and systems have constraints on the rate at which certain
/// operations should occur. Exceeding this rate can result in suboptimal
/// performance or even errors.
///
/// The provided example uses a `TokenBucket`, implemented using a semaphore, that
/// limits operations to a specific rate although token buckets allow short bursts that are faster
/// than the allowed rate. The token bucket will be refilled after each interval. When the rate is exceeded,
/// the `acquire` method will await until a token is available. Note that this implementation is suboptimal
/// when the rate is large, because it consumes a lot of cpu constantly looping and sleeping.
/// ```
/// use std::sync::Arc;
/// use tokio::sync::{AcquireError, Semaphore};
/// use tokio::time::{interval, Duration};
///
/// struct TokenBucket {
///     sem: Arc<Semaphore>,
///     jh: tokio::task::JoinHandle<()>,
/// }
///
/// impl TokenBucket {
///     fn new(duration: Duration) -> Self {
///         let rate = (1.0 / duration.as_secs_f32()) as usize;
///         let sem = Arc::new(Semaphore::new(rate));
///
///         // refills the tokens at the end of each interval
///         let jh = tokio::spawn({
///             let sem = sem.clone();
///             let mut interval = interval(duration);
///             interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
///
///             async move {
///                 loop {
///                     interval.tick().await;
///
///                     let cap = rate - sem.available_permits();
///                     sem.add_permits(cap.min(1));
///                 }
///             }
///         });
///
///         Self { jh, sem }
///     }
///
///     async fn acquire(&self) -> Result<(), AcquireError> {
///         self.sem.acquire().await.map(|p| p.forget())
///     }
///

Review comment (suggestion):
You're only using this to compute the maximum capacity. Instead, I suggest making the capacity into a separate argument.

Code after:
/// use tokio::sync::{AcquireError, Semaphore};
/// use tokio::time::{interval, Duration};
///
/// struct TokenBucket {
///     sem: Arc<Semaphore>,
///     jh: tokio::task::JoinHandle<()>,
/// }
///
/// impl TokenBucket {
///     fn new(duration: Duration, capacity: usize) -> Self {
///         let sem = Arc::new(Semaphore::new(capacity));
///
///         // refills the tokens at the end of each interval
///         let jh = tokio::spawn({
///             let sem = sem.clone();
///             let mut interval = interval(duration);
///             interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
///
///             async move {
///                 loop {
///                     interval.tick().await;
///
///                     if sem.available_permits() < capacity {
///                         sem.add_permits(1);
///                     }
///                 }
///             }
///         });
///
///         Self { jh, sem }
///     }
///
///     async fn acquire(&self) -> Result<(), AcquireError> {
///         self.sem.acquire().await.map(|p| p.forget())
///     }
///
///     async fn close(self) {
///         self.sem.close();
///         self.jh.abort();
///         let _ = self.jh.await;
///     }
/// }
///
/// #[tokio::main]
/// async fn main() {
///     let capacity = 5; // operation per second
///     let update_interval = Duration::from_secs_f32(1.0 / capacity as f32);
///     let bucket = TokenBucket::new(update_interval, capacity);
///
///     for _ in 0..5 {
///         bucket.acquire().await.unwrap();