tower_rate_limiter/store/
redis.rs

1//! Redis rate-limit store, enabled by the `redis` Cargo feature.
2
3use std::{future::Future, pin::Pin, time::Duration};
4
5use redis::aio::MultiplexedConnection;
6
7#[cfg(feature = "redis-lua")]
8use redis::Script;
9
10use crate::{RateLimitError, Store, Usage};
11
12const REDIS_PREFIX: &str = "rl:";
13
14#[cfg(feature = "redis-lua")]
15const INCREMENT_SCRIPT: &str = r#"
16local count = redis.call('INCR', KEYS[1])
17if count == 1 then
18    redis.call('PEXPIRE', KEYS[1], ARGV[1])
19end
20return {count, redis.call('PTTL', KEYS[1])}
21"#;
22
23/// Errors returned by the Redis rate-limit store.
24#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
25pub enum RedisStoreError {
26    /// The fixed window cannot be represented in Redis milliseconds.
27    #[error("window must be at least one millisecond")]
28    WindowTooShort,
29
30    /// The fixed window exceeds Redis's signed 64-bit millisecond range.
31    #[error("window is too large")]
32    WindowTooLarge,
33
34    /// Redis rejected or could not execute the atomic increment operation.
35    #[error("redis command failed: {0}")]
36    CommandFailed(String),
37
38    /// The atomic increment returned a non-positive usage count.
39    #[error("redis returned invalid usage {0}")]
40    InvalidUsage(i64),
41
42    /// The atomic increment returned a non-positive reset duration.
43    #[error("redis returned invalid reset-after milliseconds {0}")]
44    InvalidResetAfter(i64),
45}
46
47impl From<redis::RedisError> for RedisStoreError {
48    fn from(error: redis::RedisError) -> Self {
49        Self::CommandFailed(error.to_string())
50    }
51}
52
53/// Convert the RedisStoreError to a RateLimitError.
54impl From<RedisStoreError> for RateLimitError {
55    fn from(error: RedisStoreError) -> Self {
56        RateLimitError::Store("redis_store_error".into(), error.to_string())
57    }
58}
59
60/// Redis-backed implementation of the common fixed-window [`Store`] seam.
61///
62/// The connection is supplied by the caller and is cloned for each increment. A Redis
63/// `MultiplexedConnection` clone shares its underlying connection and does not transfer
64/// connection lifecycle ownership to this adapter. Window durations are truncated to whole
65/// milliseconds because Redis expiry commands use millisecond precision; values shorter than one
66/// millisecond are rejected.
67#[derive(Clone, Debug)]
68pub struct RedisStore {
69    connection: MultiplexedConnection,
70    namespace: Option<String>,
71}
72
73impl RedisStore {
74    /// Construct a store from an established Redis multiplexed connection.
75    pub fn new(connection: MultiplexedConnection) -> Self {
76        Self {
77            connection,
78            namespace: None,
79        }
80    }
81
82    /// Add an optional namespace. Empty namespaces are treated as absent.
83    pub fn with_namespace(mut self, namespace: impl Into<String>) -> Self {
84        self.namespace = Some(namespace.into());
85        self
86    }
87
88    /// Format the Redis transport key for this store's optional namespace.
89    fn redis_key(&self, key: &str) -> String {
90        format_redis_key(self.namespace.as_deref(), key)
91    }
92}
93
94impl Store for RedisStore {
95    type Future = Pin<Box<dyn Future<Output = Result<Usage, RateLimitError>> + Send>>;
96
97    fn increment(&self, key: &str, window: Duration) -> Self::Future {
98        let window_millis = match checked_window_millis(window) {
99            Ok(window_millis) => window_millis,
100            Err(error) => return Box::pin(std::future::ready(Err(error.into()))),
101        };
102        let redis_key = self.redis_key(key);
103        let mut connection = self.connection.clone();
104
105        Box::pin(async move {
106            let result = increment_counter(&mut connection, &redis_key, window_millis).await?;
107            usage_from_increment_result(result).map_err(Into::into)
108        })
109    }
110}
111
112#[cfg(feature = "redis-lua")]
113async fn increment_counter(
114    connection: &mut MultiplexedConnection,
115    redis_key: &str,
116    window_millis: i64,
117) -> Result<(i64, i64), RedisStoreError> {
118    let script = Script::new(INCREMENT_SCRIPT);
119    let mut invocation = script.key(redis_key);
120    invocation.arg(window_millis);
121    invocation.invoke_async(connection).await.map_err(Into::into)
122}
123
124#[cfg(not(feature = "redis-lua"))]
125async fn increment_counter(
126    connection: &mut MultiplexedConnection,
127    redis_key: &str,
128    window_millis: i64,
129) -> Result<(i64, i64), RedisStoreError> {
130    redis::pipe()
131        .atomic()
132        .cmd("SET")
133        .arg(redis_key)
134        .arg(0)
135        .arg("PX")
136        .arg(window_millis)
137        .arg("NX")
138        .ignore()
139        .cmd("INCR")
140        .arg(redis_key)
141        .cmd("PTTL")
142        .arg(redis_key)
143        .query_async(connection)
144        .await
145        .map_err(Into::into)
146}
147
148fn checked_window_millis(window: Duration) -> Result<i64, RedisStoreError> {
149    let window_millis = window.as_millis();
150    if window_millis == 0 {
151        return Err(RedisStoreError::WindowTooShort);
152    }
153    if window_millis > i64::MAX as u128 {
154        return Err(RedisStoreError::WindowTooLarge);
155    }
156    Ok(window_millis as i64)
157}
158
159/// Convert the atomic increment result `(count, PTTL milliseconds)` into [`Usage`].
160///
161/// Redis reports `-1` for a persistent key and `-2` for a missing key. Both, as well as a
162/// zero TTL, are errors because the fixed window cannot be trusted without a positive TTL.
163fn usage_from_increment_result((used, reset_after_millis): (i64, i64)) -> Result<Usage, RedisStoreError> {
164    if used < 1 {
165        return Err(RedisStoreError::InvalidUsage(used));
166    }
167    if reset_after_millis <= 0 {
168        return Err(RedisStoreError::InvalidResetAfter(reset_after_millis));
169    }
170
171    Ok(Usage {
172        used: used as u64,
173        reset_after: Duration::from_millis(reset_after_millis as u64),
174    })
175}
176
177/// Redis transport naming for a scoped Key already owned by the Rate Limiter.
178fn format_redis_key(namespace: Option<&str>, key: &str) -> String {
179    match namespace.filter(|value| !value.is_empty()) {
180        Some(namespace) => format!("{namespace}:{REDIS_PREFIX}{key}"),
181        None => format!("{REDIS_PREFIX}{key}"),
182    }
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188
189    #[test]
190    fn redis_store_implements_the_common_store_seam() {
191        fn assert_store<T: Store>() {}
192        assert_store::<RedisStore>();
193    }
194
195    #[test]
196    fn increment_result_requires_positive_usage_and_ttl() {
197        let usage = usage_from_increment_result((4, 1_500)).expect("valid Redis increment result");
198        assert_eq!(
199            usage,
200            Usage {
201                used: 4,
202                reset_after: Duration::from_millis(1_500),
203            }
204        );
205
206        assert!(matches!(
207            usage_from_increment_result((0, 1_500)),
208            Err(RedisStoreError::InvalidUsage(0))
209        ));
210        assert!(matches!(
211            usage_from_increment_result((4, 0)),
212            Err(RedisStoreError::InvalidResetAfter(0))
213        ));
214        assert!(matches!(
215            usage_from_increment_result((4, -1)),
216            Err(RedisStoreError::InvalidResetAfter(-1))
217        ));
218    }
219
220    #[test]
221    fn redis_transport_key_keeps_namespace_private_to_the_adapter() {
222        assert_eq!(format_redis_key(None, "policy:client"), "rl:policy:client");
223        assert_eq!(
224            format_redis_key(Some("tenant"), "policy:client"),
225            "tenant:rl:policy:client"
226        );
227        assert_eq!(format_redis_key(Some(""), "policy:client"), "rl:policy:client");
228    }
229
230    #[test]
231    fn window_must_be_representable_as_positive_redis_milliseconds() {
232        assert!(matches!(
233            checked_window_millis(Duration::ZERO),
234            Err(RedisStoreError::WindowTooShort)
235        ));
236        assert!(matches!(
237            checked_window_millis(Duration::from_nanos(1)),
238            Err(RedisStoreError::WindowTooShort)
239        ));
240        assert_eq!(
241            checked_window_millis(Duration::from_millis(1)).expect("one millisecond"),
242            1
243        );
244        assert_eq!(
245            checked_window_millis(Duration::from_micros(1_500)).expect("sub-millisecond remainder is truncated"),
246            1
247        );
248        assert_eq!(
249            checked_window_millis(Duration::from_millis(i64::MAX as u64)),
250            Ok(i64::MAX)
251        );
252        assert!(matches!(
253            checked_window_millis(Duration::from_millis(i64::MAX as u64 + 1)),
254            Err(RedisStoreError::WindowTooLarge)
255        ));
256    }
257
258    #[test]
259    fn redis_store_errors_map_to_one_store_error_code() {
260        let error = RateLimitError::from(RedisStoreError::InvalidUsage(0));
261
262        assert_eq!(
263            error,
264            RateLimitError::Store(
265                String::from("redis_store_error"),
266                String::from("redis returned invalid usage 0"),
267            )
268        );
269    }
270}
tower_rate_limiter/store/redis.rs

tower_rate_limiter/store/
redis.rs
