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, }, } #[cfg(test)] mod tests { use std::time::Duration; use crank_core::CacheBackend; use serde_json::json; use super::{ InMemoryCoordinationStateStore, InMemoryRateLimitStateStore, InMemoryReplayGuardStore, InMemoryResponseCacheStore, RedisCacheStore, RuntimeCacheConfig, RuntimeCacheConfigError, RuntimeCacheStoreInitError, RuntimeCacheStores, }; use crank_core::{ CacheScope, CacheStoreError, CachedHeader, CachedResponse, CoordinationStateStore, CoordinationStateValue, RateLimitBucketState, RateLimitStateStore, ReplayGuardStatus, ReplayGuardStore, ResponseCacheStore, }; #[test] fn defaults_to_in_memory_cache_without_url() { let config = RuntimeCacheConfig::default(); assert_eq!(config.backend, CacheBackend::Memory); assert_eq!(config.url, None); assert_eq!(config.default_ttl_ms, None); } #[test] fn loads_valkey_config_from_env() { unsafe { std::env::set_var("CRANK_CACHE_BACKEND", "valkey"); std::env::set_var("CRANK_CACHE_URL", "redis://cache:6379/0"); std::env::set_var("CRANK_CACHE_DEFAULT_TTL_MS", "15000"); } let config = RuntimeCacheConfig::from_env().unwrap(); assert_eq!(config.backend, CacheBackend::Valkey); assert_eq!(config.url.as_deref(), Some("redis://cache:6379/0")); assert_eq!(config.default_ttl_ms, Some(15_000)); unsafe { std::env::remove_var("CRANK_CACHE_BACKEND"); std::env::remove_var("CRANK_CACHE_URL"); std::env::remove_var("CRANK_CACHE_DEFAULT_TTL_MS"); } } #[test] fn rejects_external_backend_without_url() { unsafe { std::env::set_var("CRANK_CACHE_BACKEND", "redis"); std::env::remove_var("CRANK_CACHE_URL"); } let error = RuntimeCacheConfig::from_env().unwrap_err(); assert!(matches!( error, RuntimeCacheConfigError::MissingUrl { backend: CacheBackend::Redis } )); unsafe { std::env::remove_var("CRANK_CACHE_BACKEND"); } } #[test] fn rejects_zero_ttl() { unsafe { std::env::set_var("CRANK_CACHE_DEFAULT_TTL_MS", "0"); } let error = RuntimeCacheConfig::from_env().unwrap_err(); assert!(matches!( error, RuntimeCacheConfigError::ZeroTtl { name: "CRANK_CACHE_DEFAULT_TTL_MS" } )); unsafe { std::env::remove_var("CRANK_CACHE_DEFAULT_TTL_MS"); } } #[tokio::test] async fn in_memory_response_cache_roundtrips_values() { let store = InMemoryResponseCacheStore::default(); let value = CachedResponse { status: 200, headers: vec![CachedHeader { name: "content-type".to_owned(), value: "application/json".to_owned(), }], body: br#"{"ok":true}"#.to_vec(), data: br#"{"ok":true}"#.to_vec(), }; store .put("response:crm:list", value.clone(), Duration::from_secs(60)) .await .unwrap(); assert_eq!( store.get("response:crm:list").await.unwrap(), Some(value.clone()) ); store.delete("response:crm:list").await.unwrap(); assert_eq!(store.get("response:crm:list").await.unwrap(), None); } #[tokio::test] async fn in_memory_rate_limit_store_roundtrips_bucket_state() { let store = InMemoryRateLimitStateStore::default(); let bucket = RateLimitBucketState { tokens_micros: 1_500_000, last_refill_unix_ms: 1_735_689_000_000, }; store .put_bucket("tenant:alpha", bucket, Duration::from_secs(30)) .await .unwrap(); assert_eq!( store.get_bucket("tenant:alpha").await.unwrap(), Some(bucket) ); store.delete_bucket("tenant:alpha").await.unwrap(); assert_eq!(store.get_bucket("tenant:alpha").await.unwrap(), None); } #[tokio::test] async fn in_memory_replay_guard_marks_key_only_once_until_cleared() { let store = InMemoryReplayGuardStore::default(); assert_eq!( store .mark_seen("token:nonce:1", Duration::from_secs(10)) .await .unwrap(), ReplayGuardStatus::Fresh ); assert_eq!( store .mark_seen("token:nonce:1", Duration::from_secs(10)) .await .unwrap(), ReplayGuardStatus::AlreadySeen ); store.clear("token:nonce:1").await.unwrap(); assert_eq!( store .mark_seen("token:nonce:1", Duration::from_secs(10)) .await .unwrap(), ReplayGuardStatus::Fresh ); } #[tokio::test] async fn in_memory_coordination_store_scopes_keys() { let store = InMemoryCoordinationStateStore::default(); let response_value = CoordinationStateValue { payload: json!({ "cursor": "abc" }), }; let session_value = CoordinationStateValue { payload: json!({ "cursor": "xyz" }), }; store .put_value( CacheScope::Response, "job-1", response_value.clone(), Duration::from_secs(20), ) .await .unwrap(); store .put_value( CacheScope::Coordination, "job-1", session_value.clone(), Duration::from_secs(20), ) .await .unwrap(); assert_eq!( store .get_value(CacheScope::Response, "job-1") .await .unwrap(), Some(response_value) ); assert_eq!( store .get_value(CacheScope::Coordination, "job-1") .await .unwrap(), Some(session_value) ); } #[tokio::test] async fn in_memory_stores_reject_empty_keys() { let response_store = InMemoryResponseCacheStore::default(); let replay_store = InMemoryReplayGuardStore::default(); let error = response_store.get("").await.unwrap_err(); assert!(matches!(error, CacheStoreError::InvalidKey { .. })); let error = replay_store .mark_seen("", Duration::from_secs(1)) .await .unwrap_err(); assert!(matches!(error, CacheStoreError::InvalidKey { .. })); } #[tokio::test] async fn runtime_cache_stores_default_to_memory_backend() { let stores = RuntimeCacheStores::from_config(&RuntimeCacheConfig::default()) .await .unwrap(); assert_eq!(stores.backend, CacheBackend::Memory); stores .response .put( "response:health", CachedResponse { status: 200, headers: vec![], body: b"ok".to_vec(), data: br#"null"#.to_vec(), }, Duration::from_secs(5), ) .await .unwrap(); assert!( stores .response .get("response:health") .await .unwrap() .is_some() ); } #[test] fn redis_cache_store_scopes_keys_by_kind() { let key = "agent:primary"; let response_key = format!("crank:{}:{}", "response", key); let coordination_key = format!( "crank:{}:{}", RedisCacheStore::coordination_kind(CacheScope::Coordination), key ); assert_eq!(response_key, "crank:response:agent:primary"); assert_eq!( coordination_key, "crank:coordination:coordination:agent:primary" ); } #[test] fn redis_cache_store_roundtrips_serialized_values() { let response = CachedResponse { status: 202, headers: vec![CachedHeader { name: "x-cache".to_owned(), value: "hit".to_owned(), }], body: br#"{"queued":true}"#.to_vec(), data: br#"null"#.to_vec(), }; let encoded = serde_json::to_vec(&response).unwrap(); let decoded: CachedResponse = serde_json::from_slice(&encoded).unwrap(); assert_eq!(decoded, response); } #[test] fn runtime_cache_stores_report_missing_external_url() { let future = RuntimeCacheStores::from_config(&RuntimeCacheConfig { backend: CacheBackend::Valkey, url: None, default_ttl_ms: None, }); let runtime = tokio::runtime::Runtime::new().unwrap(); let error = match runtime.block_on(future) { Ok(_) => panic!("expected missing external cache url error"), Err(error) => error, }; assert_eq!( error, RuntimeCacheStoreInitError::MissingUrl { backend: CacheBackend::Valkey } ); } }