#![allow(clippy::doc_markdown)]
#![cfg_attr(docsrs, feature(doc_cfg))]
mod changehook;
mod err;
mod rawhook;
mod wrconn;
pub mod autovacuum;
pub mod utils;
use std::{fmt, mem::ManuallyDrop, path::Path, str::FromStr, sync::Arc};
use parking_lot::{Condvar, Mutex, MutexGuard};
use r2d2::{CustomizeConnection, PooledConnection};
pub use {r2d2, r2d2_sqlite::SqliteConnectionManager, rusqlite};
use rusqlite::{Connection, OpenFlags, params};
#[cfg(feature = "tpool")]
use threadpool::ThreadPool;
pub use changehook::ChangeLogHook;
pub use err::Error;
pub use rawhook::{Action, Hook};
pub use wrconn::WrConn;
pub enum RegOn<F>
where
F: Fn(&Connection) -> Result<(), rusqlite::Error>
{
RO(F),
RW(F),
Both(F)
}
type RegCb = dyn Fn(&Connection) -> Result<(), rusqlite::Error> + Send + Sync;
enum CbType {
Ro(Box<RegCb>),
Rw(Box<RegCb>),
Both(Box<RegCb>)
}
pub trait SchemaMgr {
#[allow(unused_variables)]
fn init(&self, conn: &mut Connection, newdb: bool) -> Result<(), Error> {
Ok(())
}
#[allow(unused_variables)]
fn need_upgrade(&self, conn: &Connection) -> Result<bool, Error> {
Ok(false)
}
#[allow(unused_variables)]
fn upgrade(&self, conn: &mut Connection) -> Result<(), Error> {
Ok(())
}
}
struct RoConn {
regfuncs: Vec<CbType>
}
impl fmt::Debug for RoConn {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "RoConn {{}}")
}
}
impl CustomizeConnection<rusqlite::Connection, rusqlite::Error> for RoConn {
fn on_acquire(
&self,
conn: &mut rusqlite::Connection
) -> Result<(), rusqlite::Error> {
conn.pragma_update(None, "foreign_keys", "ON")?;
for rf in &self.regfuncs {
match rf {
CbType::Ro(f) | CbType::Both(f) => {
f(conn)?;
}
CbType::Rw(_) => {}
}
}
Ok(())
}
fn on_release(&self, _conn: rusqlite::Connection) {}
}
pub struct Builder {
schmgr: Box<dyn SchemaMgr>,
full_vacuum: bool,
autovacuum: bool,
max_readers: usize,
hook: Option<Arc<dyn Hook + Send + Sync>>,
regfuncs: Option<Vec<CbType>>,
#[cfg(feature = "tpool")]
tpool: Option<Arc<ThreadPool>>
}
impl Builder {
fn open_writer(&self, fname: &Path) -> Result<Connection, rusqlite::Error> {
let conn = Connection::open(fname)?;
conn.pragma_update(None, "journal_mode", "WAL")?;
conn.pragma_update(None, "foreign_keys", "ON")?;
if self.autovacuum {
conn.pragma_update(None, "auto_vacuum", "INCREMENTAL")?;
}
Ok(conn)
}
fn full_vacuum(conn: &Connection) -> Result<(), rusqlite::Error> {
conn.execute("VACUUM;", params![])?;
Ok(())
}
fn create_ro_pool(
&self,
fname: &Path,
regfuncs: Vec<CbType>
) -> Result<r2d2::Pool<SqliteConnectionManager>, r2d2::Error> {
let fl =
OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX;
let manager = SqliteConnectionManager::file(fname).with_flags(fl);
let roconn_initterm = RoConn { regfuncs };
let max_readers = u32::try_from(self.max_readers).unwrap();
r2d2::Pool::builder()
.max_size(max_readers)
.connection_customizer(Box::new(roconn_initterm))
.build(manager)
}
}
impl Builder {
#[must_use]
pub fn new(schmgr: Box<dyn SchemaMgr>) -> Self {
Self {
schmgr,
full_vacuum: false,
autovacuum: false,
max_readers: 2,
hook: None,
regfuncs: None,
#[cfg(feature = "tpool")]
tpool: None
}
}
#[must_use]
pub const fn init_vacuum(mut self) -> Self {
self.full_vacuum = true;
self
}
pub const fn init_vacuum_r(&mut self) -> &mut Self {
self.full_vacuum = true;
self
}
#[must_use]
pub const fn autovacuum(mut self) -> Self {
self.autovacuum = true;
self
}
pub const fn autovacuum_r(&mut self) -> &mut Self {
self.autovacuum = true;
self
}
#[must_use]
pub const fn max_readers(mut self, n: usize) -> Self {
self.max_readers = n;
self
}
pub const fn max_readers_r(&mut self, n: usize) -> &mut Self {
self.max_readers = n;
self
}
#[must_use]
pub fn hook(mut self, hook: Arc<dyn Hook + Send + Sync>) -> Self {
self.hook = Some(hook);
self
}
pub fn hook_r(&mut self, hook: Arc<dyn Hook + Send + Sync>) -> &mut Self {
self.hook = Some(hook);
self
}
#[must_use]
pub fn reg_scalar_fn<F>(mut self, r: RegOn<F>) -> Self
where
F: Fn(&Connection) -> Result<(), rusqlite::Error> + Send + Sync + 'static
{
self.reg_scalar_fn_r(r);
self
}
pub fn reg_scalar_fn_r<F>(&mut self, r: RegOn<F>) -> &mut Self
where
F: Fn(&Connection) -> Result<(), rusqlite::Error> + Send + Sync + 'static
{
match r {
RegOn::RO(f) => {
self
.regfuncs
.get_or_insert(Vec::new())
.push(CbType::Ro(Box::new(f)));
}
RegOn::RW(f) => {
self
.regfuncs
.get_or_insert(Vec::new())
.push(CbType::Rw(Box::new(f)));
}
RegOn::Both(f) => {
self
.regfuncs
.get_or_insert(Vec::new())
.push(CbType::Both(Box::new(f)));
}
}
self
}
#[cfg(feature = "tpool")]
#[must_use]
pub fn thread_pool(mut self, tpool: Arc<ThreadPool>) -> Self {
self.tpool = Some(tpool);
self
}
#[cfg(feature = "tpool")]
pub fn thread_pool_r(&mut self, tpool: Arc<ThreadPool>) -> &mut Self {
self.tpool = Some(tpool);
self
}
pub fn build<P>(mut self, fname: P) -> Result<ConnPool, Error>
where
P: AsRef<Path>
{
let fname = fname.as_ref();
let db_exists = fname.exists();
let mut conn = self.open_writer(fname)?;
let regfuncs = self.regfuncs.take().unwrap_or_default();
for rf in ®funcs {
match rf {
CbType::Rw(f) | CbType::Both(f) => {
f(&conn)?;
}
CbType::Ro(_) => {}
}
}
self.schmgr.init(&mut conn, !db_exists)?;
if self.schmgr.need_upgrade(&conn)? {
self.schmgr.upgrade(&mut conn)?;
}
if self.full_vacuum {
Self::full_vacuum(&conn)?;
}
if let Some(ref hook) = self.hook {
rawhook::hook(&conn, hook)?;
}
let rpool = self.create_ro_pool(fname, regfuncs)?;
let iconn = InnerWrConn { conn, dirt: 0 };
let inner = Inner { conn: Some(iconn) };
let sh = Arc::new(Shared {
inner: Mutex::new(inner),
signal: Condvar::new()
});
Ok(ConnPool {
rpool,
sh,
#[cfg(feature = "tpool")]
tpool: self.tpool
})
}
pub fn build_with_changelog_hook<P, D, T>(
mut self,
fname: P,
hook: Box<dyn ChangeLogHook<Database = D, Table = T> + Send>
) -> Result<ConnPool, Error>
where
P: AsRef<Path>,
D: FromStr + Send + Sized + 'static,
T: FromStr + Send + Sized + 'static
{
assert!(
self.hook.is_some(),
"Can't build a connection pool with both a raw and changelog hook"
);
let fname = fname.as_ref();
let db_exists = fname.exists();
let mut conn = self.open_writer(fname)?;
let regfuncs = self.regfuncs.take().unwrap_or_default();
for rf in ®funcs {
match rf {
CbType::Rw(f) | CbType::Both(f) => {
f(&conn)?;
}
CbType::Ro(_) => {}
}
}
self.schmgr.init(&mut conn, !db_exists)?;
if self.schmgr.need_upgrade(&conn)? {
self.schmgr.upgrade(&mut conn)?;
}
if self.full_vacuum {
Self::full_vacuum(&conn)?;
}
changehook::hook(&conn, hook)?;
let rpool = self.create_ro_pool(fname, regfuncs)?;
let iconn = InnerWrConn { conn, dirt: 0 };
let inner = Inner { conn: Some(iconn) };
let sh = Arc::new(Shared {
inner: Mutex::new(inner),
signal: Condvar::new()
});
Ok(ConnPool {
rpool,
sh,
#[cfg(feature = "tpool")]
tpool: self.tpool
})
}
}
struct InnerWrConn {
conn: Connection,
dirt: usize
}
struct Inner {
conn: Option<InnerWrConn>
}
struct Shared {
inner: Mutex<Inner>,
signal: Condvar
}
impl Shared {
#[inline]
fn lock(&self) -> MutexGuard<'_, Inner> {
self.inner.lock()
}
#[inline]
fn with_lock_guard<F, R>(&self, f: F) -> R
where
F: FnOnce(MutexGuard<'_, Inner>) -> R
{
let g = self.lock();
f(g)
}
}
#[derive(Clone)]
pub struct ConnPool {
rpool: r2d2::Pool<SqliteConnectionManager>,
sh: Arc<Shared>,
#[cfg(feature = "tpool")]
tpool: Option<Arc<ThreadPool>>
}
impl ConnPool {
#[must_use]
pub fn size(&self) -> usize {
(self.rpool.max_size() + 1) as usize
}
pub fn reader(
&self
) -> Result<PooledConnection<SqliteConnectionManager>, r2d2::Error> {
self.rpool.get()
}
#[must_use]
pub fn writer(&self) -> WrConn {
let conn = self.sh.with_lock_guard(|mut g| {
loop {
if let Some(conn) = g.conn.take() {
break conn;
}
self.sh.signal.wait(&mut g);
}
});
WrConn {
sh: Arc::clone(&self.sh),
inner: ManuallyDrop::new(conn)
}
}
#[must_use]
pub fn try_writer(&self) -> Option<WrConn> {
let conn = self.sh.inner.lock().conn.take()?;
Some(WrConn {
sh: Arc::clone(&self.sh),
inner: ManuallyDrop::new(conn)
})
}
}
impl ConnPool {
pub fn freelist_count(&self) -> Result<usize, Error> {
let npages = self.reader()?.query_row_and_then(
"PRAGMA freelist_count;",
[],
|row| row.get::<_, i64>(0)
)?;
usize::try_from(npages).map_err(|_| {
Error::oob("Freelist count could not be expressed as an `usize`")
})
}
}
impl ConnPool {
#[inline]
pub fn with_ro<T, F, E>(&self, f: F) -> Result<T, E>
where
T: Send + 'static,
F: FnOnce(&Connection) -> Result<T, E> + Send + 'static,
E: From<r2d2::Error>
{
let conn = self.reader()?;
f(&conn)
}
#[cfg(feature = "tpool")]
#[inline]
pub fn with_ro_thrd<F>(&self, f: F) -> Result<(), r2d2::Error>
where
F: FnOnce(&Connection) + Send + 'static
{
let Some(ref tpool) = self.tpool else {
panic!("ConnPool does to have a thread pool");
};
let conn = self.reader()?;
tpool.execute(move || {
f(&conn);
});
Ok(())
}
#[cfg(feature = "tpool")]
#[inline]
pub fn with_ro_thrd_result<T, E, F>(
&self,
f: F
) -> Result<swctx::WaitCtx<T, (), E>, r2d2::Error>
where
T: Send + 'static,
E: fmt::Debug + Send + 'static,
F: FnOnce(&Connection) -> Result<T, E> + Send + 'static
{
let Some(ref tpool) = self.tpool else {
panic!("ConnPool does to have a thread pool");
};
let conn = self.reader()?;
let (sctx, wctx) = swctx::mkpair();
tpool.execute(move || match f(&conn) {
Ok(t) => {
let _ = sctx.set(t);
}
Err(e) => {
let _ = sctx.fail(e);
}
});
Ok(wctx)
}
}
impl ConnPool {
#[inline]
pub fn with_rw<T, E, F>(&self, f: F) -> Result<T, E>
where
T: Send + 'static,
E: fmt::Debug + Send + 'static,
F: FnOnce(&mut WrConn) -> Result<T, E> + Send + 'static
{
let mut conn = self.writer();
f(&mut conn)
}
#[cfg(feature = "tpool")]
#[inline]
pub fn with_rw_thrd<F>(&self, f: F)
where
F: FnOnce(&mut WrConn) -> Option<usize> + Send + 'static
{
let Some(ref tpool) = self.tpool else {
panic!("ConnPool does to have a thread pool");
};
let mut conn = self.writer();
tpool.execute(move || {
let dirt = f(&mut conn);
if let Some(dirt) = dirt {
conn.add_dirt(dirt);
}
});
}
#[cfg(feature = "tpool")]
#[inline]
pub fn with_rw_thrd_result<T, E, F>(&self, f: F) -> swctx::WaitCtx<T, (), E>
where
T: Send + 'static,
E: fmt::Debug + Send + 'static,
F: FnOnce(&mut WrConn) -> Result<T, E> + Send + 'static
{
let Some(ref tpool) = self.tpool else {
panic!("ConnPool does to have a thread pool");
};
let mut conn = self.writer();
let (sctx, wctx) = swctx::mkpair();
tpool.execute(move || match f(&mut conn) {
Ok(t) => {
let _ = sctx.set(t);
}
Err(e) => {
let _ = sctx.fail(e);
}
});
wctx
}
}
impl ConnPool {
#[cfg(feature = "tpool")]
#[must_use]
pub fn incremental_vacuum(
&self,
n: Option<usize>
) -> swctx::WaitCtx<(), (), Error> {
self.with_rw_thrd_result(move |conn| conn.incremental_vacuum(n))
}
}