vault-link/sync-server/src/app_state/database.rs

use core::time::Duration;
use std::{
    collections::HashMap,
    path::{Path, PathBuf},
    sync::{Arc, Weak},
};

use anyhow::{Context as _, Result};
use log::info;
use models::VaultId;
use sha2::{Digest, Sha256};
use sqlx::{ConnectOptions, Executor, sqlite::SqliteConnectOptions};
use sqlx::{Pool, Sqlite, sqlite::SqlitePoolOptions};
use tokio::sync::Mutex;
use tokio::time::Instant;

pub mod models;

mod history;
mod mutations;
mod queries;

#[cfg(test)]
mod tests;

use crate::{
    app_state::weak_slots::get_or_create,
    config::database_config::DatabaseConfig,
    utils::normalize_vault_id::{normalize_string, validate_vault_id},
};

const VAULT_DATABASE_DIRECTORY: &str = "vaults";

#[derive(Clone)]
struct PoolWithTimestamp {
    pool: Pool<Sqlite>,
    last_accessed: Instant,
}

impl std::fmt::Debug for PoolWithTimestamp {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PoolWithTimestamp")
            .field("pool", &"Pool<Sqlite>")
            .field("last_accessed", &self.last_accessed)
            .finish()
    }
}

#[derive(Clone, Debug)]
pub struct Database {
    config: DatabaseConfig,
    connection_pools: Arc<Mutex<HashMap<VaultId, PoolWithTimestamp>>>,
    opening_pools: Arc<Mutex<HashMap<VaultId, Weak<Mutex<()>>>>>,
}

pub type Transaction<'a> = sqlx::Transaction<'a, Sqlite>;

impl Database {
    pub async fn try_new(config: &DatabaseConfig) -> Result<Self> {
        let vault_directory = config
            .databases_directory_path
            .join(VAULT_DATABASE_DIRECTORY);
        tokio::fs::create_dir_all(&vault_directory)
            .await
            .with_context(|| {
                format!(
                    "Failed to create databases directory at `{}`",
                    config.databases_directory_path.to_string_lossy()
                )
            })?;

        // SQLite flushes its own files, not ancestors created by this process
        Self::flush_directory_ancestors(&vault_directory).await?;

        let database = Self {
            config: config.clone(),
            connection_pools: Arc::default(),
            opening_pools: Arc::default(),
        };

        // Start background task to cleanup idle connection pools
        database.start_idle_pool_cleanup();

        Ok(database)
    }

    fn get_database_path(config: &DatabaseConfig, vault: &str) -> PathBuf {
        // Hash the normalized vault name into a fixed-length ASCII filename so
        // filesystem case/Unicode aliasing (e.g. composed vs. decomposed é on
        // macOS) cannot make separately authorized vaults share a database.
        // This is only a storage filename; the user-facing vault name is unchanged.
        let digest = Sha256::digest(vault.as_bytes());
        config
            .databases_directory_path
            .join(VAULT_DATABASE_DIRECTORY)
            .join(format!("{digest:x}.sqlite"))
    }

    async fn open_database(config: &DatabaseConfig, file_name: &Path) -> Result<Pool<Sqlite>> {
        let connection_options = SqliteConnectOptions::new()
            .filename(file_name)
            .create_if_missing(true)
            .auto_vacuum(sqlx::sqlite::SqliteAutoVacuum::Full)
            .busy_timeout(Duration::from_secs(3600))
            .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal)
            .synchronous(sqlx::sqlite::SqliteSynchronous::Full)
            .pragma("fullfsync", "ON")
            .foreign_keys(true)
            .log_slow_statements(log::LevelFilter::Warn, Duration::from_secs(30));

        let pool = SqlitePoolOptions::new()
            .max_connections(config.max_connections_per_vault)
            .acquire_slow_threshold(Duration::from_secs(30))
            .test_before_acquire(true)
            .connect_with(connection_options)
            .await
            .with_context(|| format!("Cannot open database at `{}`", file_name.display()))?;

        Self::initialize_schema(&pool).await?;

        if let Some(directory) = file_name.parent() {
            #[cfg(unix)]
            Self::flush_directory_ancestors(directory).await?;
        }

        Ok(pool)
    }

    async fn flush_directory_ancestors(directory: &Path) -> Result<()> {
        let directory = tokio::fs::canonicalize(directory).await?;
        for ancestor in directory.ancestors() {
            tokio::fs::File::open(ancestor)
                .await?
                .sync_all()
                .await
                .with_context(|| format!("Cannot flush directory {}", ancestor.display()))?;
        }

        Ok(())
    }

    async fn initialize_schema(pool: &Pool<Sqlite>) -> Result<()> {
        let mut tx = pool.begin_with("BEGIN IMMEDIATE").await?;
        let initialized: bool = sqlx::query_scalar(
            "SELECT EXISTS(SELECT 1 FROM sqlite_schema WHERE name NOT LIKE 'sqlite_%')",
        )
        .fetch_one(&mut *tx)
        .await?;
        if !initialized {
            tx.execute(include_str!("database/schema.sql"))
                .await
                .context("Cannot initialize vault database")?;
        }
        tx.commit()
            .await
            .context("Cannot commit vault database initialization")
    }

    async fn get_connection_pool(&self, vault: &VaultId) -> Result<Pool<Sqlite>> {
        let vault = normalize_string(vault);
        validate_vault_id(&vault)?;

        if let Some(pool) = self.find_cached_pool(&vault).await {
            return Ok(pool);
        }

        let opening = get_or_create(self.opening_pools.lock(), vault.clone(), || {
            Arc::new(Mutex::new(()))
        })
        .await;

        let _opening = opening.lock().await;

        if let Some(pool) = self.find_cached_pool(&vault).await {
            return Ok(pool);
        }

        let file_name = Self::get_database_path(&self.config, &vault);
        let pool = Self::open_database(&self.config, &file_name).await?;

        self.connection_pools.lock().await.insert(
            vault,
            PoolWithTimestamp {
                pool: pool.clone(),
                last_accessed: Instant::now(),
            },
        );
        Ok(pool)
    }

    async fn find_cached_pool(&self, vault: &VaultId) -> Option<Pool<Sqlite>> {
        self.connection_pools
            .lock()
            .await
            .get_mut(vault)
            .map(|entry| {
                entry.last_accessed = Instant::now();
                entry.pool.clone()
            })
    }

    /// Attempting to write from this transaction might result in a
    /// database locked error. Use this transaction for read-only operations.
    pub async fn create_readonly_transaction(
        &self,
        vault: &VaultId,
    ) -> Result<Transaction<'static>> {
        self.get_connection_pool(vault)
            .await?
            .begin()
            .await
            .context("Cannot create transaction")
    }

    pub async fn create_write_transaction(&self, vault: &VaultId) -> Result<Transaction<'static>> {
        self.get_connection_pool(vault)
            .await?
            .begin_with("BEGIN IMMEDIATE")
            .await
            .context("Cannot create write transaction")
    }

    /// Cleanup idle connection pools that haven't been accessed in more than 5 minutes
    async fn cleanup_idle_pools(&self) {
        let mut pools = self.connection_pools.lock().await;
        let now = Instant::now();
        let idle_timeout = Duration::from_secs(5 * 60); // 5 minutes

        // Collect vaults to remove
        let vaults_to_remove: Vec<VaultId> = pools
            .iter()
            .filter(|(_, pool_with_timestamp)| {
                now.duration_since(pool_with_timestamp.last_accessed) > idle_timeout
                    && pool_with_timestamp.pool.num_idle()
                        == pool_with_timestamp.pool.size() as usize
            })
            .map(|(vault_id, _)| vault_id.clone())
            .collect();

        let closing: Vec<_> = vaults_to_remove
            .into_iter()
            .filter_map(|vault| pools.remove(&vault).map(|entry| (vault, entry.pool)))
            .collect();

        // Don't hold the pools mutex while closing the removed vaults
        drop(pools);

        // A checkout racing cleanup can delay close. It must never retain the
        // mutex used to access every other vault in the process.
        for (vault, pool) in closing {
            info!("Closing idle database connection pool for vault `{vault}`");
            pool.close().await;
        }
    }

    /// Start a background task that periodically cleans up idle connection pools
    fn start_idle_pool_cleanup(&self) {
        let database = self.clone();
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_secs(60)); // Check every minute
            interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);

            loop {
                interval.tick().await;
                database.cleanup_idle_pools().await;
            }
        });
    }
}