tower_http_cache/backend/
memory.rs1use 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#[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 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 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}