Skip to main content

ringline_ping/
pool.rs

1//! Connection pool for ringline-ping.
2//!
3//! `Pool` manages a fixed set of backend connections with round-robin dispatch
4//! and lazy reconnection. It is single-threaded (no Arc, no Mutex) — designed
5//! for use within a single ringline async worker.
6//!
7//! # Usage
8//!
9//! ```no_run
10//! # use std::net::SocketAddr;
11//! # use ringline_ping::{Pool, PoolConfig};
12//! # async fn example() -> Result<(), ringline_ping::Error> {
13//! let config = PoolConfig {
14//!     addr: "127.0.0.1:6379".parse().unwrap(),
15//!     pool_size: 4,
16//!     connect_timeout_ms: 0,
17//!     tls_server_name: None,
18//! };
19//! let mut pool = Pool::new(config);
20//! pool.connect_all().await?;
21//! pool.client().await?.ping().await?;
22//! # Ok(())
23//! # }
24//! ```
25
26use std::net::SocketAddr;
27
28use ringline::ConnCtx;
29
30use crate::{Client, Error};
31
32/// Configuration for a connection pool.
33pub struct PoolConfig {
34    /// Backend address to connect to.
35    pub addr: SocketAddr,
36    /// Number of connections in the pool.
37    pub pool_size: usize,
38    /// Connect timeout in milliseconds. 0 means no timeout.
39    pub connect_timeout_ms: u64,
40    /// TLS server name (SNI) for outbound connections. `None` means plain TCP.
41    pub tls_server_name: Option<String>,
42}
43
44enum Slot {
45    Connected(ConnCtx),
46    Disconnected,
47}
48
49/// A fixed-size connection pool with round-robin dispatch.
50///
51/// All slots start disconnected. Call [`connect_all()`](Pool::connect_all) for
52/// eager startup, or let [`client()`](Pool::client) lazily reconnect on demand.
53pub struct Pool {
54    addr: SocketAddr,
55    slots: Vec<Slot>,
56    next: usize,
57    connect_timeout_ms: u64,
58    tls_server_name: Option<String>,
59}
60
61impl Pool {
62    /// Create a new pool. All slots start disconnected.
63    pub fn new(config: PoolConfig) -> Self {
64        let mut slots = Vec::with_capacity(config.pool_size);
65        for _ in 0..config.pool_size {
66            slots.push(Slot::Disconnected);
67        }
68        Pool {
69            addr: config.addr,
70            slots,
71            next: 0,
72            connect_timeout_ms: config.connect_timeout_ms,
73            tls_server_name: config.tls_server_name,
74        }
75    }
76
77    /// Eagerly connect all slots. Returns an error if any connection fails.
78    pub async fn connect_all(&mut self) -> Result<(), Error> {
79        for i in 0..self.slots.len() {
80            let conn = self.do_connect().await?;
81            self.slots[i] = Slot::Connected(conn);
82        }
83        Ok(())
84    }
85
86    /// Get a [`Client`] bound to the next healthy connection.
87    ///
88    /// Advances the round-robin cursor and returns a client for a connected
89    /// slot. Disconnected slots are lazily reconnected. If all slots fail,
90    /// returns [`Error::AllConnectionsFailed`].
91    pub async fn client(&mut self) -> Result<Client, Error> {
92        let size = self.slots.len();
93        for _ in 0..size {
94            let idx = self.next;
95            self.next = (self.next + 1) % size;
96
97            match &self.slots[idx] {
98                Slot::Connected(conn) => return Ok(Client::new(*conn)),
99                Slot::Disconnected => {
100                    if let Ok(conn) = self.do_connect().await {
101                        self.slots[idx] = Slot::Connected(conn);
102                        return Ok(Client::new(conn));
103                    }
104                }
105            }
106        }
107        Err(Error::AllConnectionsFailed)
108    }
109
110    /// Mark a connection as dead after a `ConnectionClosed` error.
111    ///
112    /// Matches by [`ConnCtx::token()`] and sets the slot to disconnected.
113    /// The next [`client()`](Pool::client) call will lazily reconnect.
114    pub fn mark_disconnected(&mut self, conn: ConnCtx) {
115        let token = conn.token();
116        for slot in &mut self.slots {
117            if let Slot::Connected(conn) = slot
118                && conn.token() == token
119            {
120                conn.close();
121                *slot = Slot::Disconnected;
122                return;
123            }
124        }
125    }
126
127    /// Close all connections and reset slots to disconnected.
128    pub fn close_all(&mut self) {
129        for slot in &mut self.slots {
130            if let Slot::Connected(conn) = slot {
131                conn.close();
132            }
133            *slot = Slot::Disconnected;
134        }
135    }
136
137    /// Number of currently connected slots.
138    pub fn connected_count(&self) -> usize {
139        self.slots
140            .iter()
141            .filter(|s| matches!(s, Slot::Connected(_)))
142            .count()
143    }
144
145    /// Total number of slots in the pool.
146    pub fn pool_size(&self) -> usize {
147        self.slots.len()
148    }
149
150    async fn do_connect(&self) -> Result<ConnCtx, Error> {
151        let conn = if let Some(sni) = &self.tls_server_name {
152            let fut = if self.connect_timeout_ms > 0 {
153                ringline::connect_tls_with_timeout(self.addr, sni, self.connect_timeout_ms)?
154            } else {
155                ringline::connect_tls(self.addr, sni)?
156            };
157            fut.await?
158        } else {
159            let fut = if self.connect_timeout_ms > 0 {
160                ringline::connect_with_timeout(self.addr, self.connect_timeout_ms)?
161            } else {
162                ringline::connect(self.addr)?
163            };
164            fut.await?
165        };
166
167        Ok(conn)
168    }
169}