use async_trait::async_trait;
use dashmap::DashMap;
use std::sync::{
Arc,
atomic::{AtomicUsize, Ordering},
};
use std::time::Instant;
mod tenant;
pub use tenant::TenantCache;
#[derive(Debug)]
#[non_exhaustive]
pub enum CacheError {
Driver(String),
Serialization(String),
InvalidKey(String),
InspectionUnsupported,
InvalidInspectionLimit,
}
impl std::fmt::Display for CacheError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
CacheError::Driver(msg) => write!(f, "Cache driver error: {}", msg),
CacheError::Serialization(msg) => write!(f, "Cache serialization error: {}", msg),
CacheError::InvalidKey(msg) => write!(f, "Invalid cache key: {}", msg),
CacheError::InspectionUnsupported => {
write!(f, "Cache metadata inspection is unsupported by this driver")
}
CacheError::InvalidInspectionLimit => write!(
f,
"Cache inspection limit must be between 1 and {MAX_CACHE_INSPECTION_ENTRIES}"
),
}
}
}
impl std::error::Error for CacheError {}
#[async_trait]
pub trait CacheDriver: Send + Sync {
async fn get(&self, key: &str) -> Result<Option<Arc<String>>, CacheError>;
async fn put(&self, key: &str, value: &str, ttl_secs: Option<u64>) -> Result<(), CacheError>;
async fn forget(&self, key: &str) -> Result<(), CacheError>;
async fn flush(&self) -> Result<(), CacheError>;
async fn has(&self, key: &str) -> Result<bool, CacheError>;
async fn inspect(&self, _limit: usize) -> Result<CacheInspection, CacheError> {
Err(CacheError::InspectionUnsupported)
}
}
#[path = "cache/inspection.rs"]
mod inspection;
pub use inspection::{CacheEntryMetadata, CacheInspection, MAX_CACHE_INSPECTION_ENTRIES};
#[derive(Clone)]
struct CacheEntry {
value: Arc<String>,
expires_at: Option<Instant>,
}
pub struct MemoryDriver {
store: DashMap<String, CacheEntry>,
operations_since_cleanup: AtomicUsize,
}
impl MemoryDriver {
pub fn new() -> Self {
Self {
store: DashMap::new(),
operations_since_cleanup: AtomicUsize::new(0),
}
}
fn cleanup_if_due(&self) {
const CLEANUP_INTERVAL_OPERATIONS: usize = 256;
let previous = self
.operations_since_cleanup
.fetch_add(1, Ordering::Relaxed);
if previous > 0 && previous.is_multiple_of(CLEANUP_INTERVAL_OPERATIONS) {
let now = Instant::now();
self.store
.retain(|_, entry| entry.expires_at.is_none_or(|expires_at| now < expires_at));
}
}
fn get_sync(&self, key: &str) -> Option<Arc<String>> {
self.cleanup_if_due();
let entry = self.store.get(key)?;
if entry
.expires_at
.is_some_and(|expires_at| Instant::now() >= expires_at)
{
drop(entry);
self.store.remove(key);
return None;
}
Some(Arc::clone(&entry.value))
}
fn put_sync(&self, key: &str, value: &str, ttl_secs: Option<u64>) {
self.cleanup_if_due();
let expires_at = ttl_secs.map(|secs| Instant::now() + std::time::Duration::from_secs(secs));
self.store.insert(
key.to_string(),
CacheEntry {
value: Arc::new(value.to_string()),
expires_at,
},
);
}
}
impl Default for MemoryDriver {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl CacheDriver for MemoryDriver {
#[cfg_attr(mutants, mutants::skip)]
async fn get(&self, key: &str) -> Result<Option<Arc<String>>, CacheError> {
Ok(self.get_sync(key))
}
async fn put(&self, key: &str, value: &str, ttl_secs: Option<u64>) -> Result<(), CacheError> {
self.put_sync(key, value, ttl_secs);
Ok(())
}
async fn forget(&self, key: &str) -> Result<(), CacheError> {
self.store.remove(key);
Ok(())
}
async fn flush(&self) -> Result<(), CacheError> {
self.store.clear();
Ok(())
}
async fn has(&self, key: &str) -> Result<bool, CacheError> {
Ok(self.get(key).await?.is_some())
}
async fn inspect(&self, limit: usize) -> Result<CacheInspection, CacheError> {
inspection::validate_inspection_limit(limit)?;
self.cleanup_if_due();
let now = Instant::now();
let mut entries: Vec<_> = self
.store
.iter()
.filter(|entry| entry.expires_at.is_none_or(|expires_at| expires_at > now))
.map(|entry| {
CacheEntryMetadata::new(
entry.key().clone(),
entry.value.len(),
entry
.expires_at
.map(|expires_at| expires_at.saturating_duration_since(now).as_millis())
.and_then(|ttl| u64::try_from(ttl).ok()),
)
})
.collect();
entries.sort_by(|left, right| left.logical_key.cmp(&right.logical_key));
let truncated = entries.len() > limit;
entries.truncate(limit);
Ok(CacheInspection::new(entries, truncated))
}
}
pub mod memory {
use super::MemoryDriver;
use std::sync::OnceLock;
static GLOBAL_MEMO_CACHE: OnceLock<MemoryDriver> = OnceLock::new();
fn get_cache() -> &'static MemoryDriver {
GLOBAL_MEMO_CACHE.get_or_init(MemoryDriver::new)
}
pub fn get(key: &str) -> Option<String> {
get_cache().get_sync(key).map(|value| value.to_string())
}
pub fn set(key: &str, value: &str) {
get_cache().put_sync(key, value, Some(3600));
}
}
#[cfg(feature = "cache-redis")]
#[path = "cache/redis_driver.rs"]
pub mod redis_driver;
#[derive(Clone)]
pub struct Cache {
driver: Arc<Box<dyn CacheDriver>>,
}
impl Cache {
pub fn memory() -> Self {
Self {
driver: Arc::new(Box::new(MemoryDriver::new())),
}
}
#[cfg(feature = "cache-redis")]
#[cfg_attr(mutants, mutants::skip)]
pub fn redis(redis_url: impl Into<String>) -> Result<Self, CacheError> {
let driver = redis_driver::RedisDriver::new(redis_url)?;
Ok(Self {
driver: Arc::new(Box::new(driver)),
})
}
pub fn custom(driver: Box<dyn CacheDriver>) -> Self {
Self {
driver: Arc::new(driver),
}
}
pub async fn get(&self, key: &str) -> Result<Option<Arc<String>>, CacheError> {
self.driver.get(key).await
}
pub async fn put(
&self,
key: &str,
value: &str,
ttl_secs: Option<u64>,
) -> Result<(), CacheError> {
self.driver.put(key, value, ttl_secs).await
}
pub async fn forget(&self, key: &str) -> Result<(), CacheError> {
self.driver.forget(key).await
}
pub async fn flush(&self) -> Result<(), CacheError> {
self.driver.flush().await
}
pub async fn has(&self, key: &str) -> Result<bool, CacheError> {
self.driver.has(key).await
}
pub async fn inspect(&self, limit: usize) -> Result<CacheInspection, CacheError> {
self.driver.inspect(limit).await
}
pub async fn remember<F, Fut>(
&self,
key: &str,
ttl_secs: u64,
f: F,
) -> Result<Arc<String>, CacheError>
where
F: FnOnce() -> Fut,
Fut: std::future::Future<Output = Result<String, CacheError>>,
{
if let Some(cached) = self.get(key).await? {
return Ok(cached);
}
let value = f().await?;
self.put(key, &value, Some(ttl_secs)).await?;
Ok(Arc::new(value))
}
}
#[cfg(test)]
mod tests;
#[cfg(test)]
#[path = "cache_contract_tests.rs"]
mod contract_tests;