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

Code before:
///     buffer.write_all(message).await?;
///     Ok(()) // Permit goes out of scope here, and is available again for acquisition
/// }
/// ```
///
/// 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. The token bucket will be refilled gradually.
/// When the rate is exceeded, the `acquire` method will await until a token is available.
/// ```
/// use std::sync::Arc;
/// use tokio::sync::{AcquireError, Semaphore, SemaphorePermit};
/// use tokio::time::{sleep, Duration};
///
/// struct TokenBucket {
///     sem: Arc<Semaphore>,
///     jh: tokio::task::JoinHandle<()>,
/// }
///
/// impl TokenBucket {
///     fn new(rate: usize) -> Self {
///         let sem = Arc::new(Semaphore::new(rate));
///
///         // refills the permits each 1/rate seconds.
///         let jh = tokio::spawn({
///             let sem = sem.clone();
///
///             async move {
///                 let time_slice = 1.0 / (rate as f32);
///
///                 loop {
///                     sleep(Duration::from_secs_f32(time_slice)).await;
///
///                     let cap = rate - sem.available_permits();
///                     sem.add_permits(std::cmp::min(cap, 1));
///                 }
///             }
///         });
///
///         Self { jh, sem }
///     }
///
///     async fn acquire(&self) -> Result<SemaphorePermit<'_>, AcquireError> {
///         self.sem.acquire().await
///     }
///

Review comment (suggestion):
Perhaps this should just take a Duration of the amount of time there should be between each permit? That would also allow for permits being added more rarely than once per second.

Code after:
/// [token bucket]: https://en.wikipedia.org/wiki/Token_bucket
/// ```
/// 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, 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);