axum_limit/backend/
redis.rs

1use super::{apply_policy, BackendError, RateLimitBackend};
2use crate::policy::RateLimitPolicy;
3use crate::quota::Quota;
4use crate::snapshot::RateLimitSnapshot;
5use async_trait::async_trait;
6use redis::{AsyncCommands, ServerErrorKind};
7use std::sync::Arc;
8use std::time::Duration;
9use tokio::sync::Mutex;
10
11const MAX_RETRIES: usize = 8;
12
13fn ttl_seconds_for_quota(quota: Quota) -> i64 {
14    let seconds = quota.per_ms.div_ceil(1000).saturating_mul(2).max(60);
15    i64::try_from(seconds).unwrap_or(i64::MAX)
16}
17
18/// Redis-backed storage for multi-node deployments.
19#[derive(Clone)]
20pub struct RedisBackend {
21    namespace: Arc<str>,
22    client: redis::Client,
23    connection: Arc<Mutex<redis::aio::MultiplexedConnection>>,
24}
25
26impl RedisBackend {
27    /// Connects to Redis using the provided URL.
28    pub async fn connect(url: impl AsRef<str>) -> Result<Self, BackendError> {
29        Self::connect_with_namespace(url, "axum-limit").await
30    }
31
32    /// Connects to Redis and sets a custom namespace.
33    pub async fn connect_with_namespace(
34        url: impl AsRef<str>,
35        namespace: impl Into<Arc<str>>,
36    ) -> Result<Self, BackendError> {
37        let client = redis::Client::open(url.as_ref()).map_err(BackendError::Redis)?;
38        let connection = client
39            .get_multiplexed_async_connection()
40            .await
41            .map_err(BackendError::Redis)?;
42
43        Ok(Self {
44            namespace: namespace.into(),
45            client,
46            connection: Arc::new(Mutex::new(connection)),
47        })
48    }
49
50    /// Returns the underlying Redis client.
51    pub fn client(&self) -> &redis::Client {
52        &self.client
53    }
54}
55
56#[async_trait]
57impl RateLimitBackend for RedisBackend {
58    type Error = BackendError;
59
60    fn namespace(&self) -> &str {
61        &self.namespace
62    }
63
64    async fn transact<P>(
65        &self,
66        storage_key: &str,
67        quota: Quota,
68        now_ms: u64,
69    ) -> Result<RateLimitSnapshot, Self::Error>
70    where
71        P: RateLimitPolicy,
72    {
73        let ttl = ttl_seconds_for_quota(quota);
74
75        for _ in 0..MAX_RETRIES {
76            let mut connection = self.connection.lock().await;
77
78            redis::cmd("WATCH")
79                .arg(storage_key)
80                .query_async::<()>(&mut *connection)
81                .await
82                .map_err(BackendError::Redis)?;
83
84            let payload: Option<Vec<u8>> = connection
85                .get(storage_key)
86                .await
87                .map_err(BackendError::Redis)?;
88
89            let (encoded, snapshot) = apply_policy::<P>(payload.as_deref(), quota, now_ms)?;
90
91            let mut pipe = redis::pipe();
92            pipe.atomic()
93                .set(storage_key, encoded)
94                .ignore()
95                .expire(storage_key, ttl)
96                .ignore();
97
98            match pipe.query_async::<Option<()>>(&mut *connection).await {
99                Ok(Some(())) => return Ok(snapshot),
100                Ok(None) => continue,
101                Err(error) => {
102                    if error.kind() == redis::ErrorKind::Server(ServerErrorKind::ExecAbort) {
103                        tokio::time::sleep(Duration::from_millis(1)).await;
104                        continue;
105                    }
106                    return Err(BackendError::Redis(error));
107                }
108            }
109        }
110
111        Err(BackendError::Contention)
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118    use crate::policy::TokenBucketPolicy;
119    use crate::quota::Quota;
120
121    #[tokio::test]
122    #[ignore = "requires a running Redis server at redis://127.0.0.1/"]
123    async fn transact_persists_state() {
124        let backend = RedisBackend::connect("redis://127.0.0.1/")
125            .await
126            .expect("redis connection");
127        let storage_key = format!(
128            "test:{}",
129            std::time::SystemTime::now()
130                .duration_since(std::time::UNIX_EPOCH)
131                .expect("clock")
132                .as_nanos()
133        );
134        let quota = Quota::per_second(2);
135
136        let first = backend
137            .transact::<TokenBucketPolicy>(&storage_key, quota, 1_000_000)
138            .await
139            .expect("first transact");
140        assert!(first.allowed);
141
142        let second = backend
143            .transact::<TokenBucketPolicy>(&storage_key, quota, 1_000_000)
144            .await
145            .expect("second transact");
146        assert!(second.allowed);
147
148        let third = backend
149            .transact::<TokenBucketPolicy>(&storage_key, quota, 1_000_000)
150            .await
151            .expect("third transact");
152        assert!(!third.allowed);
153    }
154}
axum_limit/backend/redis.rs

axum_limit/backend/
redis.rs
