Something went wrong. Try again.
A downstream fork of Tranquil PDS for use in margin.cafe
Something went wrong. Try again.
8.6 kB · 261 lines
Rust
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262pub use tranquil_infra::{Cache, CacheError, DistributedRateLimiter, RateLimitStatus};
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]) endif redis.call('TTL', KEYS[1]) == -1 then redis.call('EXPIRE', KEYS[1], ARGV[1]) endreturn 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 consume_rate_limit( &self, key: &str, limit: u32, window_ms: u64, points: u32, ) -> RateLimitStatus { let mut conn = self.conn.clone(); let full_key = format!("rl:{}", key); let window_ms = window_ms.max(1); let result: Result<(i64, i64, i64), _> = redis::Script::new( r"local c = tonumber(redis.call('GET', KEYS[1]) or '0')local cost = tonumber(ARGV[1])local limit = tonumber(ARGV[2])if c + cost > limit then local ttl = redis.call('PTTL', KEYS[1]) return {0, c, ttl}endlocal next = redis.call('INCRBY', KEYS[1], cost)local ttl = redis.call('PTTL', KEYS[1])if ttl < 0 then redis.call('PEXPIRE', KEYS[1], ARGV[3]) ttl = tonumber(ARGV[3])endreturn {1, next, ttl}", ) .key(&full_key) .arg(points) .arg(limit) .arg(window_ms) .invoke_async(&mut conn) .await; match result { Ok((allowed, count, ttl)) => { let count = u64::try_from(count).unwrap_or(u64::MAX); let retry_after_ms = u64::try_from(ttl).unwrap_or(window_ms); if allowed == 1 { RateLimitStatus::allowed(limit, count, retry_after_ms) } else { RateLimitStatus::rejected(limit, count, retry_after_ms) } } Err(e) => { tracing::warn!(error = %e, "redis rate limit script failed"); RateLimitStatus::backend_error(window_ms) } } }
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"); let valkey_url = cache_cfg.and_then(|c| c.valkey_url.as_deref());
#[cfg(feature = "valkey")] if backend == "valkey" { if let Some(url) = valkey_url { 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))}