pub use tranquil_infra::{ Cache, CacheError, DistributedRateLimiter, cache_keys, cached_json, read_json, write_json, }; use async_trait::async_trait; use std::sync::Arc; use std::time::Duration; #[cfg(feature = "valkey")] mod valkey { use super::*; use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64}; #[derive(Clone)] pub struct ValkeyCache { conn: redis::aio::ConnectionManager, } impl ValkeyCache { pub async fn new(url: &str) -> Result { let client = redis::Client::open(url).map_err(|e| CacheError::Connection(e.to_string()))?; let manager = client .get_connection_manager() .await .map_err(|e| CacheError::Connection(e.to_string()))?; Ok(Self { conn: manager }) } pub fn connection(&self) -> redis::aio::ConnectionManager { self.conn.clone() } } #[async_trait] impl Cache for ValkeyCache { async fn get(&self, key: &str) -> Option { let mut conn = self.conn.clone(); redis::cmd("GET") .arg(key) .query_async::>(&mut conn) .await .ok() .flatten() } async fn set(&self, key: &str, value: &str, ttl: Duration) -> Result<(), CacheError> { let mut conn = self.conn.clone(); redis::cmd("SET") .arg(key) .arg(value) .arg("PX") .arg(i64::try_from(ttl.as_millis()).unwrap_or(i64::MAX)) .query_async::<()>(&mut conn) .await .map_err(|e| CacheError::Connection(e.to_string())) } async fn delete(&self, key: &str) -> Result<(), CacheError> { let mut conn = self.conn.clone(); redis::cmd("DEL") .arg(key) .query_async::<()>(&mut conn) .await .map_err(|e| CacheError::Connection(e.to_string())) } async fn get_bytes(&self, key: &str) -> Option> { self.get(key).await.and_then(|s| BASE64.decode(&s).ok()) } async fn set_bytes( &self, key: &str, value: &[u8], ttl: Duration, ) -> Result<(), CacheError> { let encoded = BASE64.encode(value); self.set(key, &encoded, ttl).await } } #[derive(Clone)] pub struct RedisRateLimiter { conn: redis::aio::ConnectionManager, } impl RedisRateLimiter { pub fn new(conn: redis::aio::ConnectionManager) -> Self { Self { conn } } } #[async_trait] impl DistributedRateLimiter for RedisRateLimiter { async fn check_rate_limit(&self, key: &str, limit: u32, window_ms: u64) -> bool { let mut conn = self.conn.clone(); let full_key = format!("rl:{}", key); let window_secs = i64::try_from(window_ms.div_ceil(1000).max(1)).unwrap_or(i64::MAX); let result: Result = redis::Script::new( r"local c = redis.call('INCR', KEYS[1]) if c == 1 then redis.call('EXPIRE', KEYS[1], ARGV[1]) end if redis.call('TTL', KEYS[1]) == -1 then redis.call('EXPIRE', KEYS[1], ARGV[1]) end return c", ) .key(&full_key) .arg(window_secs) .invoke_async(&mut conn) .await; match result { Ok(count) => count <= i64::from(limit), Err(e) => { tracing::warn!(error = %e, "redis rate limit script failed, allowing request"); true } } } async fn peek_rate_limit_count(&self, key: &str, _window_ms: u64) -> u64 { let mut conn = self.conn.clone(); let full_key = format!("rl:{}", key); redis::cmd("GET") .arg(&full_key) .query_async::>(&mut conn) .await .ok() .flatten() .unwrap_or(0) } } } #[cfg(feature = "valkey")] pub use valkey::{RedisRateLimiter, ValkeyCache}; pub struct NoOpCache; #[async_trait] impl Cache for NoOpCache { async fn get(&self, _key: &str) -> Option { None } async fn set(&self, _key: &str, _value: &str, _ttl: Duration) -> Result<(), CacheError> { Ok(()) } async fn delete(&self, _key: &str) -> Result<(), CacheError> { Ok(()) } async fn get_bytes(&self, _key: &str) -> Option> { None } async fn set_bytes(&self, _key: &str, _value: &[u8], _ttl: Duration) -> Result<(), CacheError> { Ok(()) } fn is_available(&self) -> bool { false } } #[derive(Debug, thiserror::Error)] pub enum CacheInitError { #[error("ripple config: {0}")] Config(#[from] tranquil_ripple::RippleConfigError), #[error("ripple start: {0}")] Start(#[from] tranquil_ripple::RippleStartError), } pub async fn create_cache( shutdown: tokio_util::sync::CancellationToken, ) -> Result<(Arc, Arc), CacheInitError> { let cache_cfg = tranquil_config::try_get().map(|c| &c.cache); let backend = cache_cfg.map(|c| c.backend.as_str()).unwrap_or("ripple"); #[cfg(feature = "valkey")] if backend == "valkey" { if let Some(url) = cache_cfg.and_then(|c| c.valkey_url.as_deref()) { match ValkeyCache::new(url).await { Ok(cache) => { tracing::info!("using valkey cache at {url}"); let rate_limiter = Arc::new(RedisRateLimiter::new(cache.connection())); return Ok((Arc::new(cache), rate_limiter)); } Err(e) => { tracing::warn!("failed to connect to valkey: {e}. falling back to ripple."); } } } else { tracing::warn!("cache.backend is \"valkey\" but VALKEY_URL is not set. using ripple."); } } #[cfg(not(feature = "valkey"))] if backend == "valkey" { tracing::warn!( "cache.backend is \"valkey\" but binary was compiled without valkey feature. using ripple." ); } let config = tranquil_ripple::RippleConfig::from_config()?; let peer_count = config.seed_peers.len(); let (cache, rate_limiter, _bound_addr) = tranquil_ripple::RippleEngine::start(config, shutdown).await?; match peer_count { 0 => tracing::info!("ripple cache started as a single node"), n => tracing::info!("ripple cache started with {n} seed peers"), } Ok((cache, rate_limiter)) }