sqlite_rwc/

pool.rs

use crate::drivers::Driver;
use parking_lot::{Condvar, Mutex};
#[cfg(feature = "watcher")]
use sqlite_watcher::connection::State;
#[cfg(feature = "watcher")]
use sqlite_watcher::watcher::Watcher;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;

pub struct ConnectionPoolConfig {
    pub max_read_connection_count: usize,
    pub file_path: PathBuf,
    pub connection_acquire_timeout: Option<Duration>,
    #[cfg(feature = "watcher")]
    pub watcher: Arc<Watcher>,
}

pub struct ConnectionPool<T: Driver, A: ConnectionAdapter<T>> {
    read_connections: Mutex<Vec<A>>,
    reader_condvar: Condvar,
    write_connection: Mutex<WatchedConnection<T>>,
    config: ConnectionPoolConfig,
}

#[derive(Debug, thiserror::Error)]
pub enum ConnectionPoolError<E> {
    #[error(transparent)]
    Driver(#[from] E),
    #[error("Failed to acquire connection in time")]
    ConnectionAcquireTimeout,
    #[error("Failed to setup connection watcher")]
    WatcherSetup,
}

impl<T: Driver, A: ConnectionAdapter<T>> ConnectionPool<T, A> {
    /// Create a new connection pool with the given `config`.
    ///
    /// The write connection is created first, followed by a read connection for every reader.
    ///
    /// # Errors
    ///
    /// Returns error if the connections could not be initialized.
    pub fn new(
        config: ConnectionPoolConfig,
    ) -> Result<Arc<Self>, ConnectionPoolError<T::ConnectionError>> {
        let watched_connection = T::new_write_connection(&config.file_path)
            .inspect_err(|e| tracing::error!("Failed to create write connection: {e:?}"))?;
        #[cfg(feature = "watcher")]
        let watched_connection = WatchedConnection::new(watched_connection).map_err(|e| {
            tracing::error!("Failed to setup connection watcher: {e:?}");
            ConnectionPoolError::WatcherSetup
        })?;
        #[cfg(not(feature = "watcher"))]
        let watched_connection = WatchedConnection::new(watched_connection);

        let mut read_connections = Vec::with_capacity(config.max_read_connection_count);
        for _ in 0..config.max_read_connection_count {
            read_connections.push(A::from_driver_connection(
                T::new_read_connection(&config.file_path)
                    .inspect_err(|e| tracing::error!("Failed to create read connection: {e:?}"))?,
            ));
        }
        Ok(Arc::new(Self {
            write_connection: Mutex::new(watched_connection),
            read_connections: Mutex::new(read_connections),
            reader_condvar: Condvar::new(),
            config,
        }))
    }

    /// Retrieve a connection from the pool.
    ///
    /// If all the connections are currently in use, we will wait until one is returned to the
    /// pool. If `ConnectionPoolConfig.connection_acquire_timeout` has no value, this method will
    /// block indefinitely.
    ///
    /// # Errors
    ///
    /// Return error if we could not retrieve a connection from the pool before the timeout
    /// triggered.
    pub fn connection(
        self: &Arc<Self>,
    ) -> Result<PooledConnection<T, A>, ConnectionPoolError<T::Error>> {
        let mut rd_connections = self.read_connections.lock();
        loop {
            if let Some(rd_connection) = rd_connections.pop() {
                return Ok(PooledConnection::new(self.clone(), rd_connection));
            } else if let Some(duration) = self.config.connection_acquire_timeout {
                if self
                    .reader_condvar
                    .wait_for(&mut rd_connections, duration)
                    .timed_out()
                {
                    return Err(ConnectionPoolError::ConnectionAcquireTimeout);
                }
            } else {
                self.reader_condvar.wait(&mut rd_connections);
            }
        }
    }

    pub(crate) fn transaction<F, R, E>(&self, closure: F) -> Result<R, E>
    where
        F: FnOnce(&mut T::Transaction<'_>) -> Result<R, E>,
        E: From<T::Error>,
    {
        let mut writer_connection = self.write_connection.lock();
        #[cfg(feature = "watcher")]
        {
            writer_connection.transaction(closure, &self.config.watcher)
        }
        #[cfg(not(feature = "watcher"))]
        {
            writer_connection.transaction(closure)
        }
    }

    fn return_to_pool(&self, conn: A) {
        let mut read_connections = self.read_connections.lock();
        read_connections.push(conn);
        drop(read_connections);
        self.reader_condvar.notify_one();
    }

    #[cfg(feature = "watcher")]
    pub fn watcher(&self) -> &Arc<Watcher> {
        &self.config.watcher
    }
}

pub trait ConnectionAdapter<T: Driver> {
    fn from_driver_connection(connection: T::Connection) -> Self;
}

pub struct PooledConnection<T: Driver, A: ConnectionAdapter<T>> {
    pub(crate) pool: Arc<ConnectionPool<T, A>>,
    conn: Option<A>,
}

impl<T: Driver, A: ConnectionAdapter<T>> Drop for PooledConnection<T, A> {
    fn drop(&mut self) {
        let conn = self.conn.take().expect("Connection should be set");
        self.pool.return_to_pool(conn);
    }
}

impl<T: Driver, A: ConnectionAdapter<T>> PooledConnection<T, A> {
    fn new(pool: Arc<ConnectionPool<T, A>>, connection: A) -> PooledConnection<T, A> {
        Self {
            pool,
            conn: Some(connection),
        }
    }

    pub(crate) fn connection(&self) -> &A {
        self.conn.as_ref().expect("Connection should be set")
    }

    pub(crate) fn connection_mut(&mut self) -> &mut A {
        self.conn.as_mut().expect("Connection should be set")
    }
}

struct WatchedConnection<T>
where
    T: Driver,
{
    connection: T::Connection,
    #[cfg(feature = "watcher")]
    state: State,
}

#[cfg(feature = "watcher")]
impl<T> WatchedConnection<T>
where
    T: Driver,
{
    fn new(mut connection: T::Connection) -> Result<Self, <T as Driver>::Error> {
        use sqlite_watcher::statement::Statement;
        State::start_tracking().execute_mut(&mut connection)?;
        Ok(Self {
            connection,
            state: State::new(),
        })
    }

    fn transaction<F, R, E>(&mut self, closure: F, watcher: &Watcher) -> Result<R, E>
    where
        F: FnOnce(&mut T::Transaction<'_>) -> Result<R, E>,
        E: From<T::Error>,
    {
        self.before_write(watcher)?;
        let result = T::write(&mut self.connection, closure);
        if let Err(e) = self.after_write(watcher) {
            tracing::error!("Failed to publish updates to watcher: {e:?}");
        }
        result
    }

    fn before_write(&mut self, watcher: &Watcher) -> Result<(), <T as Driver>::Error> {
        use sqlite_watcher::statement::Statement;
        if let Some(stmt) = self.state.sync_tables(watcher) {
            stmt.execute_mut(&mut self.connection)?;
        }
        Ok(())
    }

    fn after_write(&mut self, watcher: &Watcher) -> Result<(), <T as Driver>::Error> {
        use sqlite_watcher::statement::Statement;
        self.state
            .publish_changes(watcher)
            .execute_mut(&mut self.connection)?;
        Ok(())
    }
}

#[cfg(not(feature = "watcher"))]
impl<T> WatchedConnection<T>
where
    T: Driver,
{
    fn new(connection: T::Connection) -> Self {
        Self { connection }
    }

    fn transaction<F, R, E>(&mut self, closure: F) -> Result<R, E>
    where
        F: FnOnce(&mut T::Transaction<'_>) -> Result<R, E>,
        E: From<T::Error>,
    {
        T::write(&mut self.connection, closure)
    }
}