use std::{
collections::BTreeMap,
sync::Arc,
time::{Duration, Instant},
};
use async_trait::async_trait;
use tokio::sync::RwLock;
use crate::cache::{CacheKey, CacheResult, CacheStore};
#[derive(Debug, Clone)]
struct Entry {
value: Vec<u8>,
expires_at: Option<Instant>,
}
#[derive(Debug, Clone, Default)]
pub struct MemoryCacheStore {
entries: Arc<RwLock<BTreeMap<String, Entry>>>,
}
impl MemoryCacheStore {
pub fn new() -> Self {
Self::default()
}
}
#[async_trait]
impl CacheStore for MemoryCacheStore {
async fn get_raw(&self, key: &CacheKey) -> CacheResult<Option<Vec<u8>>> {
let rendered = key.render();
let mut entries = self.entries.write().await;
let Some(entry) = entries.get(&rendered) else {
return Ok(None);
};
if entry
.expires_at
.is_some_and(|deadline| deadline <= Instant::now())
{
entries.remove(&rendered);
return Ok(None);
}
Ok(Some(entry.value.clone()))
}
async fn set_raw(
&self,
key: &CacheKey,
value: Vec<u8>,
ttl: Option<Duration>,
) -> CacheResult<()> {
let expires_at = ttl.map(|ttl| Instant::now() + ttl);
self.entries
.write()
.await
.insert(key.render(), Entry { value, expires_at });
Ok(())
}
async fn delete(&self, key: &CacheKey) -> CacheResult<()> {
self.entries.write().await.remove(&key.render());
Ok(())
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use crate::cache::{CacheKey, CacheStore, MemoryCacheStore};
#[tokio::test]
async fn memory_cache_stores_json_with_ttl() {
let store = MemoryCacheStore::new();
let key = CacheKey::new("test", ["user", "1"]);
store
.set_json(
&key,
&serde_json::json!({"name":"Ada"}),
Some(Duration::from_secs(1)),
)
.await
.expect("set");
let value: serde_json::Value = store.get_json(&key).await.expect("get").expect("value");
assert_eq!(value["name"], "Ada");
store.delete(&key).await.expect("delete");
assert!(store.get_raw(&key).await.expect("get").is_none());
}
}