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<Self, CacheError> {
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<String> {
let mut conn = self.conn.clone();
redis::cmd("GET")
.arg(key)
.query_async::<Option<String>>(&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<Vec<u8>> {
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<i64, _> = 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::<Option<u64>>(&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<String> {
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<Vec<u8>> {
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<dyn Cache>, Arc<dyn DistributedRateLimiter>), 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))
}