cache: add valkey backend layer

crates/crank-runtime/src/cache.rs

@@ -12,6 +12,8 @@ use crank_core::{
     CoordinationStateValue, RateLimitBucketState, RateLimitStateStore, ReplayGuardStatus,
     ReplayGuardStore, ResponseCacheStore,
 };
+use redis::{Client, aio::ConnectionManager};
+use serde::de::DeserializeOwned;
 use thiserror::Error;
 use tokio::sync::RwLock;
 
@@ -50,6 +52,21 @@ impl RuntimeCacheConfig {
     }
 }
 
+#[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>>>>,
@@ -76,6 +93,160 @@ struct ExpiringValue<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> {
@@ -204,6 +375,105 @@ impl CoordinationStateStore for InMemoryCoordinationStateStore {
     }
 }
 
+#[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 {
@@ -293,6 +563,23 @@ pub enum RuntimeCacheConfigError {
     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;
@@ -302,7 +589,8 @@ mod tests {
 
     use super::{
         InMemoryCoordinationStateStore, InMemoryRateLimitStateStore, InMemoryReplayGuardStore,
-        InMemoryResponseCacheStore, RuntimeCacheConfig, RuntimeCacheConfigError,
+        InMemoryResponseCacheStore, RedisCacheStore, RuntimeCacheConfig, RuntimeCacheConfigError,
+        RuntimeCacheStoreInitError, RuntimeCacheStores,
     };
     use crank_core::{
         CacheScope, CacheStoreError, CachedHeader, CachedResponse, CoordinationStateStore,
@@ -520,4 +808,93 @@ mod tests {
             .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(),
+                },
+                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(),
+        };
+
+        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
+            }
+        );
+    }
 }