axum_limit/backend/
memory.rs

1use super::{apply_policy, build_storage_key, RateLimitBackend};
2use crate::policy::RateLimitPolicy;
3use crate::quota::Quota;
4use crate::snapshot::RateLimitSnapshot;
5use async_trait::async_trait;
6use dashmap::DashMap;
7use std::sync::Arc;
8
9/// In-memory storage backend suitable for single-node deployments.
10#[derive(Clone)]
11pub struct MemoryBackend {
12    namespace: Arc<str>,
13    states: Arc<DashMap<String, Vec<u8>>>,
14}
15
16impl MemoryBackend {
17    /// Creates a backend with the default namespace.
18    pub fn new() -> Self {
19        Self::with_namespace("axum-limit")
20    }
21
22    /// Creates a backend with a custom namespace.
23    pub fn with_namespace(namespace: impl Into<Arc<str>>) -> Self {
24        Self {
25            namespace: namespace.into(),
26            states: Arc::new(DashMap::new()),
27        }
28    }
29}
30
31impl Default for MemoryBackend {
32    fn default() -> Self {
33        Self::new()
34    }
35}
36
37#[async_trait]
38impl RateLimitBackend for MemoryBackend {
39    type Error = crate::codec::CodecError;
40
41    fn namespace(&self) -> &str {
42        &self.namespace
43    }
44
45    async fn transact<P>(
46        &self,
47        storage_key: &str,
48        quota: Quota,
49        now_ms: u64,
50    ) -> Result<RateLimitSnapshot, Self::Error>
51    where
52        P: RateLimitPolicy,
53    {
54        let entry = self.states.entry(storage_key.to_string());
55        match entry {
56            dashmap::mapref::entry::Entry::Occupied(mut occupied) => {
57                let (encoded, snapshot) = apply_policy::<P>(Some(occupied.get()), quota, now_ms)?;
58                *occupied.get_mut() = encoded;
59                Ok(snapshot)
60            }
61            dashmap::mapref::entry::Entry::Vacant(vacant) => {
62                let (encoded, snapshot) = apply_policy::<P>(None, quota, now_ms)?;
63                vacant.insert(encoded);
64                Ok(snapshot)
65            }
66        }
67    }
68}
69
70impl MemoryBackend {
71    /// Builds a storage key for the given subject and quota.
72    pub fn storage_key<P>(namespace: &str, subject: &str, quota: Quota) -> String
73    where
74        P: RateLimitPolicy,
75    {
76        build_storage_key::<P>(namespace, subject, quota)
77    }
78}
axum_limit/backend/memory.rs

axum_limit/backend/
memory.rs
