connpool/pools/
async_registry.rs

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

connpool/pools/
async_registry.rs
