use crate::error::RedisResultExt; use crate::redis::RedisKeyspace; use crate::DataLayerError; use std::future::Future; use std::pin::Pin; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::time::Duration; use tokio::sync::{OwnedSemaphorePermit, Semaphore}; use tracing::info; pub(crate) type RedisClient = redis::Client; pub(crate) type RedisManagedConnection = redis::aio::ConnectionManager; const DEFAULT_STREAM_LANES: usize = 4; const DEFAULT_BLOCKING_STREAM_LANES_FALLBACK: usize = 4; const DEFAULT_BLOCKING_STREAM_LANES_CAP: usize = 16; const MAX_BLOCKING_STREAM_LANES_CAP: usize = 64; pub(crate) const REDIS_COMMAND_LATENCY_BUCKETS_MS: [u64; 12] = [1, 5, 10, 25, 50, 100, 250, 500, 1_000, 2_500, 5_000, 10_000]; const REDIS_COMMAND_LATENCY_BUCKET_COUNT: usize = REDIS_COMMAND_LATENCY_BUCKETS_MS.len() + 1; #[derive(Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)] pub struct RedisClientConfig { pub url: String, pub key_prefix: Option, } impl std::fmt::Debug for RedisClientConfig { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { formatter .debug_struct("RedisClientConfig") .field("url", &redact_redis_url_for_debug(&self.url)) .field("key_prefix_len", &self.key_prefix.as_ref().map(String::len)) .finish() } } fn redact_redis_url_for_debug(raw: &str) -> String { const MAX_DEBUG_URL_CHARS: usize = 512; let raw = raw.trim(); let Ok(mut url) = url::Url::parse(raw) else { return format!("[invalid-redis-url len={}]", raw.len()); }; let _ = url.set_username(""); let _ = url.set_password(None); url.set_query(None); url.set_fragment(None); let rendered = url.to_string(); if rendered.chars().count() <= MAX_DEBUG_URL_CHARS { rendered } else { format!( "{}...", rendered .chars() .take(MAX_DEBUG_URL_CHARS.saturating_sub(3)) .collect::() ) } } impl RedisClientConfig { pub fn validate(&self) -> Result<(), DataLayerError> { let raw = self.url.trim(); if raw.is_empty() { return Err(DataLayerError::InvalidConfiguration( "redis url cannot be empty".to_string(), )); } url::Url::parse(raw).map_err(|err| { DataLayerError::InvalidConfiguration(format!("invalid redis url: {err}")) })?; Ok(()) } pub fn keyspace(&self) -> RedisKeyspace { RedisKeyspace::new(self.key_prefix.as_deref()) } } #[derive(Debug, Clone)] pub(crate) struct RedisClientFactory { config: RedisClientConfig, } impl RedisClientFactory { pub(crate) fn new(config: RedisClientConfig) -> Result { config.validate()?; Ok(Self { config }) } pub(crate) fn config(&self) -> &RedisClientConfig { &self.config } pub(crate) fn connect_lazy(&self) -> Result { RedisClient::open(self.config.url.clone()).map_redis_err() } pub(crate) async fn connect_router( &self, command_timeout_ms: Option, ) -> Result { self.connect_router_with_blocking_stream_lanes(command_timeout_ms, None) .await } pub(crate) async fn connect_router_with_blocking_stream_lanes( &self, command_timeout_ms: Option, blocking_stream_lanes: Option, ) -> Result { RedisConnectionRouter::connect( self.connect_lazy()?, command_timeout_ms, blocking_stream_lanes, ) .await } } #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum RedisConnectionLane { Fast, Stream, BlockingStream, Admin, } impl RedisConnectionLane { pub(crate) const fn as_str(self) -> &'static str { match self { Self::Fast => "fast", Self::Stream => "stream", Self::BlockingStream => "blocking_stream", Self::Admin => "admin", } } } #[derive(Clone)] pub(crate) struct RedisConnectionRouter { fast: RedisManagedConnection, stream: Arc>, stream_next: Arc, blocking_stream: Arc, usage_cleanup: Arc, admin: RedisManagedConnection, metrics: Arc, } impl std::fmt::Debug for RedisConnectionRouter { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("RedisConnectionRouter") .field("lanes", &["fast", "stream", "blocking_stream", "admin"]) .field("stream_lanes", &self.stream.len()) .field("blocking_stream_lanes", &self.blocking_stream.capacity) .finish() } } impl RedisConnectionRouter { pub(crate) async fn connect( client: RedisClient, command_timeout_ms: Option, blocking_stream_lanes: Option, ) -> Result { let fast = connect_lane( &client, connection_manager_config(command_timeout_ms), RedisConnectionLane::Fast, command_timeout_ms, ) .await?; let stream = connect_stream_lanes(&client, command_timeout_ms).await?; let blocking_stream = connect_blocking_stream_lanes(&client, command_timeout_ms, blocking_stream_lanes) .await?; let admin = connect_lane( &client, connection_manager_config(command_timeout_ms), RedisConnectionLane::Admin, command_timeout_ms, ) .await?; let stream_lanes = stream.len(); let blocking_stream_lanes = blocking_stream.capacity; info!( redis_lanes = "fast,stream,blocking_stream,admin", redis_stream_lanes = stream_lanes, redis_blocking_stream_lanes = blocking_stream_lanes, "runtime redis connection lanes initialized" ); Ok(Self { fast, stream: Arc::new(stream), stream_next: Arc::new(AtomicUsize::new(0)), blocking_stream: Arc::new(blocking_stream), usage_cleanup: Arc::new(RedisBlockingStreamPool::new( client, command_timeout_ms, vec![None, None], )), admin, metrics: Arc::new(RedisConnectionMetrics::default()), }) } pub(crate) fn connection(&self, lane: RedisConnectionLane) -> RedisManagedConnection { match lane { RedisConnectionLane::Fast => self.fast.clone(), RedisConnectionLane::Stream => { let index = next_lane_index(&self.stream_next, self.stream.len()); self.stream[index].clone() } RedisConnectionLane::BlockingStream => { unreachable!("blocking stream commands require an exclusive connection lease") } RedisConnectionLane::Admin => self.admin.clone(), } } pub(crate) async fn blocking_stream_connection( &self, ) -> Result { self.blocking_stream.checkout().await } // WATCH/MULTI state must never share a multiplexed connection with other callers. // These lazy leases own their drivers, so cancellation closes an unfinished transaction. pub(crate) async fn usage_cleanup_connection( &self, ) -> Result { self.usage_cleanup.checkout().await } pub(crate) fn record_error(&self, lane: RedisConnectionLane) { self.metrics .for_lane(lane) .errors .fetch_add(1, Ordering::Relaxed); } pub(crate) fn record_timeout(&self, lane: RedisConnectionLane) { self.metrics .for_lane(lane) .timeouts .fetch_add(1, Ordering::Relaxed); } pub(crate) fn record_latency(&self, lane: RedisConnectionLane, elapsed: Duration) { self.metrics.for_lane(lane).record_latency(elapsed); } pub(crate) fn lane_diagnostics(&self) -> Vec { [ RedisConnectionLane::Fast, RedisConnectionLane::Stream, RedisConnectionLane::BlockingStream, RedisConnectionLane::Admin, ] .into_iter() .map(|lane| { let metrics = self.metrics.for_lane(lane); RedisLaneDiagnostics { lane: lane.as_str(), command_errors: metrics.errors.load(Ordering::Relaxed), command_timeouts: metrics.timeouts.load(Ordering::Relaxed), command_count: metrics.command_count.load(Ordering::Relaxed), command_latency_total_ms: metrics.latency_total_ms.load(Ordering::Relaxed), command_latency_max_ms: metrics.latency_max_ms.load(Ordering::Relaxed), command_latency_buckets: metrics.latency_buckets(), } }) .collect() } } struct RedisBlockingConnection { connection: redis::aio::MultiplexedConnection, driver: Pin + Send>>, } impl RedisBlockingConnection { async fn query(&mut self, command: &redis::Cmd) -> Result { // Drive the connection inside its owning query, so cancellation drops the // socket immediately instead of leaving a spawned driver with an old BLOCK. tokio::select! { result = command.query_async(&mut self.connection) => result.map_redis_err(), () = self.driver.as_mut() => Err(DataLayerError::Redis( "runtime redis blocking stream connection driver terminated".to_string(), )), } } } struct RedisBlockingStreamPool { client: RedisClient, command_timeout_ms: Option, capacity: usize, available: Mutex>>, permits: Arc, } impl RedisBlockingStreamPool { fn new( client: RedisClient, command_timeout_ms: Option, available: Vec>, ) -> Self { let capacity = available.len(); Self { client, command_timeout_ms, capacity, available: Mutex::new(available), permits: Arc::new(Semaphore::new(capacity)), } } async fn checkout(self: &Arc) -> Result { let permit = Arc::clone(&self.permits) .acquire_owned() .await .map_err(|_| { DataLayerError::Redis("runtime redis blocking stream pool closed".to_string()) })?; let connection = self .available .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) .pop() .expect("blocking stream permit must have an available slot"); Ok(RedisBlockingStreamLease { pool: Arc::clone(self), connection, reusable: false, _permit: permit, }) } } pub(crate) struct RedisBlockingStreamLease { pool: Arc, connection: Option, reusable: bool, _permit: OwnedSemaphorePermit, } impl RedisBlockingStreamLease { pub(crate) async fn query( &mut self, command: &redis::Cmd, ) -> Result { self.reusable = false; if self.connection.is_none() { self.connection = Some( connect_blocking_stream_lane(&self.pool.client, self.pool.command_timeout_ms) .await?, ); } self.connection .as_mut() .expect("blocking stream connection initialized") .query(command) .await } pub(crate) fn recycle(&mut self) { self.reusable = true; } } impl Drop for RedisBlockingStreamLease { fn drop(&mut self) { let connection = self.connection.take().filter(|_| self.reusable); self.pool .available .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) .push(connection); // The permit is released after the slot is restored. An uncompleted // query has already dropped both its connection and its owned driver. } } #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)] pub struct RedisLaneDiagnostics { pub lane: &'static str, pub command_errors: u64, pub command_timeouts: u64, pub command_count: u64, pub command_latency_total_ms: u64, pub command_latency_max_ms: u64, pub command_latency_buckets: Vec, } #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)] pub struct RedisCommandLatencyBucket { pub le_ms: Option, pub count: u64, } #[derive(Default)] struct RedisConnectionMetrics { fast: RedisLaneMetrics, stream: RedisLaneMetrics, blocking_stream: RedisLaneMetrics, admin: RedisLaneMetrics, } impl RedisConnectionMetrics { fn for_lane(&self, lane: RedisConnectionLane) -> &RedisLaneMetrics { match lane { RedisConnectionLane::Fast => &self.fast, RedisConnectionLane::Stream => &self.stream, RedisConnectionLane::BlockingStream => &self.blocking_stream, RedisConnectionLane::Admin => &self.admin, } } } struct RedisLaneMetrics { errors: AtomicU64, timeouts: AtomicU64, command_count: AtomicU64, latency_total_ms: AtomicU64, latency_max_ms: AtomicU64, latency_bucket_counts: [AtomicU64; REDIS_COMMAND_LATENCY_BUCKET_COUNT], } impl Default for RedisLaneMetrics { fn default() -> Self { Self { errors: AtomicU64::new(0), timeouts: AtomicU64::new(0), command_count: AtomicU64::new(0), latency_total_ms: AtomicU64::new(0), latency_max_ms: AtomicU64::new(0), latency_bucket_counts: std::array::from_fn(|_| AtomicU64::new(0)), } } } impl RedisLaneMetrics { fn record_latency(&self, elapsed: Duration) { let elapsed_ms = u64::try_from(elapsed.as_millis()).unwrap_or(u64::MAX); self.command_count.fetch_add(1, Ordering::Relaxed); self.latency_total_ms .fetch_add(elapsed_ms, Ordering::Relaxed); update_atomic_max(&self.latency_max_ms, elapsed_ms); let bucket_index = REDIS_COMMAND_LATENCY_BUCKETS_MS .iter() .position(|upper_bound_ms| elapsed_ms <= *upper_bound_ms) .unwrap_or(REDIS_COMMAND_LATENCY_BUCKETS_MS.len()); self.latency_bucket_counts[bucket_index].fetch_add(1, Ordering::Relaxed); } fn latency_buckets(&self) -> Vec { let mut cumulative = 0u64; let mut buckets = Vec::with_capacity(REDIS_COMMAND_LATENCY_BUCKET_COUNT); for (index, upper_bound_ms) in REDIS_COMMAND_LATENCY_BUCKETS_MS.iter().enumerate() { cumulative = cumulative .saturating_add(self.latency_bucket_counts[index].load(Ordering::Relaxed)); buckets.push(RedisCommandLatencyBucket { le_ms: Some(*upper_bound_ms), count: cumulative, }); } cumulative = cumulative.saturating_add( self.latency_bucket_counts[REDIS_COMMAND_LATENCY_BUCKETS_MS.len()] .load(Ordering::Relaxed), ); buckets.push(RedisCommandLatencyBucket { le_ms: None, count: cumulative, }); buckets } } fn update_atomic_max(target: &AtomicU64, value: u64) { let mut current = target.load(Ordering::Relaxed); while value > current { match target.compare_exchange_weak(current, value, Ordering::Relaxed, Ordering::Relaxed) { Ok(_) => break, Err(next) => current = next, } } } fn connection_manager_config( command_timeout_ms: Option, ) -> redis::aio::ConnectionManagerConfig { let mut config = redis::aio::ConnectionManagerConfig::new(); if let Some(timeout_ms) = command_timeout_ms { config = config.set_connection_timeout(Duration::from_millis(timeout_ms)); } config } async fn connect_stream_lanes( client: &RedisClient, command_timeout_ms: Option, ) -> Result, DataLayerError> { let lane_count = stream_lane_count(); let mut lanes = Vec::with_capacity(lane_count); for _ in 0..lane_count { lanes.push( connect_lane( client, connection_manager_config(command_timeout_ms), RedisConnectionLane::Stream, command_timeout_ms, ) .await?, ); } Ok(lanes) } const fn stream_lane_count() -> usize { DEFAULT_STREAM_LANES } fn next_lane_index(next: &AtomicUsize, lane_count: usize) -> usize { debug_assert!(lane_count > 0, "redis connection lane must not be empty"); next.fetch_add(1, Ordering::Relaxed) % lane_count } async fn connect_blocking_stream_lanes( client: &RedisClient, command_timeout_ms: Option, requested_lanes: Option, ) -> Result { let lane_count = blocking_stream_lane_count(requested_lanes)?; let mut lanes = Vec::with_capacity(lane_count); for _ in 0..lane_count { lanes.push(Some( connect_blocking_stream_lane(client, command_timeout_ms).await?, )); } Ok(RedisBlockingStreamPool::new( client.clone(), command_timeout_ms, lanes, )) } async fn connect_blocking_stream_lane( client: &RedisClient, command_timeout_ms: Option, ) -> Result { let connect = client.create_multiplexed_tokio_connection(); let result = if let Some(timeout_ms) = command_timeout_ms { tokio::time::timeout(Duration::from_millis(timeout_ms), connect) .await .map_err(|_| { DataLayerError::TimedOut(format!( "runtime redis blocking_stream lane connection exceeded {timeout_ms}ms timeout" )) })? } else { connect.await }; let (connection, driver) = result.map_err(|err| { DataLayerError::Redis(format!( "failed to initialize runtime redis blocking_stream lane: {err}" )) })?; Ok(RedisBlockingConnection { connection, driver: Box::pin(driver), }) } fn blocking_stream_lane_count(requested_lanes: Option) -> Result { if matches!(requested_lanes, Some(0)) { return Err(DataLayerError::InvalidConfiguration( "runtime redis blocking_stream_lanes must be positive".to_string(), )); } let default_lanes = default_blocking_stream_lane_count(); Ok(requested_lanes .map(|lanes| lanes.max(default_lanes)) .unwrap_or(default_lanes) .clamp(1, MAX_BLOCKING_STREAM_LANES_CAP)) } fn default_blocking_stream_lane_count() -> usize { std::thread::available_parallelism() .map(|value| value.get()) .unwrap_or(DEFAULT_BLOCKING_STREAM_LANES_FALLBACK) .clamp( DEFAULT_BLOCKING_STREAM_LANES_FALLBACK, DEFAULT_BLOCKING_STREAM_LANES_CAP, ) } async fn connect_lane( client: &RedisClient, config: redis::aio::ConnectionManagerConfig, lane: RedisConnectionLane, command_timeout_ms: Option, ) -> Result { let connect = client.get_connection_manager_with_config(config); let result = if let Some(timeout_ms) = command_timeout_ms { match tokio::time::timeout(Duration::from_millis(timeout_ms), connect).await { Ok(result) => result, Err(_) => { return Err(DataLayerError::TimedOut(format!( "runtime redis {} lane connection exceeded {}ms timeout", lane.as_str(), timeout_ms ))); } } } else { connect.await }; result.map_err(|err| { DataLayerError::Redis(format!( "failed to initialize runtime redis {} lane: {err}", lane.as_str() )) }) } #[cfg(test)] mod tests { use super::{ blocking_stream_lane_count, default_blocking_stream_lane_count, next_lane_index, stream_lane_count, RedisBlockingStreamPool, RedisClientConfig, RedisClientFactory, RedisLaneMetrics, DEFAULT_STREAM_LANES, MAX_BLOCKING_STREAM_LANES_CAP, REDIS_COMMAND_LATENCY_BUCKETS_MS, }; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use std::time::Duration; #[test] fn factory_builds_lazy_client_from_valid_config() { let config = RedisClientConfig { url: "redis://127.0.0.1/0".to_string(), key_prefix: Some("aether".to_string()), }; let factory = RedisClientFactory::new(config.clone()).expect("factory should build"); assert_eq!(factory.config(), &config); let _client = factory .connect_lazy() .expect("lazy redis client should build"); } #[test] fn redis_config_debug_redacts_url_credentials_and_query() { let config = RedisClientConfig { url: "redis://redis-user:redis-password@redis.example/0?token=redis-secret".into(), key_prefix: Some("tenant-secret".into()), }; let debug = format!("{config:?}"); for secret in [ "redis-user", "redis-password", "redis-secret", "tenant-secret", ] { assert!(!debug.contains(secret), "debug leaked {secret}: {debug}"); } assert!(debug.contains("redis.example")); assert!(debug.contains("key_prefix_len")); } #[test] fn blocking_stream_lane_count_uses_requested_as_floor() { let default_lanes = default_blocking_stream_lane_count(); assert_eq!( blocking_stream_lane_count(None).expect("default lanes"), default_lanes ); assert_eq!( blocking_stream_lane_count(Some(1)).expect("requested below default"), default_lanes ); assert_eq!( blocking_stream_lane_count(Some(default_lanes + 1)).expect("requested above default"), default_lanes + 1 ); assert_eq!( blocking_stream_lane_count(Some(MAX_BLOCKING_STREAM_LANES_CAP + 1)) .expect("requested above cap"), MAX_BLOCKING_STREAM_LANES_CAP ); assert!(blocking_stream_lane_count(Some(0)).is_err()); } fn empty_blocking_pool(capacity: usize) -> Arc { Arc::new(RedisBlockingStreamPool::new( redis::Client::open("redis://127.0.0.1/0").expect("lazy client"), None, (0..capacity).map(|_| None).collect(), )) } #[tokio::test] async fn blocking_stream_pool_cancelled_checkout_preserves_owner_and_capacity() { use std::future::Future; use std::task::Poll; let pool = empty_blocking_pool(1); let owner = pool.checkout().await.expect("first lease"); let mut waiting = Box::pin(pool.checkout()); std::future::poll_fn(|cx| { assert!(matches!(waiting.as_mut().poll(cx), Poll::Pending)); Poll::Ready(()) }) .await; drop(waiting); assert_eq!(pool.permits.available_permits(), 0); assert!(pool.available.lock().unwrap().is_empty()); drop(owner); assert_eq!(pool.permits.available_permits(), 1); let replacement = pool .checkout() .await .expect("cancelled owner slot restored"); assert!(replacement.connection.is_none()); let panic = tokio::spawn(async move { let _lease = replacement; panic!("test owner panic"); }); assert!(panic.await.expect_err("owner panicked").is_panic()); assert_eq!(pool.permits.available_permits(), 1); assert_eq!(pool.available.lock().unwrap().len(), 1); } #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn blocking_stream_pool_concurrent_checkouts_stay_within_capacity() { let pool = empty_blocking_pool(3); let active = Arc::new(AtomicUsize::new(0)); let peak = Arc::new(AtomicUsize::new(0)); let barrier = Arc::new(tokio::sync::Barrier::new(32)); let mut tasks = tokio::task::JoinSet::new(); for _ in 0..32 { let pool = Arc::clone(&pool); let active = Arc::clone(&active); let peak = Arc::clone(&peak); let barrier = Arc::clone(&barrier); tasks.spawn(async move { barrier.wait().await; for _ in 0..32 { let lease = pool.checkout().await.expect("bounded lease"); let concurrent = active.fetch_add(1, Ordering::SeqCst) + 1; peak.fetch_max(concurrent, Ordering::SeqCst); assert!(concurrent <= 3); tokio::task::yield_now().await; active.fetch_sub(1, Ordering::SeqCst); drop(lease); } }); } tokio::time::timeout(Duration::from_secs(5), async { while let Some(result) = tasks.join_next().await { result.expect("checkout task"); } }) .await .expect("all checkouts complete without losing capacity"); assert!(peak.load(Ordering::SeqCst) <= 3); assert_eq!(active.load(Ordering::SeqCst), 0); assert_eq!(pool.permits.available_permits(), 3); assert_eq!(pool.available.lock().unwrap().len(), 3); } #[test] fn stream_lane_count_uses_fixed_default() { assert_eq!(stream_lane_count(), DEFAULT_STREAM_LANES); assert_eq!(stream_lane_count(), 4); } #[test] fn lane_index_round_robins_across_all_connections() { let next = AtomicUsize::new(0); let indexes = (0..10) .map(|_| next_lane_index(&next, stream_lane_count())) .collect::>(); assert_eq!(indexes, vec![0, 1, 2, 3, 0, 1, 2, 3, 0, 1]); } #[test] fn lane_index_round_robin_survives_counter_wraparound() { let next = AtomicUsize::new(usize::MAX - 1); let indexes = (0..3) .map(|_| next_lane_index(&next, stream_lane_count())) .collect::>(); assert_eq!(indexes, vec![2, 3, 0]); } #[test] fn lane_metrics_record_cumulative_latency_buckets() { let metrics = RedisLaneMetrics::default(); metrics.record_latency(Duration::from_millis(0)); metrics.record_latency(Duration::from_millis(12)); metrics.record_latency(Duration::from_millis(12_345)); assert_eq!(metrics.command_count.load(Ordering::Relaxed), 3); assert_eq!(metrics.latency_total_ms.load(Ordering::Relaxed), 12_357); assert_eq!(metrics.latency_max_ms.load(Ordering::Relaxed), 12_345); let buckets = metrics.latency_buckets(); let le_1 = buckets .iter() .find(|bucket| bucket.le_ms == Some(1)) .expect("1ms bucket"); let le_25 = buckets .iter() .find(|bucket| bucket.le_ms == Some(25)) .expect("25ms bucket"); let plus_inf = buckets.last().expect("+Inf bucket"); assert_eq!(buckets.len(), REDIS_COMMAND_LATENCY_BUCKETS_MS.len() + 1); assert_eq!(le_1.count, 1); assert_eq!(le_25.count, 2); assert_eq!(plus_inf.le_ms, None); assert_eq!(plus_inf.count, 3); } }