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

use std::{fmt, str::FromStr, time::Duration};

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use thiserror::Error;

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CacheBackend {
    Memory,
    Valkey,
    Redis,
}

impl CacheBackend {
    pub fn is_external(self) -> bool {
        matches!(self, Self::Valkey | Self::Redis)
    }
}

impl fmt::Display for CacheBackend {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let value = match self {
            Self::Memory => "memory",
            Self::Valkey => "valkey",
            Self::Redis => "redis",
        };
        f.write_str(value)
    }
}

impl FromStr for CacheBackend {
    type Err = ParseCacheBackendError;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        match value {
            "memory" => Ok(Self::Memory),
            "valkey" => Ok(Self::Valkey),
            "redis" => Ok(Self::Redis),
            _ => Err(ParseCacheBackendError {
                value: value.to_owned(),
            }),
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CacheScope {
    Response,
    RateLimit,
    ReplayGuard,
    Coordination,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct CachedHeader {
    pub name: String,
    pub value: String,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct CachedResponse {
    pub status: u16,
    pub headers: Vec<CachedHeader>,
    pub body: Vec<u8>,
    pub data: Vec<u8>,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct RateLimitBucketState {
    pub tokens_micros: u64,
    pub last_refill_unix_ms: i64,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ReplayGuardStatus {
    Fresh,
    AlreadySeen,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct CoordinationStateValue {
    pub payload: Value,
}

#[async_trait]
pub trait ResponseCacheStore: Send + Sync {
    async fn get(&self, key: &str) -> Result<Option<CachedResponse>, CacheStoreError>;
    async fn put(
        &self,
        key: &str,
        value: CachedResponse,
        ttl: Duration,
    ) -> Result<(), CacheStoreError>;
    async fn delete(&self, key: &str) -> Result<(), CacheStoreError>;
}

#[async_trait]
pub trait RateLimitStateStore: Send + Sync {
    async fn get_bucket(&self, key: &str) -> Result<Option<RateLimitBucketState>, CacheStoreError>;
    async fn put_bucket(
        &self,
        key: &str,
        value: RateLimitBucketState,
        ttl: Duration,
    ) -> Result<(), CacheStoreError>;
    async fn delete_bucket(&self, key: &str) -> Result<(), CacheStoreError>;
}

#[async_trait]
pub trait ReplayGuardStore: Send + Sync {
    async fn mark_seen(
        &self,
        key: &str,
        ttl: Duration,
    ) -> Result<ReplayGuardStatus, CacheStoreError>;
    async fn clear(&self, key: &str) -> Result<(), CacheStoreError>;
}

#[async_trait]
pub trait CoordinationStateStore: Send + Sync {
    async fn get_value(
        &self,
        scope: CacheScope,
        key: &str,
    ) -> Result<Option<CoordinationStateValue>, CacheStoreError>;
    async fn put_value(
        &self,
        scope: CacheScope,
        key: &str,
        value: CoordinationStateValue,
        ttl: Duration,
    ) -> Result<(), CacheStoreError>;
    async fn delete_value(&self, scope: CacheScope, key: &str) -> Result<(), CacheStoreError>;
}

#[derive(Debug, Error, PartialEq, Eq)]
pub enum CacheStoreError {
    #[error("cache backend is unavailable: {message}")]
    Unavailable { message: String },
    #[error("cache key is invalid: {message}")]
    InvalidKey { message: String },
    #[error("cache value serialization failed: {message}")]
    Serialization { message: String },
}

#[derive(Debug, Error, PartialEq, Eq)]
#[error("unsupported cache backend {value}")]
pub struct ParseCacheBackendError {
    pub value: String,
}

#[cfg(test)]
mod tests {
    use super::{CacheBackend, CacheScope, ParseCacheBackendError, ReplayGuardStatus};

    #[test]
    fn cache_backend_roundtrip_is_stable() {
        for backend in [
            CacheBackend::Memory,
            CacheBackend::Valkey,
            CacheBackend::Redis,
        ] {
            let encoded = backend.to_string();
            let decoded = encoded.parse::<CacheBackend>().unwrap();
            assert_eq!(decoded, backend);
        }
    }

    #[test]
    fn rejects_unknown_cache_backend() {
        assert_eq!(
            "memcached".parse::<CacheBackend>(),
            Err(ParseCacheBackendError {
                value: "memcached".to_owned()
            })
        );
    }

    #[test]
    fn distinguishes_external_cache_backends() {
        assert!(!CacheBackend::Memory.is_external());
        assert!(CacheBackend::Valkey.is_external());
        assert!(CacheBackend::Redis.is_external());
    }

    #[test]
    fn serialized_contracts_use_snake_case_values() {
        assert_eq!(
            serde_json::to_string(&CacheScope::ReplayGuard).unwrap(),
            "\"replay_guard\""
        );
        assert_eq!(
            serde_json::to_string(&ReplayGuardStatus::AlreadySeen).unwrap(),
            "\"already_seen\""
        );
    }
}