use crate::error::CacheResult;
use crate::traits::CacheStore;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
pub struct TaggedCache<C: CacheStore> {
cache: Arc<C>,
tags: Arc<RwLock<HashMap<String, HashSet<String>>>>,
key_tags: Arc<RwLock<HashMap<String, HashSet<String>>>>,
}
impl<C: CacheStore> TaggedCache<C> {
pub fn new(cache: Arc<C>) -> Self {
Self {
cache,
tags: Arc::new(RwLock::new(HashMap::new())),
key_tags: Arc::new(RwLock::new(HashMap::new())),
}
}
pub async fn set_with_tags(
&self,
key: &str,
value: String,
tags: &[&str],
ttl: Option<Duration>,
) -> CacheResult<()> {
self.cache.set_json(key, value, ttl).await?;
let mut tags_map = self.tags.write().await;
let mut key_tags_map = self.key_tags.write().await;
for tag in tags {
tags_map
.entry(tag.to_string())
.or_insert_with(HashSet::new)
.insert(key.to_string());
}
let tag_set: HashSet<String> = tags.iter().map(|t| t.to_string()).collect();
key_tags_map.insert(key.to_string(), tag_set);
Ok(())
}
pub async fn get(&self, key: &str) -> CacheResult<Option<String>> {
self.cache.get_json(key).await
}
pub async fn delete(&self, key: &str) -> CacheResult<()> {
self.cache.delete(key).await?;
let mut tags_map = self.tags.write().await;
let mut key_tags_map = self.key_tags.write().await;
if let Some(tag_set) = key_tags_map.remove(key) {
for tag in tag_set {
if let Some(keys) = tags_map.get_mut(&tag) {
keys.remove(key);
if keys.is_empty() {
tags_map.remove(&tag);
}
}
}
}
Ok(())
}
pub async fn invalidate_tag(&self, tag: &str) -> CacheResult<()> {
let mut tags_map = self.tags.write().await;
let mut key_tags_map = self.key_tags.write().await;
if let Some(keys) = tags_map.remove(tag) {
let key_refs: Vec<&str> = keys.iter().map(|s| s.as_str()).collect();
self.cache.delete_many(&key_refs).await?;
for key in keys {
if let Some(tag_set) = key_tags_map.get_mut(&key) {
tag_set.remove(tag);
if tag_set.is_empty() {
key_tags_map.remove(&key);
}
}
}
}
Ok(())
}
pub async fn invalidate_tags(&self, tags: &[&str]) -> CacheResult<()> {
for tag in tags {
self.invalidate_tag(tag).await?;
}
Ok(())
}
pub async fn get_keys_by_tag(&self, tag: &str) -> Vec<String> {
let tags_map = self.tags.read().await;
tags_map
.get(tag)
.map(|keys| keys.iter().cloned().collect())
.unwrap_or_default()
}
pub async fn get_tags_for_key(&self, key: &str) -> Vec<String> {
let key_tags_map = self.key_tags.read().await;
key_tags_map
.get(key)
.map(|tags| tags.iter().cloned().collect())
.unwrap_or_default()
}
pub async fn list_tags(&self) -> Vec<String> {
let tags_map = self.tags.read().await;
tags_map.keys().cloned().collect()
}
}
impl<C: CacheStore> Clone for TaggedCache<C> {
fn clone(&self) -> Self {
Self {
cache: self.cache.clone(),
tags: self.tags.clone(),
key_tags: self.key_tags.clone(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::error::CacheResult;
use async_trait::async_trait;
#[derive(Clone)]
struct MockCache {
data: Arc<RwLock<HashMap<String, String>>>,
}
impl MockCache {
fn new() -> Self {
Self {
data: Arc::new(RwLock::new(HashMap::new())),
}
}
}
#[async_trait]
impl CacheStore for MockCache {
async fn get_json(&self, key: &str) -> CacheResult<Option<String>> {
Ok(self.data.read().await.get(key).cloned())
}
async fn set_json(
&self,
key: &str,
value: String,
_ttl: Option<Duration>,
) -> CacheResult<()> {
self.data.write().await.insert(key.to_string(), value);
Ok(())
}
async fn delete(&self, key: &str) -> CacheResult<()> {
self.data.write().await.remove(key);
Ok(())
}
async fn exists(&self, key: &str) -> CacheResult<bool> {
Ok(self.data.read().await.contains_key(key))
}
async fn clear(&self) -> CacheResult<()> {
self.data.write().await.clear();
Ok(())
}
async fn ttl(&self, _key: &str) -> CacheResult<Option<Duration>> {
Ok(None)
}
async fn expire(&self, _key: &str, _ttl: Duration) -> CacheResult<()> {
Ok(())
}
async fn increment(&self, _key: &str, _delta: i64) -> CacheResult<i64> {
Ok(0)
}
async fn decrement(&self, _key: &str, _delta: i64) -> CacheResult<i64> {
Ok(0)
}
}
#[tokio::test]
async fn test_tagged_cache() {
let cache = Arc::new(MockCache::new());
let tagged = TaggedCache::new(cache);
tagged
.set_with_tags("user:1", "Alice".to_string(), &["users", "active"], None)
.await
.unwrap();
tagged
.set_with_tags("user:2", "Bob".to_string(), &["users"], None)
.await
.unwrap();
let value = tagged.get("user:1").await.unwrap();
assert_eq!(value, Some("Alice".to_string()));
let user_keys = tagged.get_keys_by_tag("users").await;
assert_eq!(user_keys.len(), 2);
tagged.invalidate_tag("users").await.unwrap();
let value = tagged.get("user:1").await.unwrap();
assert_eq!(value, None);
}
#[tokio::test]
async fn test_multiple_tags() {
let cache = Arc::new(MockCache::new());
let tagged = TaggedCache::new(cache);
tagged
.set_with_tags("key1", "value1".to_string(), &["tag1", "tag2"], None)
.await
.unwrap();
let tags = tagged.get_tags_for_key("key1").await;
assert_eq!(tags.len(), 2);
tagged.invalidate_tag("tag1").await.unwrap();
let value = tagged.get("key1").await.unwrap();
assert_eq!(value, None);
}
}