use std::{ collections::HashMap, env, num::ParseIntError, sync::Arc, time::{Duration, Instant}, }; use async_trait::async_trait; use crank_core::{ CacheBackend, CacheScope, CacheStoreError, CachedResponse, CoordinationStateStore, CoordinationStateValue, RateLimitBucketState, RateLimitStateStore, ReplayGuardStatus, ReplayGuardStore, ResponseCacheStore, }; use redis::{Client, aio::ConnectionManager}; use serde::de::DeserializeOwned; use thiserror::Error; use tokio::sync::RwLock; #[derive(Clone, Debug, PartialEq, Eq)] pub struct RuntimeCacheConfig { pub backend: CacheBackend, pub url: Option, pub default_ttl_ms: Option, } impl Default for RuntimeCacheConfig { fn default() -> Self { Self { backend: CacheBackend::Memory, url: None, default_ttl_ms: None, } } } impl RuntimeCacheConfig { pub fn from_env() -> Result { let backend = parse_backend()?; let url = parse_optional_string("CRANK_CACHE_URL")?; let default_ttl_ms = parse_optional_u64("CRANK_CACHE_DEFAULT_TTL_MS")?; if backend.is_external() && url.is_none() { return Err(RuntimeCacheConfigError::MissingUrl { backend }); } Ok(Self { backend, url, default_ttl_ms, }) } } #[derive(Clone)] pub struct RuntimeCacheStores { pub backend: CacheBackend, pub response: Arc, pub rate_limit: Arc, pub replay_guard: Arc, pub coordination: Arc, } #[derive(Clone)] pub struct RedisCacheStore { backend: CacheBackend, connection_manager: ConnectionManager, } #[derive(Clone, Debug, Default)] pub struct InMemoryResponseCacheStore { entries: Arc>>>, } #[derive(Clone, Debug, Default)] pub struct InMemoryRateLimitStateStore { entries: Arc>>>, } #[derive(Clone, Debug, Default)] pub struct InMemoryReplayGuardStore { entries: Arc>>, } #[derive(Clone, Debug, Default)] pub struct InMemoryCoordinationStateStore { entries: Arc>>>, } #[derive(Clone, Debug)] struct ExpiringValue { value: T, expires_at: Instant, } impl RuntimeCacheStores { pub async fn from_config( config: &RuntimeCacheConfig, ) -> Result { match config.backend { CacheBackend::Memory => { let response = Arc::new(InMemoryResponseCacheStore::default()); let rate_limit = Arc::new(InMemoryRateLimitStateStore::default()); let replay_guard = Arc::new(InMemoryReplayGuardStore::default()); let coordination = Arc::new(InMemoryCoordinationStateStore::default()); Ok(Self { backend: CacheBackend::Memory, response, rate_limit, replay_guard, coordination, }) } CacheBackend::Valkey | CacheBackend::Redis => { let url = config .url .as_deref() .ok_or(RuntimeCacheStoreInitError::MissingUrl { backend: config.backend, })?; let store = Arc::new(RedisCacheStore::connect(config.backend, url).await?); Ok(Self { backend: config.backend, response: store.clone(), rate_limit: store.clone(), replay_guard: store.clone(), coordination: store, }) } } } } impl RedisCacheStore { pub async fn connect( backend: CacheBackend, url: &str, ) -> Result { let client = Client::open(url).map_err(|source| RuntimeCacheStoreInitError::InvalidUrl { backend, url: url.to_owned(), details: source.to_string(), })?; let connection_manager = client.get_connection_manager().await.map_err(|source| { RuntimeCacheStoreInitError::ConnectFailed { backend, details: source.to_string(), } })?; Ok(Self { backend, connection_manager, }) } fn prefixed_key(&self, kind: &str, key: &str) -> String { format!("crank:{kind}:{key}") } fn serialize_value(&self, value: &T) -> Result, CacheStoreError> { serde_json::to_vec(value).map_err(|source| CacheStoreError::Serialization { message: source.to_string(), }) } fn deserialize_value(&self, value: &[u8]) -> Result { serde_json::from_slice(value).map_err(|source| CacheStoreError::Serialization { message: source.to_string(), }) } fn ttl_ms(&self, ttl: Duration) -> Result { if ttl.is_zero() { return Err(CacheStoreError::InvalidKey { message: "cache ttl must be greater than zero".to_owned(), }); } Ok(u64::try_from(ttl.as_millis()).unwrap_or(u64::MAX)) } fn unavailable(&self, source: redis::RedisError) -> CacheStoreError { CacheStoreError::Unavailable { message: format!("{} backend error: {source}", self.backend), } } async fn get_json( &self, kind: &str, key: &str, ) -> Result, CacheStoreError> { validate_key(key)?; let storage_key = self.prefixed_key(kind, key); let mut connection = self.connection_manager.clone(); let encoded: Option> = redis::cmd("GET") .arg(storage_key) .query_async(&mut connection) .await .map_err(|source| self.unavailable(source))?; encoded .map(|bytes| self.deserialize_value(&bytes)) .transpose() } async fn put_json( &self, kind: &str, key: &str, value: &T, ttl: Duration, ) -> Result<(), CacheStoreError> { validate_key(key)?; let storage_key = self.prefixed_key(kind, key); let encoded = self.serialize_value(value)?; let ttl_ms = self.ttl_ms(ttl)?; let mut connection = self.connection_manager.clone(); redis::cmd("PSETEX") .arg(storage_key) .arg(ttl_ms) .arg(encoded) .query_async::<()>(&mut connection) .await .map_err(|source| self.unavailable(source))?; Ok(()) } async fn delete_kind_key(&self, kind: &str, key: &str) -> Result<(), CacheStoreError> { validate_key(key)?; let storage_key = self.prefixed_key(kind, key); let mut connection = self.connection_manager.clone(); redis::cmd("DEL") .arg(storage_key) .query_async::<()>(&mut connection) .await .map_err(|source| self.unavailable(source))?; Ok(()) } fn coordination_kind(scope: CacheScope) -> &'static str { match scope { CacheScope::Response => "coordination:response", CacheScope::RateLimit => "coordination:rate_limit", CacheScope::ReplayGuard => "coordination:replay_guard", CacheScope::Coordination => "coordination:coordination", } } } #[async_trait] impl ResponseCacheStore for InMemoryResponseCacheStore { async fn get(&self, key: &str) -> Result, CacheStoreError> { validate_key(key)?; let now = Instant::now(); let mut entries = self.entries.write().await; retain_unexpired(&mut entries, now); Ok(entries.get(key).map(|entry| entry.value.clone())) } async fn put( &self, key: &str, value: CachedResponse, ttl: Duration, ) -> Result<(), CacheStoreError> { validate_key(key)?; let expires_at = expiry_from_ttl(ttl)?; let mut entries = self.entries.write().await; entries.insert(key.to_owned(), ExpiringValue { value, expires_at }); Ok(()) } async fn delete(&self, key: &str) -> Result<(), CacheStoreError> { validate_key(key)?; let mut entries = self.entries.write().await; entries.remove(key); Ok(()) } } #[async_trait] impl RateLimitStateStore for InMemoryRateLimitStateStore { async fn get_bucket(&self, key: &str) -> Result, CacheStoreError> { validate_key(key)?; let now = Instant::now(); let mut entries = self.entries.write().await; retain_unexpired(&mut entries, now); Ok(entries.get(key).map(|entry| entry.value)) } async fn put_bucket( &self, key: &str, value: RateLimitBucketState, ttl: Duration, ) -> Result<(), CacheStoreError> { validate_key(key)?; let expires_at = expiry_from_ttl(ttl)?; let mut entries = self.entries.write().await; entries.insert(key.to_owned(), ExpiringValue { value, expires_at }); Ok(()) } async fn delete_bucket(&self, key: &str) -> Result<(), CacheStoreError> { validate_key(key)?; let mut entries = self.entries.write().await; entries.remove(key); Ok(()) } } #[async_trait] impl ReplayGuardStore for InMemoryReplayGuardStore { async fn mark_seen( &self, key: &str, ttl: Duration, ) -> Result { validate_key(key)?; let expires_at = expiry_from_ttl(ttl)?; let now = Instant::now(); let mut entries = self.entries.write().await; entries.retain(|_, entry_expires_at| *entry_expires_at > now); if entries.get(key).is_some() { return Ok(ReplayGuardStatus::AlreadySeen); } entries.insert(key.to_owned(), expires_at); Ok(ReplayGuardStatus::Fresh) } async fn clear(&self, key: &str) -> Result<(), CacheStoreError> { validate_key(key)?; let mut entries = self.entries.write().await; entries.remove(key); Ok(()) } } #[async_trait] impl CoordinationStateStore for InMemoryCoordinationStateStore { async fn get_value( &self, scope: CacheScope, key: &str, ) -> Result, CacheStoreError> { validate_key(key)?; let storage_key = scoped_key(scope, key); let now = Instant::now(); let mut entries = self.entries.write().await; retain_unexpired(&mut entries, now); Ok(entries.get(&storage_key).map(|entry| entry.value.clone())) } async fn put_value( &self, scope: CacheScope, key: &str, value: CoordinationStateValue, ttl: Duration, ) -> Result<(), CacheStoreError> { validate_key(key)?; let expires_at = expiry_from_ttl(ttl)?; let storage_key = scoped_key(scope, key); let mut entries = self.entries.write().await; entries.insert(storage_key, ExpiringValue { value, expires_at }); Ok(()) } async fn delete_value(&self, scope: CacheScope, key: &str) -> Result<(), CacheStoreError> { validate_key(key)?; let storage_key = scoped_key(scope, key); let mut entries = self.entries.write().await; entries.remove(&storage_key); Ok(()) } } #[async_trait] impl ResponseCacheStore for RedisCacheStore { async fn get(&self, key: &str) -> Result, CacheStoreError> { self.get_json("response", key).await } async fn put( &self, key: &str, value: CachedResponse, ttl: Duration, ) -> Result<(), CacheStoreError> { self.put_json("response", key, &value, ttl).await } async fn delete(&self, key: &str) -> Result<(), CacheStoreError> { self.delete_kind_key("response", key).await } } #[async_trait] impl RateLimitStateStore for RedisCacheStore { async fn get_bucket(&self, key: &str) -> Result, CacheStoreError> { self.get_json("rate_limit", key).await } async fn put_bucket( &self, key: &str, value: RateLimitBucketState, ttl: Duration, ) -> Result<(), CacheStoreError> { self.put_json("rate_limit", key, &value, ttl).await } async fn delete_bucket(&self, key: &str) -> Result<(), CacheStoreError> { self.delete_kind_key("rate_limit", key).await } } #[async_trait] impl ReplayGuardStore for RedisCacheStore { async fn mark_seen( &self, key: &str, ttl: Duration, ) -> Result { validate_key(key)?; let storage_key = self.prefixed_key("replay_guard", key); let ttl_ms = self.ttl_ms(ttl)?; let mut connection = self.connection_manager.clone(); let result: Option = redis::cmd("SET") .arg(storage_key) .arg("1") .arg("PX") .arg(ttl_ms) .arg("NX") .query_async(&mut connection) .await .map_err(|source| self.unavailable(source))?; Ok(if result.is_some() { ReplayGuardStatus::Fresh } else { ReplayGuardStatus::AlreadySeen }) } async fn clear(&self, key: &str) -> Result<(), CacheStoreError> { self.delete_kind_key("replay_guard", key).await } } #[async_trait] impl CoordinationStateStore for RedisCacheStore { async fn get_value( &self, scope: CacheScope, key: &str, ) -> Result, CacheStoreError> { self.get_json(Self::coordination_kind(scope), key).await } async fn put_value( &self, scope: CacheScope, key: &str, value: CoordinationStateValue, ttl: Duration, ) -> Result<(), CacheStoreError> { self.put_json(Self::coordination_kind(scope), key, &value, ttl) .await } async fn delete_value(&self, scope: CacheScope, key: &str) -> Result<(), CacheStoreError> { self.delete_kind_key(Self::coordination_kind(scope), key) .await } } fn validate_key(key: &str) -> Result<(), CacheStoreError> { if key.trim().is_empty() { return Err(CacheStoreError::InvalidKey { message: "cache key must not be empty".to_owned(), }); } Ok(()) } fn expiry_from_ttl(ttl: Duration) -> Result { if ttl.is_zero() { return Err(CacheStoreError::InvalidKey { message: "cache ttl must be greater than zero".to_owned(), }); } Ok(Instant::now() + ttl) } fn retain_unexpired(entries: &mut HashMap>, now: Instant) { entries.retain(|_, entry| entry.expires_at > now); } fn scoped_key(scope: CacheScope, key: &str) -> String { format!("{scope:?}:{key}") } fn parse_backend() -> Result { match env::var("CRANK_CACHE_BACKEND") { Ok(raw) => raw .parse::() .map_err(|source| RuntimeCacheConfigError::InvalidBackend { value: raw, source }), Err(env::VarError::NotPresent) => Ok(CacheBackend::Memory), Err(env::VarError::NotUnicode(_)) => Err(RuntimeCacheConfigError::InvalidUnicode { name: "CRANK_CACHE_BACKEND", }), } } fn parse_optional_string(name: &'static str) -> Result, RuntimeCacheConfigError> { match env::var(name) { Ok(raw) => { let trimmed = raw.trim(); if trimmed.is_empty() { Ok(None) } else { Ok(Some(trimmed.to_owned())) } } Err(env::VarError::NotPresent) => Ok(None), Err(env::VarError::NotUnicode(_)) => Err(RuntimeCacheConfigError::InvalidUnicode { name }), } } fn parse_optional_u64(name: &'static str) -> Result, RuntimeCacheConfigError> { match env::var(name) { Ok(raw) => { let value = raw .parse::() .map_err(|source| RuntimeCacheConfigError::InvalidTtl { value: raw, source })?; if value == 0 { return Err(RuntimeCacheConfigError::ZeroTtl { name }); } Ok(Some(value)) } Err(env::VarError::NotPresent) => Ok(None), Err(env::VarError::NotUnicode(_)) => Err(RuntimeCacheConfigError::InvalidUnicode { name }), } } #[derive(Debug, Error)] pub enum RuntimeCacheConfigError { #[error("{name} must contain valid UTF-8")] InvalidUnicode { name: &'static str }, #[error("CRANK_CACHE_BACKEND must be one of memory, valkey, redis, got {value}")] InvalidBackend { value: String, source: crank_core::ParseCacheBackendError, }, #[error("CRANK_CACHE_DEFAULT_TTL_MS must be a positive integer, got {value}")] InvalidTtl { value: String, source: ParseIntError, }, #[error("{name} must be greater than zero")] ZeroTtl { name: &'static str }, #[error("{backend} backend requires CRANK_CACHE_URL")] MissingUrl { backend: CacheBackend }, } #[derive(Debug, Error, PartialEq, Eq)] pub enum RuntimeCacheStoreInitError { #[error("{backend} backend requires CRANK_CACHE_URL")] MissingUrl { backend: CacheBackend }, #[error("invalid {backend} cache url {url}: {details}")] InvalidUrl { backend: CacheBackend, url: String, details: String, }, #[error("failed to connect to {backend} cache backend: {details}")] ConnectFailed { backend: CacheBackend, details: String, }, }