armature_cache/
invalidation.rs1use crate::error::CacheResult;
4use crate::traits::CacheStore;
5use std::collections::{HashMap, HashSet};
6use std::sync::Arc;
7use std::time::Duration;
8use tokio::sync::RwLock;
9
10pub struct TaggedCache<C: CacheStore> {
12 cache: Arc<C>,
14
15 tags: Arc<RwLock<HashMap<String, HashSet<String>>>>,
17
18 key_tags: Arc<RwLock<HashMap<String, HashSet<String>>>>,
20}
21
22impl<C: CacheStore> TaggedCache<C> {
23 pub fn new(cache: Arc<C>) -> Self {
34 Self {
35 cache,
36 tags: Arc::new(RwLock::new(HashMap::new())),
37 key_tags: Arc::new(RwLock::new(HashMap::new())),
38 }
39 }
40
41 pub async fn set_with_tags(
54 &self,
55 key: &str,
56 value: String,
57 tags: &[&str],
58 ttl: Option<Duration>,
59 ) -> CacheResult<()> {
60 self.cache.set_json(key, value, ttl).await?;
62
63 let mut tags_map = self.tags.write().await;
65 let mut key_tags_map = self.key_tags.write().await;
66
67 for tag in tags {
68 tags_map
69 .entry(tag.to_string())
70 .or_insert_with(HashSet::new)
71 .insert(key.to_string());
72 }
73
74 let tag_set: HashSet<String> = tags.iter().map(|t| t.to_string()).collect();
75 key_tags_map.insert(key.to_string(), tag_set);
76
77 Ok(())
78 }
79
80 pub async fn get(&self, key: &str) -> CacheResult<Option<String>> {
82 self.cache.get_json(key).await
83 }
84
85 pub async fn delete(&self, key: &str) -> CacheResult<()> {
87 self.cache.delete(key).await?;
89
90 let mut tags_map = self.tags.write().await;
92 let mut key_tags_map = self.key_tags.write().await;
93
94 if let Some(tag_set) = key_tags_map.remove(key) {
95 for tag in tag_set {
96 if let Some(keys) = tags_map.get_mut(&tag) {
97 keys.remove(key);
98 if keys.is_empty() {
99 tags_map.remove(&tag);
100 }
101 }
102 }
103 }
104
105 Ok(())
106 }
107
108 pub async fn invalidate_tag(&self, tag: &str) -> CacheResult<()> {
117 let mut tags_map = self.tags.write().await;
118 let mut key_tags_map = self.key_tags.write().await;
119
120 if let Some(keys) = tags_map.remove(tag) {
121 let key_refs: Vec<&str> = keys.iter().map(|s| s.as_str()).collect();
123 self.cache.delete_many(&key_refs).await?;
124
125 for key in keys {
127 if let Some(tag_set) = key_tags_map.get_mut(&key) {
128 tag_set.remove(tag);
129 if tag_set.is_empty() {
130 key_tags_map.remove(&key);
131 }
132 }
133 }
134 }
135
136 Ok(())
137 }
138
139 pub async fn invalidate_tags(&self, tags: &[&str]) -> CacheResult<()> {
148 for tag in tags {
149 self.invalidate_tag(tag).await?;
150 }
151 Ok(())
152 }
153
154 pub async fn get_keys_by_tag(&self, tag: &str) -> Vec<String> {
156 let tags_map = self.tags.read().await;
157 tags_map
158 .get(tag)
159 .map(|keys| keys.iter().cloned().collect())
160 .unwrap_or_default()
161 }
162
163 pub async fn get_tags_for_key(&self, key: &str) -> Vec<String> {
165 let key_tags_map = self.key_tags.read().await;
166 key_tags_map
167 .get(key)
168 .map(|tags| tags.iter().cloned().collect())
169 .unwrap_or_default()
170 }
171
172 pub async fn list_tags(&self) -> Vec<String> {
174 let tags_map = self.tags.read().await;
175 tags_map.keys().cloned().collect()
176 }
177}
178
179impl<C: CacheStore> Clone for TaggedCache<C> {
180 fn clone(&self) -> Self {
181 Self {
182 cache: self.cache.clone(),
183 tags: self.tags.clone(),
184 key_tags: self.key_tags.clone(),
185 }
186 }
187}
188
189#[cfg(test)]
190mod tests {
191 use super::*;
192 use crate::error::CacheResult;
193 use async_trait::async_trait;
194
195 #[derive(Clone)]
197 struct MockCache {
198 data: Arc<RwLock<HashMap<String, String>>>,
199 }
200
201 impl MockCache {
202 fn new() -> Self {
203 Self {
204 data: Arc::new(RwLock::new(HashMap::new())),
205 }
206 }
207 }
208
209 #[async_trait]
210 impl CacheStore for MockCache {
211 async fn get_json(&self, key: &str) -> CacheResult<Option<String>> {
212 Ok(self.data.read().await.get(key).cloned())
213 }
214
215 async fn set_json(
216 &self,
217 key: &str,
218 value: String,
219 _ttl: Option<Duration>,
220 ) -> CacheResult<()> {
221 self.data.write().await.insert(key.to_string(), value);
222 Ok(())
223 }
224
225 async fn delete(&self, key: &str) -> CacheResult<()> {
226 self.data.write().await.remove(key);
227 Ok(())
228 }
229
230 async fn exists(&self, key: &str) -> CacheResult<bool> {
231 Ok(self.data.read().await.contains_key(key))
232 }
233
234 async fn clear(&self) -> CacheResult<()> {
235 self.data.write().await.clear();
236 Ok(())
237 }
238
239 async fn ttl(&self, _key: &str) -> CacheResult<Option<Duration>> {
240 Ok(None)
241 }
242
243 async fn expire(&self, _key: &str, _ttl: Duration) -> CacheResult<()> {
244 Ok(())
245 }
246
247 async fn increment(&self, _key: &str, _delta: i64) -> CacheResult<i64> {
248 Ok(0)
249 }
250
251 async fn decrement(&self, _key: &str, _delta: i64) -> CacheResult<i64> {
252 Ok(0)
253 }
254 }
255
256 #[tokio::test]
257 async fn test_tagged_cache() {
258 let cache = Arc::new(MockCache::new());
259 let tagged = TaggedCache::new(cache);
260
261 tagged
263 .set_with_tags("user:1", "Alice".to_string(), &["users", "active"], None)
264 .await
265 .unwrap();
266
267 tagged
268 .set_with_tags("user:2", "Bob".to_string(), &["users"], None)
269 .await
270 .unwrap();
271
272 let value = tagged.get("user:1").await.unwrap();
274 assert_eq!(value, Some("Alice".to_string()));
275
276 let user_keys = tagged.get_keys_by_tag("users").await;
278 assert_eq!(user_keys.len(), 2);
279
280 tagged.invalidate_tag("users").await.unwrap();
282
283 let value = tagged.get("user:1").await.unwrap();
285 assert_eq!(value, None);
286 }
287
288 #[tokio::test]
289 async fn test_multiple_tags() {
290 let cache = Arc::new(MockCache::new());
291 let tagged = TaggedCache::new(cache);
292
293 tagged
294 .set_with_tags("key1", "value1".to_string(), &["tag1", "tag2"], None)
295 .await
296 .unwrap();
297
298 let tags = tagged.get_tags_for_key("key1").await;
299 assert_eq!(tags.len(), 2);
300
301 tagged.invalidate_tag("tag1").await.unwrap();
302
303 let value = tagged.get("key1").await.unwrap();
304 assert_eq!(value, None);
305 }
306}