ouisync/db/

mod.rs

#[macro_use]
mod macros;

mod connection;
mod id;
mod migrations;

pub use id::DatabaseId;
pub use migrations::SCHEMA_VERSION;

use tracing::Span;

use deadlock::ExpectShortLifetime;
use ref_cast::RefCast;
use sqlx::{
    ConnectOptions, Row, SqlitePool, TransactionManager,
    sqlite::{
        Sqlite, SqliteAutoVacuum, SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions,
        SqliteSynchronous, SqliteTransactionManager,
    },
};
use std::{
    fmt, io,
    ops::{Deref, DerefMut},
    panic::Location,
    path::{Path, PathBuf},
    time::Duration,
};
#[cfg(test)]
use tempfile::TempDir;
use thiserror::Error;
use tokio::{fs, sync::Semaphore, task};

const ACQUIRE_TIMEOUT: Duration = Duration::from_secs(5 * 60);
const IDLE_TIMEOUT: Duration = Duration::from_secs(60);
const WARN_AFTER_CONNECTION_LIFETIME: Duration = Duration::from_secs(30);

pub use self::connection::Connection;

/// Database connection pool.
#[derive(Clone)]
pub struct Pool {
    // Path to the main database file
    path: PathBuf,
    // Pool with multiple read-only connections
    reads: SqlitePool,
    // Pool with a single writable connection.
    write: SqlitePool,
}

impl Pool {
    async fn create(conn_options: SqliteConnectOptions) -> Result<Self, sqlx::Error> {
        let path = conn_options.get_filename().to_path_buf();

        if fs::try_exists(&path).await.unwrap_or(false) {
            // Try to enable auto-vacuum, but if it fails [^1] it's not a critical failure as we can
            // keep using the db without auto-vacuum [^2]. Just log the error and keep going.
            //
            // [^1]: for example, because of low disk space
            // [^2]: with the downside that disk space won't be reclaimed after data deletion
            if let Err(error) = enable_auto_vacuum(conn_options.get_filename()).await {
                tracing::warn!(?error, "failed to enable auto vacuum");
            }
        }

        let conn_options = conn_options
            .journal_mode(SqliteJournalMode::Wal)
            .synchronous(SqliteSynchronous::Normal)
            .pragma("recursive_triggers", "ON");

        let pool_options = SqlitePoolOptions::new()
            // Disable the test as it breaks cancel-safety (also it's unnecessary in our case)
            .test_before_acquire(false)
            // Expire idle connections to conserve resources (threads, file descriptors)
            .idle_timeout(IDLE_TIMEOUT)
            .acquire_timeout(ACQUIRE_TIMEOUT);

        let write = pool_options
            .clone()
            .max_connections(1)
            .connect_with(
                conn_options
                    .clone()
                    .auto_vacuum(SqliteAutoVacuum::Full)
                    .optimize_on_close(true, Some(1000)),
            )
            .await?;

        let reads = pool_options
            .max_connections(8)
            .connect_with(conn_options.read_only(true))
            .await?;

        Ok(Self { path, reads, write })
    }

    /// Return path to the database file.
    pub fn store_path(&self) -> &Path {
        self.path.as_ref()
    }

    /// Acquire a read-only database connection.
    #[track_caller]
    pub fn acquire(&self) -> impl Future<Output = Result<PoolConnection, sqlx::Error>> + '_ {
        PoolConnection::acquire(&self.reads, Location::caller())
    }

    /// Begin a read-only transaction. See [`ReadTransaction`] for more details.
    #[track_caller]
    pub fn begin_read(&self) -> impl Future<Output = Result<ReadTransaction, sqlx::Error>> + '_ {
        ReadTransaction::begin(&self.reads, Location::caller())
    }

    /// Begin a write transaction. See [`WriteTransaction`] for more details.
    #[track_caller]
    pub fn begin_write(&self) -> impl Future<Output = Result<WriteTransaction, sqlx::Error>> + '_ {
        let location = Location::caller();

        async move {
            Ok(WriteTransaction {
                inner: ReadTransaction::begin(&self.write, location).await?,
            })
        }
    }

    pub(crate) async fn close(&self) -> Result<(), sqlx::Error> {
        // Make sure to first close `reads` and only then `write`. That way when closing the write
        // connection it is the last remaining connection and so it performs a WAL checkpoint and
        // removes the auxiliary db files (*-wal and *-shm).
        self.reads.close().await;
        self.write.close().await;

        Ok(())
    }
}

/// Database connection from pool
pub struct PoolConnection {
    inner: sqlx::pool::PoolConnection<Sqlite>,
    _track_lifetime: ExpectShortLifetime,
}

impl PoolConnection {
    // Internal
    async fn acquire(
        pool: &SqlitePool,
        location: &'static Location<'static>,
    ) -> Result<Self, sqlx::Error> {
        let inner = pool.acquire().await?;
        let track_lifetime = ExpectShortLifetime::new_in(WARN_AFTER_CONNECTION_LIFETIME, location);

        Ok(Self {
            inner,
            _track_lifetime: track_lifetime,
        })
    }
}

impl Deref for PoolConnection {
    type Target = Connection;

    fn deref(&self) -> &Self::Target {
        Connection::ref_cast(self.inner.deref())
    }
}

impl DerefMut for PoolConnection {
    fn deref_mut(&mut self) -> &mut Self::Target {
        Connection::ref_cast_mut(self.inner.deref_mut())
    }
}

/// Transaction that allows only reading.
///
/// This is useful if one wants to make sure the observed database content doesn't change for the
/// duration of the transaction even in the presence of concurrent writes. In other words - a read
/// transaction represents an immutable snapshot of the database at the point the transaction was
/// created. A read transaction doesn't need to be committed or rolled back - it's implicitly ended
/// when the `ReadTransaction` instance drops.
pub struct ReadTransaction {
    inner: PoolConnection,
    closed: bool,
}

impl ReadTransaction {
    // Internal
    async fn begin(
        pool: &SqlitePool,
        location: &'static Location<'static>,
    ) -> Result<Self, sqlx::Error> {
        let mut inner = PoolConnection::acquire(pool, location).await?;
        SqliteTransactionManager::begin(&mut inner.inner, None).await?;

        Ok(Self {
            inner,
            closed: false,
        })
    }

    // Internal
    async fn commit(mut self) -> Result<Committed, sqlx::Error> {
        SqliteTransactionManager::commit(&mut self.inner.inner).await?;
        self.closed = true;
        Ok(Committed(self))
    }
}

impl Deref for ReadTransaction {
    type Target = Connection;

    fn deref(&self) -> &Self::Target {
        self.inner.deref()
    }
}

impl DerefMut for ReadTransaction {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.inner.deref_mut()
    }
}

impl fmt::Debug for ReadTransaction {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("ReadTransaction").finish_non_exhaustive()
    }
}

impl_executor_by_deref!(ReadTransaction);

impl Drop for ReadTransaction {
    fn drop(&mut self) {
        if !self.closed {
            SqliteTransactionManager::start_rollback(&mut self.inner.inner);
        }
    }
}

// Wrapper for a transaction that's been committed. This allows to delay releasing the underlying
// connection to the pool while disallowing using the connection to execute any db operations.
struct Committed(#[allow(dead_code)] ReadTransaction);

/// Transaction that allows both reading and writing.
///
/// At most one task can hold a write transaction at any time. Any other tasks are blocked on
/// calling `begin_write` until the task that currently holds it is done with it (commits it or
/// rolls it back). Performing read-only operations concurrently while a write transaction is in
/// use is still allowed. Those operations will not see the writes performed via the write
/// transaction until that transaction is committed however.
pub struct WriteTransaction {
    inner: ReadTransaction,
}

impl WriteTransaction {
    /// Commits the transaction.
    ///
    /// # Cancel safety
    ///
    /// If the future returned by this function is cancelled before completion, the transaction
    /// is guaranteed to be either committed or rolled back but there is no way to tell in advance
    /// which of the two operations happens.
    pub async fn commit(self) -> Result<(), sqlx::Error> {
        self.inner.commit().await?;
        Ok(())
    }

    /// Commits the transaction and if (and only if) the commit completes successfully, runs the
    /// given closure.
    ///
    /// # Atomicity
    ///
    /// If the commit succeeds, the closure is guaranteed to complete before another write
    /// transaction begins.
    ///
    /// # Cancel safety
    ///
    /// The commits completes and if it succeeds the closure gets called. This is guaranteed to
    /// happen even if the future returned from this function is cancelled before completion.
    ///
    /// # Insufficient alternatives
    ///
    /// ## Calling `commit().await?` and then calling `f()`
    ///
    /// This is not enough because it has these possible outcomes depending on whether and when
    /// cancellation happened:
    ///
    /// 1. `commit` completes successfully and `f` is called
    /// 2. `commit` completes with error and `f` is not called
    /// 3. `commit` is cancelled but the transaction is still committed and `f` is not called
    /// 4. `commit` is cancelled and the transaction rolls back and `f` is not called
    ///
    /// Number 3 is typically not desirable.
    ///
    /// ## Calling `f` using a RAII guard
    ///
    /// This is still not enough because it has the following possible outcomes:
    ///
    /// 1. `commit` completes successfully and `f` is called
    /// 2. `commit` completes with error and `f` is called
    /// 3. `commit` is cancelled but the transaction is still committed and `f` is called
    /// 4. `commit` is cancelled and the transaction rolls back and `f` is called
    ///
    /// Numbers 2 and 4 are not desirable. Number 2 can be handled by explicitly handling the error
    /// case and disabling the guard but there is nothing to do about number 4.
    pub async fn commit_and_then<F, R>(self, f: F) -> Result<R, sqlx::Error>
    where
        F: FnOnce() -> R + Send + 'static,
        R: Send + 'static,
    {
        let span = Span::current();

        task::spawn(async move {
            // IMPORTANT: `_committed` must live until `f` completes.
            let _committed = self.inner.commit().await?;
            let result = span.in_scope(f);
            Ok(result)
        })
        .await
        .unwrap()
    }
}

impl Deref for WriteTransaction {
    type Target = ReadTransaction;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl DerefMut for WriteTransaction {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

impl std::fmt::Debug for WriteTransaction {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
        write!(f, "WriteTransaction{{ inner:{:?} }}", self.inner)
    }
}

impl_executor_by_deref!(WriteTransaction);

#[derive(Debug, Error)]
pub enum Error {
    #[error("failed to create database directory")]
    CreateDirectory(#[source] io::Error),
    #[error("database already exists")]
    Exists,
    #[error("failed to open database")]
    Open(#[source] sqlx::Error),
    #[error("failed to execute database query")]
    Query(#[from] sqlx::Error),
}

/// Creates a new database and opens a connection to it.
pub(crate) async fn create(path: impl AsRef<Path>) -> Result<Pool, Error> {
    let path = path.as_ref();

    if fs::metadata(path).await.is_ok() {
        return Err(Error::Exists);
    }

    create_directory(path).await?;

    let connect_options = SqliteConnectOptions::new()
        .filename(path)
        .create_if_missing(true);

    let pool = Pool::create(connect_options).await.map_err(Error::Open)?;

    migrations::run(&pool).await?;

    Ok(pool)
}

/// Creates a new database in a temporary directory. Useful for tests.
#[cfg(test)]
pub(crate) async fn create_temp() -> Result<(TempDir, Pool), Error> {
    let temp_dir = TempDir::new().map_err(Error::CreateDirectory)?;
    let pool = create(temp_dir.path().join("temp.db")).await?;

    Ok((temp_dir, pool))
}

/// Opens a connection to the specified database. Fails if the db doesn't exist.
pub(crate) async fn open(path: impl AsRef<Path>) -> Result<Pool, Error> {
    let connect_options = SqliteConnectOptions::new().filename(path);
    let pool = Pool::create(connect_options).await.map_err(Error::Open)?;

    migrations::run(&pool).await?;

    Ok(pool)
}

/// Opens a connection to the specified database. Fails if the db doesn't exist.
pub async fn open_without_migrations(path: impl AsRef<Path>) -> Result<Pool, Error> {
    let connect_options = SqliteConnectOptions::new().filename(path);
    let pool = Pool::create(connect_options).await.map_err(Error::Open)?;

    Ok(pool)
}

async fn create_directory(path: &Path) -> Result<(), Error> {
    if let Some(dir) = path.parent() {
        fs::create_dir_all(dir)
            .await
            .map_err(Error::CreateDirectory)?
    }

    Ok(())
}

// Explicit cast from `i64` to `u64` to work around the lack of native `u64` support in the sqlx
// crate.
pub(crate) const fn decode_u64(i: i64) -> u64 {
    i as u64
}

// Explicit cast from `u64` to `i64` to work around the lack of native `u64` support in the sqlx
// crate.
pub(crate) const fn encode_u64(u: u64) -> i64 {
    u as i64
}

// Enable auto-vacuum on the given database unless already enabled
async fn enable_auto_vacuum(db_path: &Path) -> Result<(), Error> {
    let mut conn = SqliteConnectOptions::new()
        .filename(db_path)
        .connect()
        .await?;

    let auto_vacuum: u32 = sqlx::query("PRAGMA auto_vacuum")
        .fetch_one(&mut conn)
        .await?
        .get(0);

    if auto_vacuum != 0 {
        return Ok(());
    }

    // VACUUM requires up to twice the size of the original database of free disk space. To prevent
    // exhausting available disk space when opening multiple databases in parallel, we limit the
    // number of dbs being vacuumed at the same time to one. Dbs that don't need to be vacuumed are
    // not affected by this limit.
    static SEMAPHORE: Semaphore = Semaphore::const_new(1);
    let _permit = SEMAPHORE.acquire().await.unwrap();

    sqlx::query("PRAGMA auto_vacuum=FULL")
        .execute(&mut conn)
        .await?;

    // Execute `VACUUM` command for the `auto_vacuum` pragma to take effect.
    sqlx::query("VACUUM").execute(&mut conn).await?;

    Ok(())
}

async fn get_pragma(conn: &mut Connection, name: &str) -> Result<u32, Error> {
    Ok(sqlx::query(&format!("PRAGMA {name}"))
        .fetch_one(&mut *conn)
        .await?
        .get(0))
}

async fn set_pragma(conn: &mut Connection, name: &str, value: u32) -> Result<(), Error> {
    // `bind` doesn't seem to be supported for setting PRAGMAs...
    sqlx::query(&format!("PRAGMA {name} = {value}"))
        .execute(&mut *conn)
        .await?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    // Check the casts are lossless

    #[test]
    fn decode_u64_sanity_check() {
        // [0i64,     i64::MAX] -> [0u64,             u64::MAX / 2]
        // [i64::MIN,    -1i64] -> [u64::MAX / 2 + 1,     u64::MAX]

        assert_eq!(decode_u64(0), 0);
        assert_eq!(decode_u64(1), 1);
        assert_eq!(decode_u64(-1), u64::MAX);
        assert_eq!(decode_u64(i64::MIN), u64::MAX / 2 + 1);
        assert_eq!(decode_u64(i64::MAX), u64::MAX / 2);
    }

    #[test]
    fn encode_u64_sanity_check() {
        assert_eq!(encode_u64(0), 0);
        assert_eq!(encode_u64(1), 1);
        assert_eq!(encode_u64(u64::MAX / 2), i64::MAX);
        assert_eq!(encode_u64(u64::MAX / 2 + 1), i64::MIN);
        assert_eq!(encode_u64(u64::MAX), -1);
    }
}