rustlavel_cache/redis/
pool.rs

1//! A small connection pool.
2//!
3//! A cache read is supposed to be cheaper than the work it avoids, and a TCP
4//! handshake plus `AUTH` plus `SELECT` per read is not. The pool keeps a few
5//! connections alive and hands them out, discarding any whose framing may have
6//! drifted — see [`super::connection::Connection::is_broken`].
7
8use super::config::RedisConfig;
9use super::connection::Connection;
10use rustlavel_core::{Error, Result};
11use std::collections::VecDeque;
12use std::sync::Arc;
13use tokio::sync::{Mutex, Semaphore};
14
15struct Inner {
16    config: RedisConfig,
17    idle: Mutex<VecDeque<Connection>>,
18    /// Bounds how many connections exist at once, including those in use, so a
19    /// traffic spike cannot open a thousand sockets against the cache.
20    permits: Arc<Semaphore>,
21}
22
23#[derive(Clone)]
24pub struct Pool {
25    inner: Arc<Inner>,
26}
27
28impl Pool {
29    /// Create a pool. Nothing is connected until the first command, so an
30    /// application still boots when Redis is briefly unavailable — a cache
31    /// being down should degrade a service, not stop it starting.
32    pub fn new(config: RedisConfig) -> Self {
33        let permits = Arc::new(Semaphore::new(config.max_connections.max(1)));
34        Pool {
35            inner: Arc::new(Inner { config, idle: Mutex::new(VecDeque::new()), permits }),
36        }
37    }
38
39    pub fn config(&self) -> &RedisConfig {
40        &self.inner.config
41    }
42
43    /// Open one connection immediately, so a misconfiguration is reported at
44    /// boot rather than on the first cache miss in production.
45    pub async fn verify(&self) -> Result<()> {
46        let mut connection = self.acquire().await?;
47        connection.command(&[b"PING"]).await?.into_result()?;
48        Ok(())
49    }
50
51    pub async fn acquire(&self) -> Result<PooledConnection> {
52        let permit = Arc::clone(&self.inner.permits)
53            .acquire_owned()
54            .await
55            .map_err(|_| Error::msg("the Redis pool has been closed"))?;
56
57        if let Some(connection) = self.inner.idle.lock().await.pop_front() {
58            return Ok(PooledConnection {
59                connection: Some(connection),
60                pool: Arc::clone(&self.inner),
61                _permit: permit,
62            });
63        }
64
65        let connection = Connection::connect(&self.inner.config).await?;
66        Ok(PooledConnection {
67            connection: Some(connection),
68            pool: Arc::clone(&self.inner),
69            _permit: permit,
70        })
71    }
72
73    /// Run one command on a borrowed connection.
74    pub async fn command(&self, args: &[&[u8]]) -> Result<super::resp::Value> {
75        let mut connection = self.acquire().await?;
76        connection.command(args).await
77    }
78
79    /// How many connections are idle. For tests and diagnostics.
80    pub async fn idle_count(&self) -> usize {
81        self.inner.idle.lock().await.len()
82    }
83
84    pub async fn close(&self) {
85        let mut idle = self.inner.idle.lock().await;
86        while let Some(connection) = idle.pop_front() {
87            connection.close().await;
88        }
89    }
90}
91
92/// A connection borrowed from the pool, returned when dropped.
93pub struct PooledConnection {
94    connection: Option<Connection>,
95    pool: Arc<Inner>,
96    /// Held for the lifetime of the borrow; releasing it lets another caller in.
97    _permit: tokio::sync::OwnedSemaphorePermit,
98}
99
100impl std::ops::Deref for PooledConnection {
101    type Target = Connection;
102
103    fn deref(&self) -> &Connection {
104        self.connection.as_ref().expect("connection is present until drop")
105    }
106}
107
108impl std::ops::DerefMut for PooledConnection {
109    fn deref_mut(&mut self) -> &mut Connection {
110        self.connection.as_mut().expect("connection is present until drop")
111    }
112}
113
114impl Drop for PooledConnection {
115    fn drop(&mut self) {
116        let Some(connection) = self.connection.take() else { return };
117
118        if connection.is_broken() {
119            // Returning this would hand the next borrower somebody else's
120            // reply. Closing it costs one socket; reusing it costs correctness.
121            tokio::spawn(async move { connection.close().await });
122            return;
123        }
124
125        let pool = Arc::clone(&self.pool);
126        tokio::spawn(async move {
127            pool.idle.lock().await.push_back(connection);
128        });
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135    use std::time::Duration;
136
137    #[tokio::test]
138    async fn a_pool_opens_nothing_until_it_is_used() {
139        let pool = Pool::new(RedisConfig { port: 1, ..RedisConfig::default() });
140        assert_eq!(pool.idle_count().await, 0);
141    }
142
143    #[tokio::test]
144    async fn acquiring_reports_a_connection_failure_rather_than_hanging() {
145        let pool = Pool::new(RedisConfig {
146            port: 1,
147            connect_timeout: Duration::from_secs(2),
148            ..RedisConfig::default()
149        });
150        assert!(pool.acquire().await.is_err());
151        assert!(pool.verify().await.is_err());
152    }
153}
rustlavel_cache/redis/pool.rs

rustlavel_cache/redis/
pool.rs
