use std::collections::HashMap;
use std::sync::{Arc, OnceLock};
use std::time::Duration;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use http::header::{CACHE_CONTROL, HeaderValue, VARY};
use serde::{Serialize, de::DeserializeOwned};
use sqlx::SqlitePool;
use tokio::sync::Mutex;
use tower_http::compression::CompressionLayer;
use tower_http::set_header::SetResponseHeaderLayer;
use umbral::prelude::*;
pub mod cache_page;
pub use cache_page::cache_page;
static AMBIENT_CACHE: OnceLock<Cache> = OnceLock::new();
pub fn ambient() -> Option<&'static Cache> {
AMBIENT_CACHE.get()
}
#[derive(Debug)]
pub enum CacheError {
#[cfg(feature = "redis")]
Redis(redis::RedisError),
Sqlx(sqlx::Error),
Other(String),
}
impl std::fmt::Display for CacheError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
#[cfg(feature = "redis")]
CacheError::Redis(e) => write!(f, "cache redis error: {e}"),
CacheError::Sqlx(e) => write!(f, "cache sqlite error: {e}"),
CacheError::Other(s) => write!(f, "cache error: {s}"),
}
}
}
impl std::error::Error for CacheError {}
#[cfg(feature = "redis")]
impl From<redis::RedisError> for CacheError {
fn from(e: redis::RedisError) -> Self {
CacheError::Redis(e)
}
}
impl From<sqlx::Error> for CacheError {
fn from(e: sqlx::Error) -> Self {
CacheError::Sqlx(e)
}
}
#[async_trait]
pub trait CacheBackend: Send + Sync {
async fn get(&self, key: &str) -> Option<Vec<u8>>;
async fn set(&self, key: &str, value: Vec<u8>, ttl: Option<Duration>);
async fn delete(&self, key: &str);
async fn clear(&self);
}
#[derive(Clone)]
pub struct Cache {
backend: Arc<dyn CacheBackend>,
}
impl Cache {
pub fn memory() -> Self {
Self {
backend: Arc::new(MemoryBackend::default()),
}
}
pub async fn sqlite(pool: SqlitePool) -> Result<Self, CacheError> {
let backend = SqliteBackend::new(pool).await?;
Ok(Self {
backend: Arc::new(backend),
})
}
#[cfg(feature = "redis")]
pub async fn redis(url: &str) -> Result<Self, CacheError> {
let backend = RedisBackend::connect(url).await?;
Ok(Self {
backend: Arc::new(backend),
})
}
pub fn with_backend(backend: Arc<dyn CacheBackend>) -> Self {
Self { backend }
}
pub async fn get<T: DeserializeOwned>(&self, key: &str) -> Option<T> {
let bytes = self.backend.get(key).await?;
serde_json::from_slice(&bytes).ok()
}
pub async fn set<T: Serialize + ?Sized>(
&self,
key: &str,
value: &T,
ttl: Option<Duration>,
) -> Result<(), serde_json::Error> {
let bytes = serde_json::to_vec(value)?;
self.backend.set(key, bytes, ttl).await;
Ok(())
}
pub async fn delete(&self, key: &str) {
self.backend.delete(key).await;
}
pub async fn clear(&self) {
self.backend.clear().await;
}
pub(crate) async fn get_bytes_raw(&self, key: &str) -> Option<Vec<u8>> {
self.backend.get(key).await
}
pub(crate) async fn set_bytes_raw(&self, key: &str, bytes: Vec<u8>, ttl: Option<Duration>) {
self.backend.set(key, bytes, ttl).await;
}
}
struct MemoryEntry {
value: Vec<u8>,
expires_at: Option<DateTime<Utc>>,
}
#[derive(Default)]
pub struct MemoryBackend {
inner: Mutex<HashMap<String, MemoryEntry>>,
}
#[async_trait]
impl CacheBackend for MemoryBackend {
async fn get(&self, key: &str) -> Option<Vec<u8>> {
let mut map = self.inner.lock().await;
if let Some(entry) = map.get(key) {
if let Some(exp) = entry.expires_at {
if Utc::now() >= exp {
map.remove(key);
return None;
}
}
return Some(entry.value.clone());
}
None
}
async fn set(&self, key: &str, value: Vec<u8>, ttl: Option<Duration>) {
let expires_at = ttl.and_then(|d| {
chrono::Duration::from_std(d)
.ok()
.and_then(|cd| Utc::now().checked_add_signed(cd))
});
self.inner
.lock()
.await
.insert(key.to_string(), MemoryEntry { value, expires_at });
}
async fn delete(&self, key: &str) {
self.inner.lock().await.remove(key);
}
async fn clear(&self) {
self.inner.lock().await.clear();
}
}
pub struct SqliteBackend {
pool: SqlitePool,
}
impl SqliteBackend {
pub async fn new(pool: SqlitePool) -> Result<Self, CacheError> {
sqlx::query(
"CREATE TABLE IF NOT EXISTS umbral_cache (
key TEXT PRIMARY KEY,
value BLOB NOT NULL,
expires_at TIMESTAMP NULL
)",
)
.execute(&pool)
.await
.map_err(CacheError::Sqlx)?;
Ok(Self { pool })
}
pub async fn sweep(&self) -> Result<u64, CacheError> {
let result = sqlx::query(
"DELETE FROM umbral_cache WHERE expires_at IS NOT NULL AND expires_at <= ?",
)
.bind(Utc::now())
.execute(&self.pool)
.await
.map_err(CacheError::Sqlx)?;
Ok(result.rows_affected())
}
}
#[async_trait]
impl CacheBackend for SqliteBackend {
async fn get(&self, key: &str) -> Option<Vec<u8>> {
let row: Option<(Vec<u8>, Option<DateTime<Utc>>)> =
sqlx::query_as("SELECT value, expires_at FROM umbral_cache WHERE key = ?")
.bind(key)
.fetch_optional(&self.pool)
.await
.ok()?;
let (value, expires_at) = row?;
if let Some(exp) = expires_at {
if Utc::now() >= exp {
return None;
}
}
Some(value)
}
async fn set(&self, key: &str, value: Vec<u8>, ttl: Option<Duration>) {
let expires_at = ttl.and_then(|d| {
chrono::Duration::from_std(d)
.ok()
.and_then(|cd| Utc::now().checked_add_signed(cd))
});
if let Err(e) = sqlx::query(
"INSERT INTO umbral_cache (key, value, expires_at) VALUES (?, ?, ?)
ON CONFLICT(key) DO UPDATE SET value = excluded.value, expires_at = excluded.expires_at",
)
.bind(key)
.bind(value)
.bind(expires_at)
.execute(&self.pool)
.await
{
tracing::warn!(error = %e, key, "umbral-cache: SQLite cache set failed (swallowed)");
}
}
async fn delete(&self, key: &str) {
if let Err(e) = sqlx::query("DELETE FROM umbral_cache WHERE key = ?")
.bind(key)
.execute(&self.pool)
.await
{
tracing::warn!(error = %e, key, "umbral-cache: SQLite cache delete failed (swallowed)");
}
}
async fn clear(&self) {
if let Err(e) = sqlx::query("DELETE FROM umbral_cache")
.execute(&self.pool)
.await
{
tracing::warn!(error = %e, "umbral-cache: SQLite cache clear failed (swallowed)");
}
}
}
#[cfg(feature = "redis")]
pub struct RedisBackend {
client: redis::aio::ConnectionManager,
prefix: String,
}
#[cfg(feature = "redis")]
impl RedisBackend {
pub async fn connect(url: &str) -> Result<Self, CacheError> {
Self::connect_with_prefix(url, Self::DEFAULT_PREFIX).await
}
pub const DEFAULT_PREFIX: &'static str = "umbral:cache:";
pub async fn connect_with_prefix(url: &str, prefix: &str) -> Result<Self, CacheError> {
let client = redis::Client::open(url).map_err(CacheError::Redis)?;
let manager = redis::aio::ConnectionManager::new(client)
.await
.map_err(CacheError::Redis)?;
Ok(Self {
client: manager,
prefix: prefix.to_string(),
})
}
fn k(&self, key: &str) -> String {
format!("{}{key}", self.prefix)
}
}
#[cfg(feature = "redis")]
#[async_trait]
impl CacheBackend for RedisBackend {
async fn get(&self, key: &str) -> Option<Vec<u8>> {
use redis::AsyncCommands;
let mut conn = self.client.clone();
conn.get::<_, Option<Vec<u8>>>(self.k(key))
.await
.ok()
.flatten()
}
async fn set(&self, key: &str, value: Vec<u8>, ttl: Option<Duration>) {
use redis::AsyncCommands;
let mut conn = self.client.clone();
let key = self.k(key);
let res: Result<(), _> = if let Some(dur) = ttl {
let secs = dur.as_secs().max(1);
conn.set_ex(&key, value, secs).await
} else {
conn.set(&key, value).await
};
if let Err(e) = res {
tracing::warn!(error = %e, key, "umbral-cache: Redis cache set failed (swallowed)");
}
}
async fn delete(&self, key: &str) {
use redis::AsyncCommands;
let mut conn = self.client.clone();
let key = self.k(key);
if let Err(e) = conn.del::<_, ()>(&key).await {
tracing::warn!(error = %e, key, "umbral-cache: Redis cache delete failed (swallowed)");
}
}
async fn clear(&self) {
use redis::AsyncCommands;
let mut conn = self.client.clone();
let pattern = format!("{}*", self.prefix);
let mut cursor: u64 = 0;
loop {
let scan: redis::RedisResult<(u64, Vec<String>)> = redis::cmd("SCAN")
.arg(cursor)
.arg("MATCH")
.arg(&pattern)
.arg("COUNT")
.arg(512)
.query_async(&mut conn)
.await;
let (next, keys) = match scan {
Ok(v) => v,
Err(e) => {
tracing::warn!(error = %e, "umbral-cache: Redis cache clear SCAN failed (swallowed)");
return;
}
};
if !keys.is_empty() {
if let Err(e) = conn.unlink::<_, ()>(keys).await {
tracing::warn!(error = %e, "umbral-cache: Redis cache clear UNLINK failed (swallowed)");
}
}
cursor = next;
if cursor == 0 {
break;
}
}
}
}
#[derive(Debug, Clone, Default)]
pub struct CacheHeaders {
pub compression: bool,
pub cache_control: Option<String>,
pub vary: Option<String>,
}
#[derive(Default)]
pub struct CachePlugin {
cache: Option<Cache>,
headers: CacheHeaders,
}
impl CachePlugin {
pub fn new(cache: Cache) -> Self {
Self {
cache: Some(cache),
headers: CacheHeaders::default(),
}
}
pub fn init(cache: Cache) {
if AMBIENT_CACHE.set(cache).is_err() {
panic!("CachePlugin::init called more than once");
}
}
pub fn with_compression(mut self) -> Self {
self.headers.compression = true;
self
}
pub fn cache_control(mut self, value: impl Into<String>) -> Self {
self.headers.cache_control = Some(value.into());
self
}
pub fn vary(mut self, value: impl Into<String>) -> Self {
self.headers.vary = Some(value.into());
self
}
}
impl Plugin for CachePlugin {
fn name(&self) -> &'static str {
"cache"
}
fn wrap_router(&self, router: Router) -> Router {
let h = &self.headers;
let mut router = router;
if let Some(ref val) = h.cache_control {
if let Ok(hv) = HeaderValue::from_str(val) {
router = router.layer(SetResponseHeaderLayer::overriding(CACHE_CONTROL, hv));
} else {
tracing::warn!(
value = %val,
"CachePlugin: cache_control value contains invalid header characters; \
Cache-Control header will NOT be emitted"
);
}
}
if let Some(ref val) = h.vary {
if let Ok(hv) = HeaderValue::from_str(val) {
router = router.layer(SetResponseHeaderLayer::overriding(VARY, hv));
} else {
tracing::warn!(
value = %val,
"CachePlugin: vary value contains invalid header characters; \
Vary header will NOT be emitted"
);
}
}
if h.compression {
router = router.layer(CompressionLayer::new());
}
router
}
fn on_ready(
&self,
_ctx: &umbral::plugin::AppContext,
) -> Result<(), umbral::plugin::PluginError> {
match &self.cache {
Some(cache) => {
if AMBIENT_CACHE.set(cache.clone()).is_err() {
tracing::warn!(
"CachePlugin::new: an ambient cache was already installed (via \
CachePlugin::init or another CachePlugin); ignoring this one."
);
}
}
None if AMBIENT_CACHE.get().is_none() => {
tracing::warn!(
"CachePlugin registered with no cache and none set via CachePlugin::init — \
cache_page layers will silently no-op. Use \
CachePlugin::new(Cache::memory())."
);
}
None => {}
}
Ok(())
}
}