Skip to main content

tokio_rate_limit/algorithm/
mod.rs

1//! Rate limiting algorithms.
2
3use crate::error::Result;
4use crate::limiter::RateLimitDecision;
5use async_trait::async_trait;
6
7mod leaky_bucket;
8mod token_bucket;
9
10// v0.6.0 experimental optimizations
11mod cached_token_bucket;
12#[allow(deprecated)]
13mod simd_token_bucket;
14#[allow(deprecated)]
15mod zerocopy_token_bucket;
16
17// v0.7.0 probabilistic rate limiting
18mod probabilistic_token_bucket;
19
20pub use leaky_bucket::LeakyBucket;
21pub use token_bucket::TokenBucket;
22
23// Experimental exports (v0.6.0)
24pub use cached_token_bucket::CachedTokenBucket;
25#[allow(deprecated)]
26pub use simd_token_bucket::SimdTokenBucket;
27#[allow(deprecated)]
28pub use zerocopy_token_bucket::ZeroCopyTokenBucket;
29
30// Probabilistic exports (v0.7.0)
31pub use probabilistic_token_bucket::ProbabilisticTokenBucket;
32
33/// Private module for the sealed trait pattern.
34///
35/// This prevents external implementations of the Algorithm trait while maintaining
36/// flexibility for internal algorithm implementations. This allows us to make
37/// breaking changes to the trait in the future without requiring a semver major bump.
38mod private {
39    pub trait Sealed {}
40}
41
42/// Trait for rate limiting algorithms.
43///
44/// Implementations of this trait define how rate limiting decisions are made.
45/// The trait is async to allow for potential I/O operations in custom implementations.
46///
47/// # Sealed Trait
48///
49/// This trait is sealed and cannot be implemented outside of this crate. This design
50/// allows us to add new methods or change the trait in minor version updates without
51/// breaking semver guarantees. If you need a custom algorithm, please open an issue
52/// to discuss adding it to the library.
53///
54/// # Available Algorithms
55///
56/// - [`TokenBucket`] - Allows bursts up to capacity, refills at constant rate (zero-copy optimized in v0.4.0)
57/// - [`LeakyBucket`] - Enforces steady rate, smooths traffic (v0.3.0)
58/// - [`CachedTokenBucket`] - Thread-local cached token bucket for hot-key workloads (v0.4.0)
59///
60/// # Experimental Algorithms (Not Recommended for Production)
61///
62/// - [`ZeroCopyTokenBucket`] - Zero-copy prototype (integrated into TokenBucket in v0.4.0)
63/// - [`SimdTokenBucket`] - SIMD prototype (deferred, no performance benefit)
64///
65/// # Future Algorithms
66///
67/// - Sliding Window - More precise rate limiting
68#[async_trait]
69pub trait Algorithm: Send + Sync + private::Sealed {
70    /// Checks if a request for the given key should be permitted.
71    ///
72    /// # Arguments
73    ///
74    /// * `key` - A string identifier for the client/resource being rate limited
75    ///
76    /// # Returns
77    ///
78    /// A `RateLimitDecision` indicating whether the request is permitted and
79    /// additional metadata about the rate limit status.
80    async fn check(&self, key: &str) -> Result<RateLimitDecision>;
81
82    /// Checks if a request with the given cost should be permitted.
83    ///
84    /// Cost represents the number of tokens to consume. The default implementation
85    /// uses a cost of 1 (equivalent to `check()`), but algorithms can override this
86    /// to support weighted rate limiting.
87    ///
88    /// # Arguments
89    ///
90    /// * `key` - A string identifier for the client/resource being rate limited
91    /// * `cost` - Number of tokens to consume (must be > 0)
92    ///
93    /// # Returns
94    ///
95    /// A `RateLimitDecision` indicating whether the request is permitted and
96    /// additional metadata about the rate limit status.
97    ///
98    /// # Default Behavior
99    ///
100    /// The default implementation delegates to `check()` for cost == 1.
101    /// For cost > 1, it falls back to `check()` as well — concrete algorithms
102    /// should override this method for proper weighted rate limiting support.
103    async fn check_with_cost(&self, key: &str, _cost: u64) -> Result<RateLimitDecision> {
104        // Default: delegate to check() for all costs.
105        // Concrete algorithms override this with proper cost-aware consumption.
106        self.check(key).await
107    }
108}