ringline_memcache/pool.rs
1//! Connection pool for ringline-memcache.
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_memcache::{Pool, PoolConfig};
12//! # async fn example() -> Result<(), ringline_memcache::Error> {
13//! let config = PoolConfig {
14//! addr: "127.0.0.1:11211".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?.set("k", "v").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}