ratelock 0.1.0

Zero-dependency, zero-allocation, lock-free token bucket rate limiter for Rust (std + no_std).
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133

use core::array::from_fn;
use core::hash::{Hash, Hasher};

#[cfg(feature = "std")]
use crate::StdClock;
use crate::{Clock, RateLimiter, Snapshot};

/// A fixed-size shard set of independent [`RateLimiter`] instances.
///
/// Sharding is a recommended pattern for high-contention workloads
/// (for example per-key or per-tenant limits).
///
/// `N` is the number of shards. `N=0` is allowed and treated as disabled;
/// all `allow*` calls return `false`, and `remaining*` returns `0`.
pub struct ShardedRateLimiter<C: Clock + Clone, const N: usize> {
    shards: [RateLimiter<C>; N],
}

impl<C: Clock + Clone, const N: usize> ShardedRateLimiter<C, N> {
    /// Creates a sharded limiter with a custom clock.
    ///
    /// Each shard has identical `(capacity, refill_per_sec)` configuration.
    pub fn with_clock(capacity: u64, refill_per_sec: u64, clock: C) -> Self {
        Self {
            shards: from_fn(|_| RateLimiter::with_clock(capacity, refill_per_sec, clock.clone())),
        }
    }

    /// Returns the number of shards.
    #[inline]
    pub const fn shard_count(&self) -> usize {
        N
    }

    /// Tries to take one token from the shard selected by `hash`.
    #[inline]
    pub fn allow_by_hash(&self, hash: u64) -> bool {
        self.shard_for_hash(hash).is_some_and(RateLimiter::allow)
    }

    /// Tries to take `n` tokens atomically from the shard selected by `hash`.
    #[inline]
    pub fn allow_n_by_hash(&self, hash: u64, n: u64) -> bool {
        self.shard_for_hash(hash)
            .is_some_and(|limiter| limiter.allow_n(n))
    }

    /// Returns remaining tokens from the shard selected by `hash`.
    #[inline]
    pub fn remaining_by_hash(&self, hash: u64) -> u64 {
        self.shard_for_hash(hash).map_or(0, RateLimiter::remaining)
    }

    /// Returns a snapshot for the shard selected by `hash`.
    #[inline]
    pub fn snapshot_by_hash(&self, hash: u64) -> Option<Snapshot> {
        self.shard_for_hash(hash).map(RateLimiter::snapshot)
    }

    /// Returns a shard by index.
    #[inline]
    pub fn shard(&self, index: usize) -> Option<&RateLimiter<C>> {
        self.shards.get(index)
    }

    /// Returns a shard selected by a precomputed hash.
    #[inline]
    pub fn shard_for_hash(&self, hash: u64) -> Option<&RateLimiter<C>> {
        self.shard(self.index_for_hash(hash)?)
    }

    #[inline]
    fn index_for_hash(&self, hash: u64) -> Option<usize> {
        if N == 0 {
            return None;
        }

        let idx = if N.is_power_of_two() {
            (hash as usize) & (N - 1)
        } else {
            (hash as usize) % N
        };
        Some(idx)
    }

    /// Hashes a key to `u64` using a caller-provided hasher implementation.
    #[inline]
    pub fn hash_key_with<H: Hasher + Default, K: Hash>(&self, key: &K) -> u64 {
        let mut hasher = H::default();
        key.hash(&mut hasher);
        hasher.finish()
    }
}

#[cfg(feature = "std")]
impl<const N: usize> ShardedRateLimiter<StdClock, N> {
    /// Creates a sharded limiter using [`StdClock`].
    #[inline]
    pub fn new(capacity: u64, refill_per_sec: u64) -> Self {
        Self::with_clock(capacity, refill_per_sec, StdClock)
    }

    /// Tries to take one token from a shard selected by a key hash.
    #[inline]
    pub fn allow_by_key<K: Hash>(&self, key: &K) -> bool {
        self.allow_by_hash(self.hash_key_with::<std::collections::hash_map::DefaultHasher, K>(key))
    }

    /// Tries to take `n` tokens atomically from a shard selected by a key hash.
    #[inline]
    pub fn allow_n_by_key<K: Hash>(&self, key: &K, n: u64) -> bool {
        self.allow_n_by_hash(
            self.hash_key_with::<std::collections::hash_map::DefaultHasher, K>(key),
            n,
        )
    }

    /// Returns remaining tokens from a shard selected by a key hash.
    #[inline]
    pub fn remaining_by_key<K: Hash>(&self, key: &K) -> u64 {
        self.remaining_by_hash(
            self.hash_key_with::<std::collections::hash_map::DefaultHasher, K>(key),
        )
    }

    /// Returns a snapshot from a shard selected by a key hash.
    #[inline]
    pub fn snapshot_by_key<K: Hash>(&self, key: &K) -> Option<Snapshot> {
        self.snapshot_by_hash(
            self.hash_key_with::<std::collections::hash_map::DefaultHasher, K>(key),
        )
    }
}