15280118615

// Copyright 2020. The Tari Project
//
// Redistribution and use in source and binary forms, with or without modification, are permitted provided that the
// following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following
// disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the
// following disclaimer in the documentation and/or other materials provided with the distribution.
//
// 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote
// products derived from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
// USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

use core::time::Duration;
use std::{convert::TryFrom, path::PathBuf};

use diesel::{
    r2d2::{ConnectionManager, Pool, PooledConnection},
    SqliteConnection,
};
use log::*;

use crate::{connection_options::ConnectionOptions, error::SqliteStorageError};

const LOG_TARGET: &str = "common_sqlite::sqlite_connection_pool";

#[derive(Clone)]
pub struct SqliteConnectionPool {
    pool: Option<Pool<ConnectionManager<SqliteConnection>>>,
    db_path: String,
    pool_size: usize,
    connection_options: ConnectionOptions,
}

impl SqliteConnectionPool {
    pub fn new(
        db_path: String,
        pool_size: usize,
        enable_wal: bool,
        enable_foreign_keys: bool,
        busy_timeout: Duration,
    ) -> Self {
        Self {
            pool: None,
            db_path,
            pool_size,
            connection_options: ConnectionOptions::new(enable_wal, enable_foreign_keys, busy_timeout),
        }
    }

    /// Create an sqlite connection pool managed by the pool connection manager
    pub fn create_pool(&mut self) -> Result<(), SqliteStorageError> {
        if self.pool.is_none() {
            let pool = Pool::builder()
                .max_size(u32::try_from(self.pool_size)?)
                .connection_customizer(Box::new(self.connection_options.clone()))
                .build(ConnectionManager::<SqliteConnection>::new(self.db_path.as_str()))
                .map_err(|e| SqliteStorageError::DieselR2d2Error(e.to_string()));
            self.pool = Some(pool?);
        } else {
            warn!(
                target: LOG_TARGET,
                "Connection pool for {} already exists", self.db_path
            );
        }
        Ok(())
    }

    /// Return a pooled sqlite connection managed by the pool connection manager, waits for at most the configured
    /// connection timeout before returning an error.
    pub fn get_pooled_connection(
        &self,
    ) -> Result<PooledConnection<ConnectionManager<SqliteConnection>>, SqliteStorageError> {
        if let Some(pool) = self.pool.as_ref() {
            pool.get().map_err(|e| {
                warn!(
                    target: LOG_TARGET,
                    "Connection pool state {:?}: {}",
                    pool.state(),
                    e.to_string()
                );
                SqliteStorageError::DieselR2d2Error(e.to_string())
            })
        } else {
            Err(SqliteStorageError::DieselR2d2Error("Pool does not exist".to_string()))
        }
    }

    /// Return a pooled sqlite connection managed by the pool connection manager, waits for at most supplied
    /// connection timeout before returning an error.
    pub fn get_pooled_connection_timeout(
        &self,
        timeout: Duration,
    ) -> Result<PooledConnection<ConnectionManager<SqliteConnection>>, SqliteStorageError> {
        if let Some(pool) = self.pool.clone() {
            pool.get_timeout(timeout).map_err(|e| {
                warn!(
                    target: LOG_TARGET,
                    "Connection pool state {:?}: {}",
                    pool.state(),
                    e.to_string()
                );
                SqliteStorageError::DieselR2d2Error(e.to_string())
            })
        } else {
            Err(SqliteStorageError::DieselR2d2Error("Pool does not exist".to_string()))
        }
    }

    /// Return a pooled sqlite connection managed by the pool connection manager, returns None if there are no idle
    /// connections available in the pool. This method will not block waiting to establish a new connection.
    pub fn try_get_pooled_connection(
        &self,
    ) -> Result<Option<PooledConnection<ConnectionManager<SqliteConnection>>>, SqliteStorageError> {
        if let Some(pool) = self.pool.clone() {
            let connection = pool.try_get();
            if connection.is_none() {
                warn!(
                    target: LOG_TARGET,
                    "No connections available, pool state {:?}",
                    pool.state()
                );
            };
            Ok(connection)
        } else {
            Err(SqliteStorageError::DieselR2d2Error("Pool does not exist".to_string()))
        }
    }

    /// Return the database path
    pub fn db_path(&self) -> PathBuf {
        PathBuf::from(&self.db_path)
    }

    /// Perform cleanup on the connection pool. This will drop the pool and return the state of the pool.
    pub fn cleanup(&mut self) -> Option<String> {
        if let Some(pool) = self.pool.take() {
            let state = format!("{:?}", pool.state());
            drop(pool);
            return Some(state);
        }
        None
    }
}

pub trait PooledDbConnection: Send + Sync + Clone {
    type Error;

    fn get_pooled_connection(&self) -> Result<PooledConnection<ConnectionManager<SqliteConnection>>, Self::Error>;
}

1	// Copyright 2020. The Tari Project
2	//
3	// Redistribution and use in source and binary forms, with or without modification, are permitted provided that the
4	// following conditions are met:
5	//
6	// 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following
7	// disclaimer.
8	//
9	// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the
10	// following disclaimer in the documentation and/or other materials provided with the distribution.
11	//
12	// 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote
13	// products derived from this software without specific prior written permission.
14	//
15	// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
16	// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
17	// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
18	// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
19	// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
20	// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
21	// USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
22
23	use core::time::Duration;
24	use std::{convert::TryFrom, path::PathBuf};
25
26	use diesel::{
27	r2d2::{ConnectionManager, Pool, PooledConnection},
28	SqliteConnection,
29	};
30	use log::*;
31
32	use crate::{connection_options::ConnectionOptions, error::SqliteStorageError};
33
34	const LOG_TARGET: &str = "common_sqlite::sqlite_connection_pool";
35
36	#[derive(Clone)]
37	pub struct SqliteConnectionPool {
38	pool: Option<Pool<ConnectionManager<SqliteConnection>>>,
39	db_path: String,
40	pool_size: usize,
41	connection_options: ConnectionOptions,
42	}
43
44	impl SqliteConnectionPool {
45	pub fn new(	848✔
46	db_path: String,	848✔
47	pool_size: usize,	848✔
48	enable_wal: bool,	848✔
49	enable_foreign_keys: bool,	848✔
50	busy_timeout: Duration,	848✔
51	) -> Self {	848✔
52	Self {	848✔
53	pool: None,	848✔
54	db_path,	848✔
55	pool_size,	848✔
56	connection_options: ConnectionOptions::new(enable_wal, enable_foreign_keys, busy_timeout),	848✔
57	}	848✔
58	}	848✔
59
60	/// Create an sqlite connection pool managed by the pool connection manager
61	pub fn create_pool(&mut self) -> Result<(), SqliteStorageError> {	848✔
62	if self.pool.is_none() {	848✔
63	let pool = Pool::builder()	848✔
64	.max_size(u32::try_from(self.pool_size)?)	848✔
65	.connection_customizer(Box::new(self.connection_options.clone()))	848✔
66	.build(ConnectionManager::<SqliteConnection>::new(self.db_path.as_str()))	848✔
67	.map_err(\|e\| SqliteStorageError::DieselR2d2Error(e.to_string()));	848✔
68	self.pool = Some(pool?);	848✔
69	} else {
70	warn!(	×
71	target: LOG_TARGET,	×
72	"Connection pool for {} already exists", self.db_path	×
73	);
74	}
75	Ok(())	848✔
76	}	848✔
77
78	/// Return a pooled sqlite connection managed by the pool connection manager, waits for at most the configured
79	/// connection timeout before returning an error.
80	pub fn get_pooled_connection(	3,739,573✔
81	&self,	3,739,573✔
82	) -> Result<PooledConnection<ConnectionManager<SqliteConnection>>, SqliteStorageError> {	3,739,573✔
83	if let Some(pool) = self.pool.as_ref() {	3,739,573✔
84	pool.get().map_err(\|e\| {	3,739,573✔
85	warn!(	381✔
86	target: LOG_TARGET,	×
87	"Connection pool state {:?}: {}",	×
88	pool.state(),	×
89	e.to_string()	×
90	);
91	SqliteStorageError::DieselR2d2Error(e.to_string())	381✔
92	})	3,739,573✔
93	} else {
94	Err(SqliteStorageError::DieselR2d2Error("Pool does not exist".to_string()))	×
95	}
96	}	3,739,573✔
97
98	/// Return a pooled sqlite connection managed by the pool connection manager, waits for at most supplied
99	/// connection timeout before returning an error.
100	pub fn get_pooled_connection_timeout(	×
101	&self,	×
102	timeout: Duration,	×
103	) -> Result<PooledConnection<ConnectionManager<SqliteConnection>>, SqliteStorageError> {	×
104	if let Some(pool) = self.pool.clone() {	×
105	pool.get_timeout(timeout).map_err(\|e\| {	×
106	warn!(	×
107	target: LOG_TARGET,	×
108	"Connection pool state {:?}: {}",	×
109	pool.state(),	×
110	e.to_string()	×
111	);
112	SqliteStorageError::DieselR2d2Error(e.to_string())	×
113	})	×
114	} else {
115	Err(SqliteStorageError::DieselR2d2Error("Pool does not exist".to_string()))	×
116	}
117	}	×
118
119	/// Return a pooled sqlite connection managed by the pool connection manager, returns None if there are no idle
120	/// connections available in the pool. This method will not block waiting to establish a new connection.
121	pub fn try_get_pooled_connection(	×
122	&self,	×
123	) -> Result<Option<PooledConnection<ConnectionManager<SqliteConnection>>>, SqliteStorageError> {	×
124	if let Some(pool) = self.pool.clone() {	×
125	let connection = pool.try_get();	×
126	if connection.is_none() {	×
127	warn!(	×
128	target: LOG_TARGET,	×
129	"No connections available, pool state {:?}",	×
130	pool.state()	×
131	);
132	};	×
133	Ok(connection)	×
134	} else {
135	Err(SqliteStorageError::DieselR2d2Error("Pool does not exist".to_string()))	×
136	}
137	}	×
138
139	/// Return the database path
140	pub fn db_path(&self) -> PathBuf {	920,245✔
141	PathBuf::from(&self.db_path)	920,245✔
142	}	920,245✔
143
144	/// Perform cleanup on the connection pool. This will drop the pool and return the state of the pool.
145	pub fn cleanup(&mut self) -> Option<String> {	184✔
146	if let Some(pool) = self.pool.take() {	184✔
147	let state = format!("{:?}", pool.state());	184✔
148	drop(pool);	184✔
149	return Some(state);	184✔
UNCOV 150	}	×
UNCOV 151	None	×
152	}	184✔
153	}
154
155	pub trait PooledDbConnection: Send + Sync + Clone {
156	type Error;
157
158	fn get_pooled_connection(&self) -> Result<PooledConnection<ConnectionManager<SqliteConnection>>, Self::Error>;
159	}

tari-project / tari / 15280118615

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous