touch_ratelimit 0.1.0

A composable, extensible rate limiting crate for Rust
Documentation
use std::time::Instant;

use crate::bucket::RateLimiter;

/// A token bucket rate limiting algorithm.
///
/// The token bucket algorithm allows requests to proceed as long as
/// sufficient tokens are available. Tokens are replenished over time
/// at a fixed rate, up to a maximum capacity.
///
/// This implementation:
///
/// - Supports fractional tokens
/// - Uses a monotonic clock (`Instant`)
/// - Is intended for **single-identity usage**
///
/// ## Typical usage
///
/// A `TokenBucket` is usually not used directly. Instead, it is created
/// by a store (e.g. `InMemoryStore`) for each unique request key.
#[derive(Clone)]
pub struct TokenBucket {
    /// Maximum number of tokens in the bucket
    capacity: f64,

    /// Current number of tokens
    tokens: f64,

    /// Tokens added per second
    refill_rate: f64,

    /// Last time the bucket was refilled
    last_refill: Instant,
}

impl TokenBucket {
    /// Create a new token bucket.
    ///
    /// # Arguments
    ///
    /// - `capacity`: Maximum number of tokens that can be accumulated.
    /// - `refill_rate`: Number of tokens added per second.
    ///
    /// The bucket starts full.
    pub fn new(capacity: f64, refill_rate: f64) -> Self {
        Self {
            capacity,
            refill_rate,
            tokens: capacity,
            last_refill: Instant::now(),
        }
    }

    /// Return the number of currently available tokens.
    ///
    /// This method updates the internal state by performing a refill
    /// before returning the token count.
    pub fn available_tokens(&mut self) -> f64 {
        self.refill();
        self.tokens
    }

    /// Attempt to consume `n` tokens.
    ///
    /// Returns `true` if enough tokens were available, or `false` otherwise.
    pub fn allow_n(&mut self, n: f64) -> bool {
        self.refill();

        if self.tokens >= n {
            self.tokens -= n;
            true
        } else {
            false
        }
    }

    /// Attempt to consume a single token.
    pub fn allow(&mut self) -> bool {
        self.allow_n(1.0)
    }

    /// Refill the bucket based on elapsed time.
    fn refill(&mut self) {
        let now = Instant::now();
        let elapsed = now.duration_since(self.last_refill).as_secs_f64();

        let added = elapsed * self.refill_rate;
        self.tokens = (self.tokens + added).min(self.capacity);
        self.last_refill = now;
    }
}

impl RateLimiter for TokenBucket {
    fn allow(&mut self) -> bool {
        TokenBucket::allow(self)
    }

    fn allow_n(&mut self, n: f64) -> bool {
        TokenBucket::allow_n(self, n)
    }

    fn available_tokens(&mut self) -> f64 {
        TokenBucket::available_tokens(self)
    }
}