Skip to main content

tank_core/
pool.rs

1use crate::{Connection, Driver, Error, Result};
2use deadpool::managed::{Manager, Metrics, Object, Pool, RecycleResult, Timeouts};
3use std::{borrow::Cow, fmt::Debug, future, time::Duration};
4
5#[derive(Debug)]
6pub struct DBConnectionManager<D: Driver> {
7    driver: D,
8    url: Cow<'static, str>,
9}
10
11impl<D: Driver> DBConnectionManager<D> {
12    pub fn new(driver: D, url: Cow<'static, str>) -> Self {
13        Self { driver, url }
14    }
15}
16
17impl<D: Driver> Manager for DBConnectionManager<D> {
18    type Type = D::Connection;
19    type Error = Error;
20    async fn create(&self) -> Result<Self::Type> {
21        Ok(D::Connection::connect(&self.driver, self.url.clone()).await?)
22    }
23    fn recycle(
24        &self,
25        _: &mut Self::Type,
26        _: &Metrics,
27    ) -> impl Future<Output = RecycleResult<Self::Error>> + Send {
28        future::ready(RecycleResult::Ok(()))
29    }
30}
31
32pub trait ConnectionPool<D: Driver>: Debug {
33    fn get(
34        &self,
35    ) -> impl Future<Output = Result<impl AsRef<D::Connection> + AsMut<D::Connection>>> + Send;
36    fn timeout_get(
37        &self,
38        timeout: Duration,
39    ) -> impl Future<Output = Result<impl AsRef<D::Connection> + AsMut<D::Connection>>> + Send;
40    fn detach(&self) -> impl Future<Output = Result<D::Connection>> + Send
41    where
42        Self: Sized;
43    fn resize(&self, max_size: usize) -> Result<()>;
44    fn close(self) -> impl Future<Output = Result<()>> + Send;
45}
46
47impl<D: Driver> ConnectionPool<D> for Pool<DBConnectionManager<D>>
48where
49    <D as Driver>::Connection: Debug,
50{
51    async fn get(&self) -> Result<impl AsRef<D::Connection> + AsMut<D::Connection>> {
52        Ok(Pool::<DBConnectionManager<D>>::get(self)
53            .await
54            .map_err(|e| Error::msg(format!("{e:#?}")))?)
55    }
56
57    async fn timeout_get(
58        &self,
59        timeout: Duration,
60    ) -> Result<impl AsRef<D::Connection> + AsMut<D::Connection>> {
61        Ok(Pool::<DBConnectionManager<D>>::timeout_get(
62            self,
63            &Timeouts::wait_millis(timeout.as_millis() as u64),
64        )
65        .await
66        .map_err(|e| Error::msg(format!("{e:#?}")))?)
67    }
68
69    async fn detach(&self) -> Result<D::Connection>
70    where
71        Self: Sized,
72    {
73        let v = Pool::<DBConnectionManager<D>>::get(self)
74            .await
75            .map_err(|e| Error::msg(format!("{e:#?}")))?;
76        Ok(Object::<DBConnectionManager<D>>::take(v))
77    }
78
79    fn resize(&self, max_size: usize) -> Result<()> {
80        Ok(self.resize(max_size))
81    }
82
83    fn close(self) -> impl Future<Output = Result<()>> + Send {
84        Self::close(&self);
85        future::ready(Ok(()))
86    }
87}
88
89#[derive(Clone, Copy)]
90pub enum QueueMode {
91    Fifo,
92    Lifo,
93}
94
95#[derive(Clone, Copy)]
96pub struct PoolConfig {
97    pub max_size: usize,
98    pub wait_timeout: Option<Duration>,
99    pub create_timeout: Option<Duration>,
100    pub recycle_timeout: Option<Duration>,
101    pub queue_mode: QueueMode,
102}
103
104impl PoolConfig {
105    pub fn new() -> Self {
106        deadpool::managed::PoolConfig::default().into()
107    }
108}
109
110impl From<PoolConfig> for deadpool::managed::PoolConfig {
111    fn from(value: PoolConfig) -> Self {
112        Self {
113            max_size: value.max_size,
114            timeouts: deadpool::managed::Timeouts {
115                wait: value.wait_timeout,
116                create: value.create_timeout,
117                recycle: value.recycle_timeout,
118            },
119            queue_mode: match value.queue_mode {
120                QueueMode::Fifo => deadpool::managed::QueueMode::Fifo,
121                QueueMode::Lifo => deadpool::managed::QueueMode::Lifo,
122            },
123        }
124    }
125}
126
127impl From<deadpool::managed::PoolConfig> for PoolConfig {
128    fn from(value: deadpool::managed::PoolConfig) -> Self {
129        Self {
130            max_size: value.max_size,
131            wait_timeout: value.timeouts.wait,
132            create_timeout: value.timeouts.create,
133            recycle_timeout: value.timeouts.recycle,
134            queue_mode: match value.queue_mode {
135                deadpool::managed::QueueMode::Fifo => QueueMode::Fifo,
136                deadpool::managed::QueueMode::Lifo => QueueMode::Lifo,
137            },
138        }
139    }
140}
141
142impl Default for PoolConfig {
143    fn default() -> Self {
144        Self::new()
145    }
146}