dragoman 0.5.12

//! Read connections to the stores, opened once and reused.
//!
//! Opening a store costs a file open, a read of the schema, the WAL index and
//! the pragmas below — more than the point lookup most requests make. Each
//! store gets one pool, found by its path, holding the connections idle
//! between requests; a request takes one for as long as its blocking task
//! runs and gives it back after. A request never waits for a connection: with
//! none idle it opens another, and one returned to a full pool is closed, so
//! the pool holds what steady load needs and a burst costs only its opens.

use std::collections::HashMap;
use std::ops::Deref;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, OnceLock};

use commonmeta::rusqlite::Connection;

use crate::error::AppError;

/// Connections kept idle per store: enough for every blocking task a burst of
/// requests runs at once, small next to the blocking pool's 512 threads.
const MAX_IDLE: usize = 64;

/// The idle read connections to one store.
pub struct Pool {
    path: PathBuf,
    idle: Mutex<Vec<Connection>>,
}

/// A connection taken from a [`Pool`], given back when dropped.
pub struct Pooled {
    conn: Option<Connection>,
    pool: Arc<Pool>,
}

impl Deref for Pooled {
    type Target = Connection;

    fn deref(&self) -> &Connection {
        self.conn
            .as_ref()
            .expect("a pooled connection is held until drop")
    }
}

impl Drop for Pooled {
    fn drop(&mut self) {
        let Some(conn) = self.conn.take() else {
            return;
        };
        let mut idle = self.pool.idle.lock().unwrap_or_else(|e| e.into_inner());
        if idle.len() < MAX_IDLE {
            idle.push(conn);
        }
    }
}

impl Pool {
    fn new(path: &Path) -> Self {
        Self {
            path: path.to_path_buf(),
            idle: Mutex::new(Vec::new()),
        }
    }

    /// An idle connection, or a new one.
    fn take(self: &Arc<Self>) -> Result<Pooled, AppError> {
        let idle = self.idle.lock().unwrap_or_else(|e| e.into_inner()).pop();
        let conn = match idle {
            Some(conn) => conn,
            None => open(&self.path)?,
        };
        Ok(Pooled {
            conn: Some(conn),
            pool: self.clone(),
        })
    }

    #[cfg(test)]
    fn idle(&self) -> usize {
        self.idle.lock().unwrap().len()
    }
}

/// The pools, one per store path.
fn pools() -> &'static Mutex<HashMap<PathBuf, Arc<Pool>>> {
    static POOLS: OnceLock<Mutex<HashMap<PathBuf, Arc<Pool>>>> = OnceLock::new();
    POOLS.get_or_init(|| Mutex::new(HashMap::new()))
}

/// A read connection to the store at `path`, from its pool.
pub fn get(path: &Path) -> Result<Pooled, AppError> {
    let pool = {
        let mut pools = pools().lock().unwrap_or_else(|e| e.into_inner());
        pools
            .entry(path.to_path_buf())
            .or_insert_with(|| Arc::new(Pool::new(path)))
            .clone()
    };
    pool.take()
}

/// Open a read connection the way commonmeta's own reads do — read-only, a
/// 30 s busy timeout, an immutable open where the WAL index can't be written —
/// and tune it for serving.
fn open(path: &Path) -> Result<Connection, AppError> {
    let conn = commonmeta::open_readonly(path)
        .map_err(|e| AppError::Internal(format!("sqlite open '{}': {e}", path.display())))?;
    // WAL admits readers alongside one writer, but a read still takes a read mark
    // in the wal-index, and that collides with WAL recovery and with a checkpoint
    // — over a large WAL left behind by an interrupted writer, a long window. The
    // default timeout is 0, so such a read failed outright with "database is
    // locked" rather than waiting a moment. commonmeta's open sets 30 s; it is
    // set again here so the guarantee does not rest on another crate's default.
    conn.busy_timeout(std::time::Duration::from_secs(30))
        .map_err(|e| {
            AppError::Internal(format!("sqlite busy_timeout '{}': {e}", path.display()))
        })?;
    // Map the full database into virtual address space so the kernel's page
    // cache (not SQLite's tiny internal cache) serves all index lookups via
    // page-faults. On a server with 128 GB RAM the hot indexes stay resident.
    // Virtual address space on 64-bit Linux is ~128 TB so a large mmap_size
    // costs nothing until pages are actually touched.
    // SQLite page cache is kept small because mmap bypasses it for mapped pages.
    conn.execute_batch(
        "PRAGMA mmap_size = 549755813888; PRAGMA cache_size = -8192; PRAGMA temp_store = MEMORY;",
    )
    .ok();
    Ok(conn)
}

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

    fn store() -> (tempfile::TempDir, PathBuf) {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("store.sqlite3");
        Connection::open(&path)
            .unwrap()
            .execute_batch("CREATE TABLE t (x INTEGER); INSERT INTO t VALUES (1);")
            .unwrap();
        (dir, path)
    }

    #[test]
    fn a_connection_is_given_back_and_taken_again() {
        let (_dir, path) = store();
        let first = get(&path).unwrap();
        // A pragma lives with its connection: set on this one, it marks it.
        first.execute_batch("PRAGMA cache_size = -1234").unwrap();
        let pool = first.pool.clone();
        drop(first);
        assert_eq!(pool.idle(), 1);
        // The same connection comes back, not a new one.
        let again = get(&path).unwrap();
        let mark: i64 = again
            .query_row("PRAGMA cache_size", [], |r| r.get(0))
            .unwrap();
        assert_eq!(mark, -1234);
        assert_eq!(pool.idle(), 0);
        let x: i64 = again
            .query_row("SELECT x FROM t", [], |r| r.get(0))
            .unwrap();
        assert_eq!(x, 1);
    }

    #[test]
    fn a_burst_opens_more_and_the_pool_keeps_at_most_its_cap() {
        let (_dir, path) = store();
        let held: Vec<Pooled> = (0..MAX_IDLE + 3).map(|_| get(&path).unwrap()).collect();
        let pool = held[0].pool.clone();
        drop(held);
        assert_eq!(pool.idle(), MAX_IDLE);
    }

    #[test]
    fn a_pooled_connection_is_read_only_and_waits_for_locks() {
        let (_dir, path) = store();
        let conn = get(&path).unwrap();
        assert!(conn.execute("INSERT INTO t VALUES (2)", []).is_err());
        let timeout: i64 = conn
            .query_row("PRAGMA busy_timeout", [], |r| r.get(0))
            .unwrap();
        assert_eq!(timeout, 30_000);
    }
}