Skip to main content

tower_http_cache/backend/
memory.rs

1use moka::future::Cache;
2use std::sync::Arc;
3use std::time::{Duration, SystemTime};
4
5use super::{CacheBackend, CacheEntry, CacheRead};
6use crate::error::CacheError;
7use crate::tags::TagIndex;
8
9/// An in-memory [`CacheBackend`] implementation backed by [`moka`].
10///
11/// The backend is cheap to clone and shares a single underlying cache.
12#[derive(Clone)]
13pub struct InMemoryBackend {
14    cache: Cache<String, StoredEntry>,
15    tag_index: Arc<TagIndex>,
16}
17
18#[derive(Clone)]
19struct StoredEntry {
20    entry: CacheEntry,
21    expires_at: SystemTime,
22    stale_until: SystemTime,
23}
24
25impl InMemoryBackend {
26    /// Creates a new in-memory cache with the provided `max_capacity`.
27    ///
28    /// The capacity is expressed in number of cached entries, not bytes.
29    pub fn new(max_capacity: u64) -> Self {
30        let cache = Cache::builder().max_capacity(max_capacity).build();
31        Self {
32            cache,
33            tag_index: Arc::new(TagIndex::new()),
34        }
35    }
36}
37
38impl CacheBackend for InMemoryBackend {
39    async fn get(&self, key: &str) -> Result<Option<CacheRead>, CacheError> {
40        if let Some(stored) = self.cache.get(key).await {
41            let now = SystemTime::now();
42            if now > stored.stale_until {
43                self.cache.invalidate(key).await;
44                return Ok(None);
45            }
46
47            Ok(Some(CacheRead {
48                entry: stored.entry.clone(),
49                expires_at: Some(stored.expires_at),
50                stale_until: Some(stored.stale_until),
51            }))
52        } else {
53            Ok(None)
54        }
55    }
56
57    async fn set(
58        &self,
59        key: String,
60        entry: CacheEntry,
61        ttl: Duration,
62        stale_for: Duration,
63    ) -> Result<(), CacheError> {
64        if ttl.is_zero() {
65            return Ok(());
66        }
67
68        let now = SystemTime::now();
69        let expires_at = now + ttl;
70        let stale_until = expires_at + stale_for;
71
72        // Index tags if present
73        if let Some(ref tags) = entry.tags {
74            if !tags.is_empty() {
75                self.tag_index.index(key.clone(), tags.clone());
76            }
77        }
78
79        let stored = StoredEntry {
80            entry,
81            expires_at,
82            stale_until,
83        };
84        self.cache.insert(key, stored).await;
85        Ok(())
86    }
87
88    async fn invalidate(&self, key: &str) -> Result<(), CacheError> {
89        self.cache.invalidate(key).await;
90        self.tag_index.remove(key);
91        Ok(())
92    }
93
94    async fn get_keys_by_tag(&self, tag: &str) -> Result<Vec<String>, CacheError> {
95        Ok(self.tag_index.get_keys_by_tag(tag))
96    }
97
98    async fn list_tags(&self) -> Result<Vec<String>, CacheError> {
99        Ok(self.tag_index.list_tags())
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106    use crate::backend::CacheEntry;
107    use bytes::Bytes;
108    use http::{StatusCode, Version};
109    use tokio::time::{Duration, sleep};
110
111    fn entry_with_body(body: &'static [u8]) -> CacheEntry {
112        CacheEntry::new(
113            StatusCode::OK,
114            Version::HTTP_11,
115            Vec::new(),
116            Bytes::from_static(body),
117        )
118    }
119
120    #[tokio::test]
121    async fn set_and_get_returns_cached_entry() {
122        let backend = InMemoryBackend::new(16);
123        let entry = entry_with_body(b"alpha");
124
125        backend
126            .set(
127                "key".into(),
128                entry.clone(),
129                Duration::from_secs(1),
130                Duration::from_secs(1),
131            )
132            .await
133            .expect("set succeeds");
134
135        let read = backend.get("key").await.expect("get succeeds");
136        let cached = read.expect("entry present");
137
138        assert_eq!(cached.entry.body, entry.body);
139        assert!(cached.expires_at.is_some());
140        assert!(cached.stale_until.is_some());
141    }
142
143    #[tokio::test]
144    async fn entry_invalidated_after_stale_window() {
145        let backend = InMemoryBackend::new(16);
146
147        backend
148            .set(
149                "key".into(),
150                entry_with_body(b"stale"),
151                Duration::from_millis(20),
152                Duration::from_millis(30),
153            )
154            .await
155            .expect("set succeeds");
156
157        sleep(Duration::from_millis(35)).await;
158        let read = backend.get("key").await.expect("get succeeds");
159        assert!(read.is_some(), "entry available during stale window");
160
161        sleep(Duration::from_millis(40)).await;
162        let read = backend.get("key").await.expect("get succeeds");
163        assert!(read.is_none(), "entry removed after stale window");
164    }
165}