rate_guard/limits/

token_bucket.rs

//! Token bucket rate limiter implementation with integrated time source and Builder pattern.
//!
//! This module provides a token bucket rate limiter that allows bursts up to
//! capacity with sustained average rate. The implementation integrates time source
//! management directly and uses a Builder pattern for configuration.
//!
//! # Examples
//!
//! ```rust
//! use rate_guard::{Nanos, MockTimeSource, RateLimit};
//! use rate_guard::limits::TokenBucketBuilder;
//! use std::time::Duration;
//!
//! let bucket = TokenBucketBuilder::builder()
//!     .capacity(100)
//!     .refill_amount(10)
//!     .refill_every(Duration::from_millis(100))
//!     .with_time(MockTimeSource::new())
//!     .with_precision::<Nanos>()
//!     .build()
//!     .unwrap();
//!
//! assert!(bucket.try_acquire(50).is_ok());
//! ```

use std::time::Duration;
use std::marker::PhantomData;
use crate::limits::RateLimit;
use crate::{Uint, SimpleRateLimitResult, RateLimitResult, BuildResult};
use crate::precision::Precision;
use crate::time_source::TimeSource;
use crate::error::{RateLimitError, BuildError};
use rate_guard_core::cores::TokenBucketCore;
use rate_guard_core::SimpleRateLimitError;

/// Token bucket rate limiter with integrated time source and precision handling.
///
/// The token bucket algorithm allows bursts up to the bucket capacity while
/// maintaining a sustained average rate through periodic refilling. This
/// implementation integrates the time source directly, eliminating the need
/// for external time management.
///
/// # Type Parameters
/// * `P` - The precision type that determines how Duration values are converted to internal ticks
/// * `T` - The time source type that provides elapsed time since the rate limiter's reference point
pub struct TokenBucket<P: Precision, T: TimeSource> {
    /// The underlying core implementation that handles the token bucket logic.
    core: TokenBucketCore,
    /// The time source for getting current elapsed time.
    time_source: T,
    /// Phantom data to associate the Precision type with this instance.
    _precision: PhantomData<P>,
}

/// Builder for configuring and creating TokenBucket instances.
///
/// TokenBucketBuilder provides a fluent interface for configuring all aspects
/// of a token bucket rate limiter. The precision and time source types are
/// determined through the builder chain.
#[derive(Debug)]
pub struct TokenBucketBuilder {
    /// Maximum number of tokens the bucket can hold.
    capacity: Option<Uint>,
    /// Number of tokens added per refill operation.
    refill_amount: Option<Uint>,
    /// Time interval between refill operations.
    refill_every: Option<Duration>,
}

/// Builder with time source configured.
///
/// This intermediate builder is created after calling `with_time()` and allows
/// setting the precision type.
pub struct TokenBucketBuilderWithTime<T: TimeSource> {
    /// Maximum number of tokens the bucket can hold.
    capacity: Option<Uint>,
    /// Number of tokens added per refill operation.
    refill_amount: Option<Uint>,
    /// Time interval between refill operations.
    refill_every: Option<Duration>,
    /// Time source instance for elapsed time management.
    time_source: T,
}

/// Fully configured builder for TokenBucket with precision and time source set.
///
/// This builder is created after calling `with_precision()` and is ready to build
/// the final TokenBucket instance.
pub struct ConfiguredTokenBucketBuilder<P: Precision, T: TimeSource> {
    /// Maximum number of tokens the bucket can hold.
    capacity: Option<Uint>,
    /// Number of tokens added per refill operation.
    refill_amount: Option<Uint>,
    /// Time interval between refill operations.
    refill_every: Option<Duration>,
    /// Time source instance for elapsed time management.
    time_source: T,
    /// Phantom data to associate the Precision type with this instance.
    _precision: PhantomData<P>,
}

impl<P: Precision, T: TimeSource> std::fmt::Debug for TokenBucket<P, T>
where
    T: std::fmt::Debug,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TokenBucket")
            .field("time_source", &self.time_source)
            .field("_precision", &std::any::type_name::<P>())
            .finish_non_exhaustive()
    }
}

impl<T: TimeSource> std::fmt::Debug for TokenBucketBuilderWithTime<T>
where
    T: std::fmt::Debug,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TokenBucketBuilderWithTime")
            .field("capacity", &self.capacity)
            .field("refill_amount", &self.refill_amount)
            .field("refill_every", &self.refill_every)
            .field("time_source", &self.time_source)
            .finish()
    }
}

impl<P: Precision, T: TimeSource> std::fmt::Debug for ConfiguredTokenBucketBuilder<P, T>
where
    T: std::fmt::Debug,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ConfiguredTokenBucketBuilder")
            .field("capacity", &self.capacity)
            .field("refill_amount", &self.refill_amount)
            .field("refill_every", &self.refill_every)
            .field("time_source", &self.time_source)
            .field("_precision", &std::any::type_name::<P>())
            .finish()
    }
}

impl<P: Precision, T: TimeSource> RateLimit for TokenBucket<P, T> {
    #[inline(always)]
    fn try_acquire(&self, tokens: Uint) -> SimpleRateLimitResult {
        let elapsed = self.time_source.now();
        let current_tick = P::to_ticks(elapsed);
        self.core.try_acquire_at(current_tick, tokens)
    }

    #[inline(always)]
    fn try_acquire_verbose(&self, tokens: Uint) -> RateLimitResult {
        let elapsed = self.time_source.now();
        let current_tick = P::to_ticks(elapsed);
        self.core.try_acquire_verbose_at(current_tick, tokens)
            .map_err(|e| RateLimitError::from_core_error(e, |ticks| P::from_ticks(ticks)))
    }

    #[inline(always)]
    fn capacity_remaining(&self) -> Result<Uint, SimpleRateLimitError> {
        let elapsed = self.time_source.now();
        let current_tick = P::to_ticks(elapsed);
        self.core.capacity_remaining(current_tick)
    }
}

impl TokenBucketBuilder {
    /// Creates a new builder for configuring a TokenBucket.
    ///
    /// This is the preferred way to start building a TokenBucket.
    ///
    /// # Returns
    /// A new TokenBucketBuilder ready for configuration
    /// 
    /// # Examples
    ///
    /// ```rust
    /// use rate_guard::{Nanos, MockTimeSource, RateLimit};
    /// use rate_guard::limits::TokenBucketBuilder;
    /// use std::time::Duration;
    ///
    /// let bucket = TokenBucketBuilder::builder()
    ///     .capacity(100)
    ///     .refill_amount(10)
    ///     .refill_every(Duration::from_millis(100))
    ///     .with_time(MockTimeSource::new())
    ///     .with_precision::<Nanos>()
    ///     .build()
    ///     .unwrap();
    ///
    /// assert!(bucket.try_acquire(50).is_ok());
    /// ```
    pub fn builder() -> Self {
        Self {
            capacity: None,
            refill_amount: None,
            refill_every: None,
        }
    }

    /// Sets the maximum capacity of the token bucket.
    ///
    /// # Arguments
    /// * `capacity` - Maximum number of tokens the bucket can hold
    ///
    /// # Returns
    /// Self for method chaining
    pub fn capacity(mut self, capacity: Uint) -> Self {
        self.capacity = Some(capacity);
        self
    }

    /// Sets the number of tokens added per refill operation.
    ///
    /// # Arguments
    /// * `amount` - Number of tokens to add per refill operation
    ///
    /// # Returns
    /// Self for method chaining
    pub fn refill_amount(mut self, amount: Uint) -> Self {
        self.refill_amount = Some(amount);
        self
    }

    /// Sets the time interval between refill operations.
    ///
    /// # Arguments
    /// * `interval` - Duration between refill operations
    ///
    /// # Returns
    /// Self for method chaining
    pub fn refill_every(mut self, interval: Duration) -> Self {
        self.refill_every = Some(interval);
        self
    }

    /// Configures the time source for the token bucket.
    ///
    /// This method sets the time source and returns a builder that can
    /// then configure the precision type.
    ///
    /// # Arguments
    /// * `time_source` - The time source instance to use
    ///
    /// # Returns
    /// A builder with time source configured
    pub fn with_time<T: TimeSource>(self, time_source: T) -> TokenBucketBuilderWithTime<T> {
        TokenBucketBuilderWithTime {
            capacity: self.capacity,
            refill_amount: self.refill_amount,
            refill_every: self.refill_every,
            time_source,
        }
    }
}

impl Default for TokenBucketBuilder {
    fn default() -> Self {
        Self::builder()
    }
}

impl<T: TimeSource> TokenBucketBuilderWithTime<T> {
    /// Sets the precision type for time calculations.
    ///
    /// This method configures the precision type and returns a fully
    /// configured builder ready to build the TokenBucket.
    ///
    /// # Type Parameters
    /// * `P` - The precision type to use for time calculations
    ///
    /// # Returns
    /// A fully configured builder ready to build the TokenBucket
    pub fn with_precision<P: Precision>(self) -> ConfiguredTokenBucketBuilder<P, T> {
        ConfiguredTokenBucketBuilder {
            capacity: self.capacity,
            refill_amount: self.refill_amount,
            refill_every: self.refill_every,
            time_source: self.time_source,
            _precision: PhantomData,
        }
    }
}

impl<P: Precision, T: TimeSource> ConfiguredTokenBucketBuilder<P, T> {
    /// Builds the TokenBucket with the configured parameters.
    ///
    /// All required configuration must be provided before calling this method.
    /// This method also validates that all parameter values are valid.
    ///
    /// # Returns
    /// * `Ok(TokenBucket)` - Successfully created token bucket
    /// * `Err(BuildError)` - Missing required configuration or invalid parameter values
    ///
    /// # Errors
    ///
    /// Returns `BuildError` if:
    /// - Any required configuration is missing
    /// - Any parameter has an invalid value (e.g., capacity = 0)
    /// - Parameter combination is invalid (e.g., refill_amount > capacity)
    pub fn build(self) -> BuildResult<TokenBucket<P, T>> {
        let capacity = self.capacity.ok_or(BuildError::MissingArgument("capacity"))?;
        let refill_amount = self.refill_amount.ok_or(BuildError::MissingArgument("refill_amount"))?;
        let refill_every = self.refill_every.ok_or(BuildError::MissingArgument("refill_every"))?;

        // Validate individual parameter values
        if capacity == 0 {
            return Err(BuildError::InvalidArgument {
                field: "capacity",
                reason: "must be greater than 0"
            });
        }
        if refill_amount == 0 {
            return Err(BuildError::InvalidArgument {
                field: "refill_amount",
                reason: "must be greater than 0"
            });
        }
        if refill_every == Duration::ZERO {
            return Err(BuildError::InvalidArgument {
                field: "refill_every",
                reason: "must be greater than zero"
            });
        }

        // Validate parameter combinations
        if refill_amount > capacity {
            return Err(BuildError::InvalidArgument {
                field: "refill_amount",
                reason: "should not exceed capacity for optimal rate limiting behavior"
            });
        }

        let refill_every_ticks = P::to_ticks(refill_every);
        let core = TokenBucketCore::new(capacity, refill_every_ticks, refill_amount);

        Ok(TokenBucket {
            core,
            time_source: self.time_source,
            _precision: PhantomData,
        })
    }
}