Aether/crates/aether-runtime/state/src/redis/client.rs

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<String>,
}

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::<String>()
        )
    }
}

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<Self, DataLayerError> {
        config.validate()?;
        Ok(Self { config })
    }

    pub(crate) fn config(&self) -> &RedisClientConfig {
        &self.config
    }

    pub(crate) fn connect_lazy(&self) -> Result<RedisClient, DataLayerError> {
        RedisClient::open(self.config.url.clone()).map_redis_err()
    }

    pub(crate) async fn connect_router(
        &self,
        command_timeout_ms: Option<u64>,
    ) -> Result<RedisConnectionRouter, DataLayerError> {
        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<u64>,
        blocking_stream_lanes: Option<usize>,
    ) -> Result<RedisConnectionRouter, DataLayerError> {
        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<Vec<RedisManagedConnection>>,
    stream_next: Arc<AtomicUsize>,
    blocking_stream: Arc<RedisBlockingStreamPool>,
    usage_cleanup: Arc<RedisBlockingStreamPool>,
    admin: RedisManagedConnection,
    metrics: Arc<RedisConnectionMetrics>,
}

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<u64>,
        blocking_stream_lanes: Option<usize>,
    ) -> Result<Self, DataLayerError> {
        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<RedisBlockingStreamLease, DataLayerError> {
        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<RedisBlockingStreamLease, DataLayerError> {
        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<RedisLaneDiagnostics> {
        [
            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<Box<dyn Future<Output = ()> + Send>>,
}

impl RedisBlockingConnection {
    async fn query(&mut self, command: &redis::Cmd) -> Result<redis::Value, DataLayerError> {
        // 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<u64>,
    capacity: usize,
    available: Mutex<Vec<Option<RedisBlockingConnection>>>,
    permits: Arc<Semaphore>,
}

impl RedisBlockingStreamPool {
    fn new(
        client: RedisClient,
        command_timeout_ms: Option<u64>,
        available: Vec<Option<RedisBlockingConnection>>,
    ) -> 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<Self>) -> Result<RedisBlockingStreamLease, DataLayerError> {
        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<RedisBlockingStreamPool>,
    connection: Option<RedisBlockingConnection>,
    reusable: bool,
    _permit: OwnedSemaphorePermit,
}

impl RedisBlockingStreamLease {
    pub(crate) async fn query(
        &mut self,
        command: &redis::Cmd,
    ) -> Result<redis::Value, DataLayerError> {
        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<RedisCommandLatencyBucket>,
}

#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub struct RedisCommandLatencyBucket {
    pub le_ms: Option<u64>,
    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<RedisCommandLatencyBucket> {
        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<u64>,
) -> 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<u64>,
) -> Result<Vec<RedisManagedConnection>, 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<u64>,
    requested_lanes: Option<usize>,
) -> Result<RedisBlockingStreamPool, DataLayerError> {
    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<u64>,
) -> Result<RedisBlockingConnection, DataLayerError> {
    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<usize>) -> Result<usize, DataLayerError> {
    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<u64>,
) -> Result<RedisManagedConnection, DataLayerError> {
    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:[email protected]/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<RedisBlockingStreamPool> {
        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::<Vec<_>>();

        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::<Vec<_>>();

        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);
    }
}