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

tokio_rate_limit/algorithm/
mod.rs
