connpool/pools/
registry.rs

1//! `PoolRegistry`: one [`ConnectionPool`] per key (per host, tenant, database…), created on
2//! demand.
3
4use std::collections::HashMap;
5use std::fmt::Display;
6use std::hash::Hash;
7use std::sync::{Arc, Mutex};
8
9use crate::core::clock::Clock;
10use crate::core::errors::PoolError;
11use crate::core::timeouts::Timeout;
12use crate::factories::sync_factory::ConnectionFactory;
13use crate::observability::events::PoolListener;
14use crate::policy::config::PoolConfig;
15use crate::pools::sync_pool::{ConnectionPool, PoolOptions};
16
17type FactoryFor<K, T> = Box<dyn Fn(&K) -> Box<dyn ConnectionFactory<T>> + Send + Sync>;
18
19struct RegState<K, T> {
20    pools: HashMap<K, ConnectionPool<T>>,
21    closed: bool,
22}
23
24/// Hands out one [`ConnectionPool`] per key, built on first use. It is **lazy**: a pool is
25/// created on first access and is NOT opened, so it does no background maintenance unless you
26/// open it yourself.
27pub struct PoolRegistry<K, T> {
28    factory_for: FactoryFor<K, T>,
29    config: PoolConfig,
30    clock: Option<Arc<dyn Clock>>,
31    listeners: Vec<Arc<dyn PoolListener<T>>>,
32    inner: Mutex<RegState<K, T>>,
33}
34
35impl<K, T> PoolRegistry<K, T>
36where
37    K: Clone + Eq + Hash + Display,
38    T: Clone + Eq + Hash + Send + Sync + 'static,
39{
40    /// Builds a registry. `factory_for` makes the factory for a given key; `config` is shared by
41    /// every pool.
42    pub fn new<F>(factory_for: F, config: PoolConfig) -> Self
43    where
44        F: Fn(&K) -> Box<dyn ConnectionFactory<T>> + Send + Sync + 'static,
45    {
46        Self {
47            factory_for: Box::new(factory_for),
48            config,
49            clock: None,
50            listeners: Vec::new(),
51            inner: Mutex::new(RegState {
52                pools: HashMap::new(),
53                closed: false,
54            }),
55        }
56    }
57
58    /// Shares `clock` with every pool the registry creates.
59    pub fn with_clock(mut self, clock: Arc<dyn Clock>) -> Self {
60        self.clock = Some(clock);
61        self
62    }
63
64    /// Shares `listeners` with every pool the registry creates.
65    pub fn with_listeners(mut self, listeners: Vec<Arc<dyn PoolListener<T>>>) -> Self {
66        self.listeners = listeners;
67        self
68    }
69
70    /// Returns the pool for `key`, creating it on first use. The returned pool is a cheap handle
71    /// to the registry-owned pool.
72    pub fn pool(&self, key: &K) -> Result<ConnectionPool<T>, PoolError> {
73        let mut g = self.inner.lock().unwrap();
74        if g.closed {
75            return Err(PoolError::Closed("registry is closed".to_string()));
76        }
77        if let Some(pool) = g.pools.get(key) {
78            return Ok(pool.clone());
79        }
80        let options = PoolOptions {
81            name: Some(key.to_string()),
82            clock: self.clock.clone(),
83            rng: None,
84            listeners: self.listeners.clone(),
85        };
86        let pool =
87            ConnectionPool::new((self.factory_for)(key), self.config.clone(), Some(options))?;
88        g.pools.insert(key.clone(), pool.clone());
89        Ok(pool)
90    }
91
92    /// Borrows a connection from the pool for `key` for the duration of `f`.
93    pub fn with_connection<R>(
94        &self,
95        key: &K,
96        timeout: Timeout,
97        f: impl FnOnce(&T) -> R,
98    ) -> Result<R, PoolError> {
99        self.pool(key)?.with_connection(timeout, f)
100    }
101
102    /// Closes every pool and rejects further use.
103    pub fn close(&self) {
104        let pools = {
105            let mut g = self.inner.lock().unwrap();
106            g.closed = true;
107            std::mem::take(&mut g.pools)
108        };
109        for (_key, pool) in pools {
110            pool.close(None);
111        }
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118    use crate::test_support::*;
119
120    fn registry(h: &Harness) -> PoolRegistry<String, Conn> {
121        let fac = h.fac.clone();
122        let rec: Arc<dyn PoolListener<Conn>> = Recorder::new();
123        PoolRegistry::new(
124            move |_key: &String| {
125                Box::new(TestFactory(fac.clone())) as Box<dyn ConnectionFactory<Conn>>
126            },
127            PoolConfig {
128                maintenance_interval: None,
129                breaker: None,
130                ..Default::default()
131            },
132        )
133        .with_clock(h.clock.clone())
134        .with_listeners(vec![rec])
135    }
136
137    #[test]
138    fn pool_is_created_once_per_key() {
139        let h = Harness::new();
140        let reg = registry(&h);
141        let p1 = reg.pool(&"a".to_string()).unwrap();
142        let p2 = reg.pool(&"a".to_string()).unwrap();
143        assert_eq!(p1.name(), "a");
144        assert_eq!(p2.name(), "a");
145        let _b = reg.pool(&"b".to_string()).unwrap();
146        // with_connection routes through the per-key pool.
147        let id = reg
148            .with_connection(&"a".to_string(), Timeout::Unset, |c| c.id)
149            .unwrap();
150        assert!(id >= 1);
151    }
152
153    #[test]
154    fn close_rejects_further_use() {
155        let h = Harness::new();
156        let reg = registry(&h);
157        let p = reg.pool(&"a".to_string()).unwrap();
158        reg.close();
159        assert!(p.is_closed());
160        assert!(matches!(
161            reg.pool(&"a".to_string()),
162            Err(PoolError::Closed(_))
163        ));
164    }
165}
connpool/pools/registry.rs

connpool/pools/
registry.rs
