use crate::{
BoxFuture, Container, Injectable, InternalServerErrorException, Module, ModuleMetadata,
PayloadTooLargeException,
};
use serde::{Serialize, de::DeserializeOwned};
use serde_json::Value;
use std::{
collections::HashMap,
sync::{Arc, RwLock},
time::{Duration, Instant},
};
const DEFAULT_MAX_ENTRIES: usize = 1024;
const DEFAULT_MAX_VALUE_BYTES: usize = 1024 * 1024;
pub trait CacheStore: Send + Sync + 'static {
fn get(&self, key: String) -> BoxFuture<'_, crate::Result<Option<Value>>>;
fn set(
&self,
key: String,
value: Value,
ttl: Option<Duration>,
) -> BoxFuture<'_, crate::Result<()>>;
fn delete(&self, key: String) -> BoxFuture<'_, crate::Result<()>>;
fn clear(&self) -> BoxFuture<'_, crate::Result<()>>;
}
#[derive(Clone)]
struct CacheEntry {
value: Value,
inserted_at: Instant,
expires_at: Option<Instant>,
}
impl CacheEntry {
fn new(value: Value, ttl: Option<Duration>) -> Self {
let now = Instant::now();
Self {
value,
inserted_at: now,
expires_at: ttl.and_then(|ttl| now.checked_add(ttl)),
}
}
fn is_expired(&self) -> bool {
self.expires_at
.is_some_and(|expires_at| Instant::now() >= expires_at)
}
}
#[derive(Clone, Debug)]
pub struct MemoryCacheOptions {
pub max_entries: usize,
pub max_value_bytes: usize,
pub default_ttl: Option<Duration>,
}
impl Default for MemoryCacheOptions {
fn default() -> Self {
Self {
max_entries: DEFAULT_MAX_ENTRIES,
max_value_bytes: DEFAULT_MAX_VALUE_BYTES,
default_ttl: None,
}
}
}
pub struct MemoryCacheStore {
options: MemoryCacheOptions,
entries: RwLock<HashMap<String, CacheEntry>>,
}
impl MemoryCacheStore {
pub fn new() -> Self {
Self::with_options(MemoryCacheOptions::default())
}
pub fn with_options(options: MemoryCacheOptions) -> Self {
Self {
options,
entries: RwLock::new(HashMap::new()),
}
}
pub fn len(&self) -> usize {
self.entries
.read()
.expect("cache store lock poisoned")
.len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
fn remove_expired(entries: &mut HashMap<String, CacheEntry>) {
entries.retain(|_, entry| !entry.is_expired());
}
fn evict_to_capacity(&self, entries: &mut HashMap<String, CacheEntry>) {
if self.options.max_entries == 0 {
entries.clear();
return;
}
while entries.len() > self.options.max_entries {
let Some(oldest_key) = entries
.iter()
.min_by_key(|(_, entry)| entry.inserted_at)
.map(|(key, _)| key.clone())
else {
return;
};
entries.remove(&oldest_key);
}
}
}
impl Default for MemoryCacheStore {
fn default() -> Self {
Self::new()
}
}
impl Injectable for MemoryCacheStore {
fn create(_container: &Container) -> BoxFuture<'_, Self> {
Box::pin(async { Self::new() })
}
}
impl CacheStore for MemoryCacheStore {
fn get(&self, key: String) -> BoxFuture<'_, crate::Result<Option<Value>>> {
Box::pin(async move {
let mut entries = self.entries.write().expect("cache store lock poisoned");
let Some(entry) = entries.get(&key) else {
return Ok(None);
};
if entry.is_expired() {
entries.remove(&key);
return Ok(None);
}
Ok(Some(entry.value.clone()))
})
}
fn set(
&self,
key: String,
value: Value,
ttl: Option<Duration>,
) -> BoxFuture<'_, crate::Result<()>> {
Box::pin(async move {
let value_size = serde_json::to_vec(&value)
.map_err(InternalServerErrorException::new)?
.len();
if value_size > self.options.max_value_bytes {
return Err(PayloadTooLargeException::new(format!(
"cache value exceeds the configured limit of {} bytes",
self.options.max_value_bytes
)));
}
let ttl = ttl.or(self.options.default_ttl);
let mut entries = self.entries.write().expect("cache store lock poisoned");
Self::remove_expired(&mut entries);
entries.insert(key, CacheEntry::new(value, ttl));
self.evict_to_capacity(&mut entries);
Ok(())
})
}
fn delete(&self, key: String) -> BoxFuture<'_, crate::Result<()>> {
Box::pin(async move {
self.entries
.write()
.expect("cache store lock poisoned")
.remove(&key);
Ok(())
})
}
fn clear(&self) -> BoxFuture<'_, crate::Result<()>> {
Box::pin(async move {
self.entries
.write()
.expect("cache store lock poisoned")
.clear();
Ok(())
})
}
}
pub struct Cache {
store: Arc<dyn CacheStore>,
}
impl Cache {
pub fn new(store: Arc<dyn CacheStore>) -> Self {
Self { store }
}
pub async fn get<T>(&self, key: impl Into<String>) -> crate::Result<Option<T>>
where
T: DeserializeOwned,
{
let Some(value) = self.store.get(key.into()).await? else {
return Ok(None);
};
serde_json::from_value(value)
.map(Some)
.map_err(InternalServerErrorException::new)
}
pub async fn set<T>(&self, key: impl Into<String>, value: T) -> crate::Result<()>
where
T: Serialize,
{
self.set_with_optional_ttl(key, value, None).await
}
pub async fn set_with_ttl<T>(
&self,
key: impl Into<String>,
value: T,
ttl: Duration,
) -> crate::Result<()>
where
T: Serialize,
{
self.set_with_optional_ttl(key, value, Some(ttl)).await
}
pub async fn set_with_optional_ttl<T>(
&self,
key: impl Into<String>,
value: T,
ttl: Option<Duration>,
) -> crate::Result<()>
where
T: Serialize,
{
let value = serde_json::to_value(value).map_err(InternalServerErrorException::new)?;
self.store.set(key.into(), value, ttl).await
}
pub async fn delete(&self, key: impl Into<String>) -> crate::Result<()> {
self.store.delete(key.into()).await
}
pub async fn clear(&self) -> crate::Result<()> {
self.store.clear().await
}
}
impl Injectable for Cache {
fn create(container: &Container) -> BoxFuture<'_, Self> {
Box::pin(async move {
let store = container.resolve::<MemoryCacheStore>() as Arc<dyn CacheStore>;
Self::new(store)
})
}
}
pub struct CacheModule;
impl Module for CacheModule {
fn register() -> ModuleMetadata {
ModuleMetadata::new()
.provider::<MemoryCacheStore>()
.provider::<Cache>()
}
}