ctoolbox/workspace/ipc_old/dispatch/

db.rs

//! Low-level database operations. Not a public API.
//! Locking is basically untested and probably flaky. Use at your own risk.

use anyhow::{Result, anyhow};

use crate::storage::get_storage_dir;
use crate::utilities::resource_lock::ResourceLock;
use crate::{debug, warn};

// NOTE: adjust the import below to wherever ResourceLock actually lives in your crate.

use redb::{Database, ReadableDatabase, ReadableTable, TableDefinition};

use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, OnceLock, RwLock};

/// The kind of database we are opening. Wrapped is the default mode that we
/// will later extend with compression/encryption + buffered in-memory storage.
/// Unwrapped (`_u`) is always a plain redb file with no extra wrapping.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum DbKind {
    Wrapped,
    Unwrapped,
}

/// Global, thread-safe pool of lazily-opened databases, keyed by (`DbKind`, name).
///
/// Note: We store `Arc<Database>` here so multiple `TableConnections` can
/// cheaply share the same underlying database handle within the process.
static DB_POOL: OnceLock<RwLock<HashMap<(DbKind, String), Arc<Database>>>> =
    OnceLock::new();

fn db_pool() -> &'static RwLock<HashMap<(DbKind, String), Arc<Database>>> {
    DB_POOL.get_or_init(|| RwLock::new(HashMap::new()))
}

// -------------------------------------------------------------------------------------
// Cross-process locking
// -------------------------------------------------------------------------------------

// Global DB open guard to serialize DB open/creation when desired.
// IMPORTANT: Only acquire this lock around DB create/open calls and NEVER
// acquire other ResourceLocks while holding it (to avoid deadlock via lock inversion).
/*const DB_GLOBAL_LOCK: &str = "db_global";

fn db_global_lock() -> Result<ResourceLock> {
    ResourceLock::acquire(DB_GLOBAL_LOCK, &"global")
}*/

// Per-database (per redb file) lock family. This isolates processes by a lock file
// and is thread-reentrant (same thread can re-enter).
const DB_FILE_LOCK_FAMILY: &str = "db_file";

fn db_file_lock(name: &str) -> Result<ResourceLock> {
    // Use the database name as the lock identity within the family.
    // This ensures exclusive access to a given redb file across processes.
    ResourceLock::acquire(DB_FILE_LOCK_FAMILY, &name.replace('/', "_"))
}

// A registry to hold long-lived lock sessions for IPC-style lock/unlock.
static LOCK_SESSIONS: OnceLock<RwLock<HashMap<String, Vec<ResourceLock>>>> =
    OnceLock::new();

fn lock_sessions() -> &'static RwLock<HashMap<String, Vec<ResourceLock>>> {
    LOCK_SESSIONS.get_or_init(|| RwLock::new(HashMap::new()))
}

static NEXT_LOCK_ID: AtomicU64 = AtomicU64::new(1);

fn new_lock_id() -> String {
    let id = NEXT_LOCK_ID.fetch_add(1, Ordering::Relaxed);
    format!("db-lock-{id}")
}

/// Acquire locks for one or more databases and return an opaque session ID to
/// hold those locks until explicitly released. This is suitable for use over
/// an IPC boundary where a separate "database process" manages locks on behalf
/// of clients.
///
/// - Locks are acquired in a stable sorted order to avoid deadlocks.
/// - If any lock acquisition fails, all acquired locks are released and an error
///   is returned.
/// - Re-entrant calls from the same thread are supported by `ResourceLock`.
///
/// Returns a `session_id` string which must be used to release the locks.
///
/// Example usage from an IPC handler:
///   let session = `lock_databases_session`(&["users/auth", "users/uuids"])?;
///   // ... do work ...
///   `unlock_databases_session(&session)`?;
pub fn lock_databases_session<I, S>(names: I) -> Result<String>
where
    I: IntoIterator<Item = S>,
    S: AsRef<str>,
{
    // Normalize and sort names to avoid deadlock from inconsistent acquisition order.
    let mut dbs: Vec<String> =
        names.into_iter().map(|s| s.as_ref().to_string()).collect();
    dbs.sort();

    let mut acquired: Vec<ResourceLock> = Vec::with_capacity(dbs.len());
    for db in &dbs {
        match db_file_lock(db) {
            Ok(lock) => acquired.push(lock),
            Err(e) => {
                // Drop all previously acquired locks before returning error.
                drop(acquired);
                return Err(e);
            }
        }
    }

    let session_id = new_lock_id();
    let mut sessions = lock_sessions().write().expect("LOCK_SESSIONS poisoned");
    sessions.insert(session_id.clone(), acquired);
    Ok(session_id)
}

/// Release a previously created lock session by its session ID. All locks
/// associated with the session are released (dropped).
pub fn unlock_databases_session(session_id: &str) -> Result<()> {
    let mut sessions = lock_sessions().write().expect("LOCK_SESSIONS poisoned");
    if let Some(locks) = sessions.remove(session_id) {
        drop(locks);
        Ok(())
    } else {
        Err(anyhow!("Unknown lock session id: {session_id}"))
    }
}

/// Convenience: lock a single database (redb file) and return a session id.
/// Use `unlock_databases_session` to release.
pub fn lock_database_session(name: &str) -> Result<String> {
    lock_databases_session([name])
}

// -------------------------------------------------------------------------------------
// Database open/get/put/delete with locking
// -------------------------------------------------------------------------------------

/// Open or fetch from the pool a Database for the given name and kind.
///
/// Today both kinds share the same on-disk implementation. In the future,
/// `DbKind::Wrapped` can be switched to use an additional in-memory layer with
/// compression/encryption and periodic flushes to disk while keeping the same
/// pooling API surface.
fn get_or_open_database(name: &str, kind: DbKind) -> Result<Arc<Database>> {
    // Fast path: check if already in pool
    if let Some(db) = db_pool()
        .read()
        .expect("DB_POOL poisoned")
        .get(&(kind, name.to_string()))
        .cloned()
    {
        return Ok(db);
    }

    // Slow path: serialize open/creation of this specific database across processes.
    // We prefer a per-database lock so other databases can still open concurrently.
    let _db_file_guard = db_file_lock(name)?;

    // Re-check after acquiring the lock in case another thread in this process raced us.
    if let Some(db) = db_pool()
        .read()
        .expect("DB_POOL poisoned")
        .get(&(kind, name.to_string()))
        .cloned()
    {
        return Ok(db);
    }

    // Optionally, if you prefer to serialize ALL opens across all DBs, uncomment:
    // let _global = db_global_lock()?;

    // Open/create database
    let db = open_redb_file(name)?;
    let arc = Arc::new(db);

    // Insert into pool (double-checked in case of a race)
    let mut map = db_pool().write().expect("DB_POOL poisoned");
    let entry = map
        .entry((kind, name.to_string()))
        .or_insert_with(|| arc.clone());
    Ok(entry.clone())
}

/// Create or open a redb database file on disk (no wrapping).
fn open_redb_file(name: &str) -> Result<Database> {
    let path_to_database = db_path(name)?;
    // Try to create if it doesn't exist, tolerating races
    if !path_to_database.exists() {
        // Serialize the create/open on this file to avoid multiple concurrent creates across processes.
        let _db_file_guard = db_file_lock(name)?;
        if !path_to_database.exists() {
            let _ = Database::create(&path_to_database);
        }
    }
    let db = Database::open(path_to_database)?;
    Ok(db)
}

fn db_path(name: &str) -> Result<PathBuf> {
    Ok(get_storage_dir()?.join(format!("{name}.redb")))
}

/// Open a thread-safe, but not multi-process-safe, connection to a specific table in a wrapped redb database.
///
/// Wrapped databases are pooled and will later support in-memory buffering with
/// compression/encryption and periodic debounced flushing to disk.
pub fn open<K, V>(table_name: &str) -> Result<TableConnection<K, V>>
where
    K: redb::Key + Sized + 'static,
    V: redb::Value + Sized + 'static,
{
    let conn = TableConnection::open(table_name)?;
    Ok(conn)
}

/// Open a thread-safe, but not multi-process-safe, connection to a specific table in an unwrapped redb database.
///
/// Unwrapped databases are also pooled, but will always be plain on-disk redb files.
pub fn open_u<K, V>(table_name: &str) -> Result<TableConnection<K, V>>
where
    K: redb::Key + Sized + 'static,
    V: redb::Value + Sized + 'static,
{
    let conn = TableConnection::open_u(table_name)?;
    Ok(conn)
}

/// A connection to a specific table in a redb database.
pub struct TableConnection<K, V>
where
    K: redb::Key + Sized + 'static,
    V: redb::Value + Sized + 'static,
{
    // NOTE: Arc-wrapped and pooled Database handle
    db: Arc<Database>,
    table_def: TableDefinition<'static, K, V>,
    table_name: String,
}

const UNWRAPPED_TABLES: [&str; 8] = [
    "users/auth",
    "users/ids",
    "users/ids_rev",
    "users/key_encryption_key_params",
    "users/pictures",
    "users/pubkeys",
    "users/uuids",
    "users/wrapped_dek",
];

impl<K, V> TableConnection<K, V>
where
    K: redb::Key + Sized + 'static,
    V: redb::Value + Sized + 'static,
{
    // `table_name` is leaked to satisfy redb's 'static requirement for TableDefinition.
    fn open_table(db: Arc<Database>, table_name: &str) -> Result<Self> {
        let leaked: &'static str =
            Box::leak(table_name.to_string().into_boxed_str());
        let table_def = TableDefinition::new(leaked);
        Ok(Self {
            db,
            table_def,
            table_name: table_name.to_string(),
        })
    }

    /// Open pooled unwrapped database and create `TableConnection`.
    fn open_u(table_name: &str) -> Result<Self> {
        let db = get_or_open_database(table_name, DbKind::Unwrapped)?;
        Self::open_table(db, table_name)
    }

    /// Open pooled (TODO: wrapped) database and create `TableConnection`.
    pub fn open(table_name: &str) -> Result<Self> {
        if UNWRAPPED_TABLES.contains(&table_name) {
            return Self::open_u(table_name);
        }
        let db = get_or_open_database(table_name, DbKind::Wrapped)?;
        Self::open_table(db, table_name)
    }

    /// Acquire the per-database cross-process lock for this connection's database.
    /// Returns a guard that releases on drop.
    fn acquire_db_lock(&self) -> Result<ResourceLock> {
        db_file_lock(&self.table_name)
    }

    /// Read-only get: returns an owned R produced by mapper while txn is alive.
    ///
    /// This operation is serialized across processes per database via `ResourceLock`.
    fn get<'k, R, F>(
        &self,
        key: <K as redb::Value>::SelfType<'k>,
        map: F,
    ) -> Option<R>
    where
        F: for<'v> FnOnce(<V as redb::Value>::SelfType<'v>) -> R,
    {
        // Serialize this DB operation across processes.
        let _lock = match self.acquire_db_lock() {
            Ok(guard) => guard,
            Err(e) => {
                warn!(format!(
                    "db: GET lock acquisition failed for {}: {e}: {e:?}",
                    self.table_name
                ));
                return None;
            }
        };

        let tx = self.db.begin_read().ok()?;
        let table = tx.open_table(self.table_def).ok()?;
        table.get(key).ok()?.map(|acc| map(acc.value()))
    }

    /// If V is &[u8] or similar, get and convert to `Vec<u8>`
    pub fn get_vec<'k>(
        &self,
        key: <K as redb::Value>::SelfType<'k>,
    ) -> Option<Vec<u8>>
    where
        V: redb::Value,
        for<'v> <V as redb::Value>::SelfType<'v>: AsRef<[u8]>,
    {
        let key_temp = format!("{key:?}");
        let res = self.get(key, |x| x.as_ref().to_vec());
        debug!(format!(
            "db: GET {}/{:?} -> {:?} (Vec<u8>)",
            self.table_name, key_temp, res
        ));
        res
    }

    /// If V is &str or similar, get and convert to String
    pub fn get_str<'k>(
        &self,
        key: <K as redb::Value>::SelfType<'k>,
    ) -> Option<String>
    where
        V: redb::Value,
        for<'v> <V as redb::Value>::SelfType<'v>: AsRef<str>,
    {
        let key_temp = format!("{key:?}");
        let res = self.get(key, |x| x.as_ref().to_string());
        debug!(format!(
            "db: GET {}/{:?} -> {:?} (str)",
            self.table_name, key_temp, res
        ));
        res
    }

    /// If V is u64, get and return it directly
    pub fn get_u64<'k>(
        &self,
        key: <K as redb::Value>::SelfType<'k>,
    ) -> Option<u64>
    where
        V: redb::Value,
        for<'a> <V as redb::Value>::SelfType<'a>: Into<u64>,
    {
        let key_temp = format!("{key:?}");
        let res = self.get(key, |x: <V as redb::Value>::SelfType<'_>| x.into());
        debug!(format!(
            "db: GET {}/{:?} -> {:?} (u64)",
            self.table_name, key_temp, res
        ));
        res
    }

    /// Put: opens a write txn and inserts value.
    ///
    /// This operation is serialized across processes per database via `ResourceLock`.
    pub fn put<'k>(
        &self,
        key: <K as redb::Value>::SelfType<'_>,
        value: <V as redb::Value>::SelfType<'_>,
    ) -> Result<()>
    where
        K: redb::Key + Sized,
        V: redb::Value + Sized,
    {
        // Serialize this DB operation across processes.
        let _lock = self.acquire_db_lock()?;

        debug!(format!(
            "db: PUT {}/{:?} -> {:?}",
            self.table_name, key, value
        ));
        let tx = self.db.begin_write()?;
        {
            let mut table = tx.open_table(self.table_def)?;
            table.insert(key, value)?;
        }
        tx.commit()?;
        Ok(())
    }

    /// Delete
    ///
    /// This operation is serialized across processes per database via `ResourceLock`.
    pub fn delete(&self, key: <K as redb::Value>::SelfType<'_>) -> Result<()> {
        // Serialize this DB operation across processes.
        let _lock = self.acquire_db_lock()?;

        debug!(format!("db: DELETE {}/{:?}", self.table_name, key));
        let tx = self.db.begin_write()?;
        {
            let mut table = tx.open_table(self.table_def)?;
            table.remove(key)?;
        }
        tx.commit()?;
        Ok(())
    }
}

/// Get a u64 value by string key.
pub fn get_str_u64(db: &str, key: &str) -> Option<u64> {
    let conn: TableConnection<&str, u64> = TableConnection::open(db).ok()?;
    conn.get_u64(key)
}

pub fn get_all_u64_keys(db: &str) -> Result<Vec<u64>> {
    let conn: TableConnection<u64, &str> = TableConnection::open(db)?;
    // Serialize this DB operation across processes.
    let _lock = db_file_lock(db)?;

    let tx = conn.db.begin_read()?;
    let table = tx.open_table(conn.table_def)?;
    let mut keys = Vec::new();
    if let Ok(entry_result) = table.iter() {
        for entry in entry_result.flatten() {
            keys.push(entry.0.value().to_owned());
        }
    } else {
        warn!("Could not get iter of users/uuids table");
        return Err(anyhow::anyhow!("Could not get iter of users/uuids table"));
    }
    Ok(keys)
}

/// Put a u64 value by string key.
pub fn put_str_u64(db: &str, key: &str, value: u64) -> Result<()> {
    let conn: TableConnection<&str, u64> = TableConnection::open(db)?;
    conn.put(key, value)
}

/// Delete a u64 value by string key.
pub fn delete_str_u64(db: &str, key: &str) -> Result<()> {
    let conn: TableConnection<&str, u64> = TableConnection::open(db)?;
    conn.delete(key)
}

/// Get a string value by u64 key.
pub fn get_u64_str(db: &str, key: u64) -> Option<String> {
    let conn: TableConnection<u64, &str> = TableConnection::open(db).ok()?;
    conn.get_str(key)
}

/// Put a string value by u64 key.
pub fn put_u64_str(db: &str, key: u64, value: &str) -> Result<()> {
    let conn: TableConnection<u64, &str> = TableConnection::open(db)?;
    conn.put(key, value)
}

/// Delete a string value by u64 key.
pub fn delete_u64_str(db: &str, key: u64) -> Result<()> {
    let conn: TableConnection<u64, &str> = TableConnection::open(db)?;
    conn.delete(key)
}

/// Get a byte vector by u64 key.
pub fn get_u64_bytes(db: &str, key: u64) -> Option<Vec<u8>> {
    let conn: TableConnection<u64, &[u8]> = TableConnection::open(db).ok()?;
    conn.get_vec(key)
}

/// Put a byte slice by u64 key.
pub fn put_u64_bytes(db: &str, key: u64, value: &[u8]) -> Result<()> {
    let conn: TableConnection<u64, &[u8]> = TableConnection::open(db)?;
    conn.put(key, value)
}

/// Delete a byte slice by u64 key.
pub fn delete_u64_bytes(db: &str, key: u64) -> Result<()> {
    let conn: TableConnection<u64, &[u8]> = TableConnection::open(db)?;
    conn.delete(key)
}

#[cfg(test)]
#[allow(clippy::unwrap_in_result, clippy::panic_in_result_fn)]
mod tests {
    //! Tests for get/put/delete helpers for redb database.

    use super::*;
    use anyhow::Result;

    fn test_db_name(suffix: &str) -> String {
        format!("test_db_{suffix}")
    }

    #[crate::ctb_test]
    fn test_str_u64_roundtrip() -> Result<()> {
        let db = test_db_name("str_u64");
        let key = "foo";
        let value = 12345u64;

        // Acquire lock session explicitly (simulating IPC usage)
        let session = lock_database_session(&db)?;

        put_str_u64(&db, key, value)?;
        assert_eq!(get_str_u64(&db, key), Some(value));

        delete_str_u64(&db, key)?;
        assert_eq!(get_str_u64(&db, key), None);

        // Release session
        unlock_databases_session(&session)?;

        std::fs::remove_file(db_path(&db)?.as_path()).ok();

        Ok(())
    }

    #[crate::ctb_test]
    fn test_u64_str_roundtrip() -> Result<()> {
        let db = test_db_name("u64_str");
        let key = 42u64;
        let value = "bar";

        put_u64_str(&db, key, value)?;
        assert_eq!(get_u64_str(&db, key), Some(value.to_string()));

        delete_u64_str(&db, key)?;
        assert_eq!(get_u64_str(&db, key), None);

        std::fs::remove_file(db_path(&db)?.as_path()).ok();

        Ok(())
    }

    #[crate::ctb_test]
    fn test_u64_bytes_roundtrip() -> Result<()> {
        let db = test_db_name("u64_bytes");
        let key = 99u64;
        let value = vec![1u8, 2, 3, 4];

        put_u64_bytes(&db, key, &value)?;
        assert_eq!(get_u64_bytes(&db, key), Some(value.clone()));

        delete_u64_bytes(&db, key)?;
        assert_eq!(get_u64_bytes(&db, key), None);

        std::fs::remove_file(db_path(&db)?.as_path()).ok();

        Ok(())
    }
}