crank/crates/crank-runtime/src/cache.rs

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<String>,
    pub default_ttl_ms: Option<u64>,
}

impl Default for RuntimeCacheConfig {
    fn default() -> Self {
        Self {
            backend: CacheBackend::Memory,
            url: None,
            default_ttl_ms: None,
        }
    }
}

impl RuntimeCacheConfig {
    pub fn from_env() -> Result<Self, RuntimeCacheConfigError> {
        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<dyn ResponseCacheStore>,
    pub rate_limit: Arc<dyn RateLimitStateStore>,
    pub replay_guard: Arc<dyn ReplayGuardStore>,
    pub coordination: Arc<dyn CoordinationStateStore>,
}

#[derive(Clone)]
pub struct RedisCacheStore {
    backend: CacheBackend,
    connection_manager: ConnectionManager,
}

#[derive(Clone, Debug, Default)]
pub struct InMemoryResponseCacheStore {
    entries: Arc<RwLock<HashMap<String, ExpiringValue<CachedResponse>>>>,
}

#[derive(Clone, Debug, Default)]
pub struct InMemoryRateLimitStateStore {
    entries: Arc<RwLock<HashMap<String, ExpiringValue<RateLimitBucketState>>>>,
}

#[derive(Clone, Debug, Default)]
pub struct InMemoryReplayGuardStore {
    entries: Arc<RwLock<HashMap<String, Instant>>>,
}

#[derive(Clone, Debug, Default)]
pub struct InMemoryCoordinationStateStore {
    entries: Arc<RwLock<HashMap<String, ExpiringValue<CoordinationStateValue>>>>,
}

#[derive(Clone, Debug)]
struct ExpiringValue<T> {
    value: T,
    expires_at: Instant,
}

impl RuntimeCacheStores {
    pub async fn from_config(
        config: &RuntimeCacheConfig,
    ) -> Result<Self, RuntimeCacheStoreInitError> {
        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<Self, RuntimeCacheStoreInitError> {
        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<T: serde::Serialize>(&self, value: &T) -> Result<Vec<u8>, CacheStoreError> {
        serde_json::to_vec(value).map_err(|source| CacheStoreError::Serialization {
            message: source.to_string(),
        })
    }

    fn deserialize_value<T: DeserializeOwned>(&self, value: &[u8]) -> Result<T, CacheStoreError> {
        serde_json::from_slice(value).map_err(|source| CacheStoreError::Serialization {
            message: source.to_string(),
        })
    }

    fn ttl_ms(&self, ttl: Duration) -> Result<u64, CacheStoreError> {
        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<T: DeserializeOwned>(
        &self,
        kind: &str,
        key: &str,
    ) -> Result<Option<T>, CacheStoreError> {
        validate_key(key)?;
        let storage_key = self.prefixed_key(kind, key);
        let mut connection = self.connection_manager.clone();
        let encoded: Option<Vec<u8>> = 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<T: serde::Serialize>(
        &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<Option<CachedResponse>, 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<Option<RateLimitBucketState>, 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<ReplayGuardStatus, CacheStoreError> {
        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<Option<CoordinationStateValue>, 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<Option<CachedResponse>, 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<Option<RateLimitBucketState>, 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<ReplayGuardStatus, CacheStoreError> {
        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<String> = 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<Option<CoordinationStateValue>, 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<Instant, CacheStoreError> {
    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<T>(entries: &mut HashMap<String, ExpiringValue<T>>, now: Instant) {
    entries.retain(|_, entry| entry.expires_at > now);
}

fn scoped_key(scope: CacheScope, key: &str) -> String {
    format!("{scope:?}:{key}")
}

fn parse_backend() -> Result<CacheBackend, RuntimeCacheConfigError> {
    match env::var("CRANK_CACHE_BACKEND") {
        Ok(raw) => raw
            .parse::<CacheBackend>()
            .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<Option<String>, 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<Option<u64>, RuntimeCacheConfigError> {
    match env::var(name) {
        Ok(raw) => {
            let value = raw
                .parse::<u64>()
                .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,
    },
}