char/avril

git clone https://git.t4t.associates/char/avril

Charlotte Somadd `avril:kv` module78f47482026-10-21

use std::io;
use std::ops::Bound;
use std::path::Path;
use std::sync::{Arc, Mutex, MutexGuard};

use redb::{
    Database, DatabaseError, ReadOnlyDatabase, ReadableDatabase, ReadableTable,
    ReadableTableMetadata, TableDefinition, TableHandle,
};

// Keyspace `name` lives in table "ks:{name}", mapping key bytes to (commit
// counter, value bytes). The prefix keeps keyspaces apart from META, which
// holds the commit counter under "" and adapter metadata under any other key.
// Not "meta": pre-unification Web Storage files have a table by that name.
const META: TableDefinition<&str, u64> = TableDefinition::new("kv_meta");
const COUNTER_KEY: &str = "";
const KEYSPACE_PREFIX: &str = "ks:";

type Bytes = &'static [u8];
type Entry = (u64, Bytes);

fn keyspace_table_name(name: &str) -> String {
    format!("{KEYSPACE_PREFIX}{name}")
}

pub fn stamp_string(counter: u64) -> String {
    counter.to_string()
}

pub fn parse_stamp(stamp: &str) -> io::Result<u64> {
    // Only the canonical rendering parses, so equal stamps are equal strings.
    // (u64::from_str alone would also accept a leading '+' or zeros.)
    if !stamp.bytes().all(|b| b.is_ascii_digit()) || (stamp.len() > 1 && stamp.starts_with('0')) {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "Invalid kv versionstamp",
        ));
    }
    stamp
        .parse()
        .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "Invalid kv versionstamp"))
}

pub fn check_keyspace_name(name: &str) -> io::Result<()> {
    if name.is_empty() {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "Keyspace name must not be empty",
        ));
    }
    Ok(())
}

// "" is the commit counter's slot, so adapter keys cannot touch it.
fn check_meta_key(key: &str) -> io::Result<()> {
    if key.is_empty() {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "kv metadata keys must be non-empty",
        ));
    }
    Ok(())
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Access {
    Shared,
    Exclusive,
    ReadOnly,
}

impl Access {
    pub fn parse(value: &str) -> io::Result<Self> {
        match value {
            "shared" => Ok(Access::Shared),
            "exclusive" => Ok(Access::Exclusive),
            "read-only" => Ok(Access::ReadOnly),
            _ => Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "access must be \"shared\", \"exclusive\" or \"read-only\"",
            )),
        }
    }

    pub fn writes(self) -> bool {
        self != Access::ReadOnly
    }
}

enum Held {
    // Shared mode never holds the file: every operation opens and closes it,
    // serialized by the store mutex so same-process opens cannot collide.
    Shared { path: std::path::PathBuf },
    Exclusive { db: Database },
    // A missing file stays unopened so read-only never creates anything;
    // reads then behave as empty.
    ReadOnly { db: Option<ReadOnlyDatabase> },
}

// A stored entry: raw value bytes plus the opaque versionstamp of the
// commit that wrote it.
pub type StoredEntry = (Vec<u8>, String);

// A listed entry: key plus StoredEntry.
pub type ListedEntry = (Vec<u8>, Vec<u8>, String);

/// Outcome of a [`Store::write_tx`] closure: commit the transaction or drop
/// it uncommitted. Aborting changes nothing, including keyspace registration.
pub enum Decision<T> {
    Commit(T),
    Abort(T),
}

struct StoreInner {
    // Serializes all operations on this database. Shared mode relies on this
    // for same-process safety; clones of one Store share the same mutex, so
    // closures must not re-enter the Store or they will deadlock.
    mutex: Mutex<()>,
    held: Held,
}

fn lock(mutex: &Mutex<()>) -> MutexGuard<'_, ()> {
    mutex.lock().unwrap_or_else(|error| error.into_inner())
}

/// Synchronous byte-oriented redb backend. Cheap to clone and share across
/// threads; every operation runs to completion under one mutex.
#[derive(Clone)]
pub struct Store {
    inner: Arc<StoreInner>,
}

impl std::fmt::Debug for Store {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Store").finish_non_exhaustive()
    }
}

fn db_error(path: &Path, error: DatabaseError) -> io::Error {
    match error {
        DatabaseError::DatabaseAlreadyOpen => io::Error::new(
            io::ErrorKind::AlreadyExists,
            format!("kv database is already open: {}", path.display()),
        ),
        DatabaseError::Storage(redb::StorageError::Io(error))
            if error.kind() == io::ErrorKind::NotFound =>
        {
            io::Error::new(
                io::ErrorKind::NotFound,
                format!("kv database does not exist: {}", path.display()),
            )
        }
        error => io::Error::other(format!("kv database at {}: {error}", path.display())),
    }
}

fn is_empty_file(path: &Path) -> bool {
    std::fs::metadata(path).is_ok_and(|metadata| metadata.len() == 0)
}

fn create_database(path: &Path) -> Result<Database, DatabaseError> {
    match Database::create(path) {
        // A crash between creating the file and initializing it leaves an
        // empty file; retry from scratch rather than reporting corruption.
        Err(_) if is_empty_file(path) => {
            let _ = std::fs::remove_file(path);
            Database::create(path)
        }
        result => result,
    }
}

fn open_database(path: &Path, create: bool) -> io::Result<Database> {
    if let Some(parent) = path.parent()
        && !parent.as_os_str().is_empty()
        && !parent.exists()
    {
        if !create {
            return Err(io::Error::new(
                io::ErrorKind::NotFound,
                format!("kv database does not exist: {}", path.display()),
            ));
        }
        return Err(io::Error::new(
            io::ErrorKind::NotFound,
            format!(
                "kv database parent directory does not exist: {}",
                parent.display()
            ),
        ));
    }
    if path.is_dir() {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("kv database path is a directory: {}", path.display()),
        ));
    }
    let result = if create {
        create_database(path)
    } else if !path.exists() || is_empty_file(path) {
        return Err(io::Error::new(
            io::ErrorKind::NotFound,
            format!("kv database does not exist: {}", path.display()),
        ));
    } else {
        Database::open(path)
    };
    result.map_err(|error| db_error(path, error))
}

// Missing or empty files stay unopened so read-only never creates anything.
fn read_only_database(path: &Path) -> Result<Option<ReadOnlyDatabase>, DatabaseError> {
    if !path.exists() || is_empty_file(path) {
        return Ok(None);
    }
    ReadOnlyDatabase::open(path).map(Some)
}

impl Store {
    pub fn open(path: &Path, access: Access, create: bool) -> io::Result<Self> {
        // Lock conflicts surface as DatabaseAlreadyOpen via db_error; redb
        // locks are non-blocking, so there is no spin or wait here.
        let held = match access {
            Access::ReadOnly if create => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    "read-only kv databases cannot be created",
                ));
            }
            Access::Shared => {
                // Probe once so open() itself reports missing, locked or
                // corrupt files instead of deferring the failure to the first
                // operation.
                open_database(path, create)?;
                Held::Shared { path: path.to_owned() }
            }
            Access::Exclusive => Held::Exclusive { db: open_database(path, create)? },
            Access::ReadOnly => match read_only_database(path) {
                Err(DatabaseError::RepairAborted) => {
                    return Err(io::Error::other(format!(
                        "kv database at {} needs repair; open it with write access first",
                        path.display()
                    )));
                }
                result => Held::ReadOnly {
                    db: result.map_err(|error| db_error(path, error))?,
                },
            },
        };
        Ok(Store::new(held))
    }

    /// Open read-only, repairing an unclean shutdown first. ReadOnlyDatabase
    /// cannot repair on its own, so a RepairAborted open is followed by
    /// exactly one writable open (which repairs) and a single read-only
    /// retry; if the file still needs repair afterwards, that error is
    /// returned. Missing or empty files open as empty without creating
    /// anything, like [`Store::open`] with [`Access::ReadOnly`].
    pub fn open_read_only_repairing(path: &Path) -> io::Result<Self> {
        let db = match read_only_database(path) {
            Err(DatabaseError::RepairAborted) => {
                Database::open(path).map_err(|error| db_error(path, error))?;
                read_only_database(path)
            }
            result => result,
        };
        Ok(Store::new(Held::ReadOnly {
            db: db.map_err(|error| db_error(path, error))?,
        }))
    }

    fn new(held: Held) -> Self {
        Store {
            inner: Arc::new(StoreInner { mutex: Mutex::new(()), held }),
        }
    }
}

/// Read-only view of one transaction. Constructed only by [`Store::read_tx`].
pub struct ReadTx<'a> {
    // None for a database file that does not exist: every read is empty.
    txn: Option<&'a redb::ReadTransaction>,
}

/// Read-write view of one transaction. Constructed only by
/// [`Store::write_tx`]; entry mutations allocate one commit stamp per
/// transaction, lazily on the first mutation, so a batch of writes shares
/// exactly one stamp.
pub struct WriteTx<'a> {
    txn: &'a redb::WriteTransaction,
    counter: Option<u64>,
}

fn commit_counter(meta: &redb::Table<'_, &str, u64>) -> io::Result<u64> {
    Ok(meta
        .get(COUNTER_KEY)
        .map_err(io::Error::other)?
        .map_or(0, |counter| counter.value()))
}

// Allocates the transaction's stamp on first use. Entry writes share one
// counter bump per transaction; metadata-only work never allocates one.
fn bump_counter(meta: &mut redb::Table<'_, &str, u64>, slot: &mut Option<u64>) -> io::Result<u64> {
    if let Some(counter) = *slot {
        return Ok(counter);
    }
    let counter = commit_counter(meta)? + 1;
    meta.insert(COUNTER_KEY, counter)
        .map_err(io::Error::other)?;
    *slot = Some(counter);
    Ok(counter)
}

fn read_meta(meta: &impl ReadableTable<&'static str, u64>, key: &str) -> io::Result<Option<u64>> {
    check_meta_key(key)?;
    Ok(meta
        .get(key)
        .map_err(io::Error::other)?
        .map(|value| value.value()))
}

// The entry type is dictated by redb's range iterator.
#[allow(clippy::type_complexity)]
fn decode_entry(
    entry: Result<(redb::AccessGuard<'_, Bytes>, redb::AccessGuard<'_, Entry>), redb::StorageError>,
) -> io::Result<ListedEntry> {
    let (key, stored) = entry.map_err(io::Error::other)?;
    let (counter, value) = stored.value();
    Ok((key.value().to_vec(), value.to_vec(), stamp_string(counter)))
}

// Exclusive end bound for a prefix: increment the last non-0xff byte. A
// prefix of all 0xff has no successor, so the range stays unbounded above.
fn prefix_successor(prefix: &[u8]) -> Option<Vec<u8>> {
    let mut end = prefix.to_vec();
    while end.last() == Some(&0xff) {
        end.pop();
    }
    let last = end.last_mut()?;
    *last += 1;
    Some(end)
}

const MAX_LIST: usize = 10_000;

fn open_or_absent<T>(result: Result<T, redb::TableError>) -> io::Result<Option<T>> {
    match result {
        Ok(table) => Ok(Some(table)),
        Err(redb::TableError::TableDoesNotExist(_)) => Ok(None),
        Err(error) => Err(io::Error::other(error)),
    }
}

// The transaction kinds ReadTx and WriteTx read through, so each read is
// written once. An absent table reads as empty and is never created.
trait Tables {
    fn table_names(&self) -> io::Result<Vec<String>>;

    type Table<'t, K: redb::Key + 'static, V: redb::Value + 'static>: ReadableTable<K, V>
    where
        Self: 't;

    fn table<K: redb::Key + 'static, V: redb::Value + 'static>(
        &self,
        definition: TableDefinition<K, V>,
    ) -> io::Result<Option<Self::Table<'_, K, V>>>;
}

impl Tables for Option<&redb::ReadTransaction> {
    fn table_names(&self) -> io::Result<Vec<String>> {
        let Some(txn) = self else {
            return Ok(Vec::new());
        };
        let tables = txn.list_tables().map_err(io::Error::other)?;
        Ok(tables.map(|table| table.name().to_owned()).collect())
    }

    type Table<'t, K: redb::Key + 'static, V: redb::Value + 'static>
        = redb::ReadOnlyTable<K, V>
    where
        Self: 't;

    fn table<K: redb::Key + 'static, V: redb::Value + 'static>(
        &self,
        definition: TableDefinition<K, V>,
    ) -> io::Result<Option<redb::ReadOnlyTable<K, V>>> {
        match self {
            Some(txn) => open_or_absent(txn.open_table(definition)),
            None => Ok(None),
        }
    }
}

impl Tables for redb::WriteTransaction {
    type Table<'t, K: redb::Key + 'static, V: redb::Value + 'static> = redb::Table<'t, K, V>;

    fn table_names(&self) -> io::Result<Vec<String>> {
        let tables = self.list_tables().map_err(io::Error::other)?;
        Ok(tables.map(|table| table.name().to_owned()).collect())
    }

    // open_table creates a missing table, so check first: reads, checks and
    // deletes inside a write transaction must not materialize anything.
    fn table<K: redb::Key + 'static, V: redb::Value + 'static>(
        &self,
        definition: TableDefinition<K, V>,
    ) -> io::Result<Option<redb::Table<'_, K, V>>> {
        if !self
            .table_names()?
            .iter()
            .any(|name| name == definition.name())
        {
            return Ok(None);
        }
        self.open_table(definition)
            .map(Some)
            .map_err(io::Error::other)
    }
}

fn keyspace_table<'t, T: Tables>(
    txn: &'t T,
    keyspace: &str,
) -> io::Result<Option<T::Table<'t, Bytes, Entry>>> {
    check_keyspace_name(keyspace)?;
    txn.table(TableDefinition::new(&keyspace_table_name(keyspace)))
}

fn get_entry(txn: &impl Tables, keyspace: &str, key: &[u8]) -> io::Result<Option<StoredEntry>> {
    let Some(table) = keyspace_table(txn, keyspace)? else {
        return Ok(None);
    };
    let Some(entry) = table.get(key).map_err(io::Error::other)? else {
        return Ok(None);
    };
    let (counter, value) = entry.value();
    Ok(Some((value.to_vec(), stamp_string(counter))))
}

fn scan(
    txn: &impl Tables,
    keyspace: &str,
    bounds: (Bound<&[u8]>, Bound<&[u8]>),
    limit: usize,
    reverse: bool,
) -> io::Result<Vec<ListedEntry>> {
    let Some(table) = keyspace_table(txn, keyspace)? else {
        return Ok(Vec::new());
    };
    let range = table.range::<&[u8]>(bounds).map_err(io::Error::other)?;
    let range: Box<dyn Iterator<Item = _>> = if reverse {
        Box::new(range.rev())
    } else {
        Box::new(range)
    };
    range.take(limit).map(decode_entry).collect()
}

fn entries(txn: &impl Tables, keyspace: &str) -> io::Result<Vec<ListedEntry>> {
    scan(
        txn,
        keyspace,
        (Bound::Unbounded, Bound::Unbounded),
        usize::MAX,
        false,
    )
}

fn count(txn: &impl Tables, keyspace: &str) -> io::Result<u64> {
    keyspace_table(txn, keyspace)?.map_or(Ok(0), |table| table.len().map_err(io::Error::other))
}

fn list_keyspaces(txn: &impl Tables) -> io::Result<Vec<String>> {
    let mut names: Vec<String> = txn
        .table_names()?
        .into_iter()
        .filter_map(|name| name.strip_prefix(KEYSPACE_PREFIX).map(str::to_owned))
        .collect();
    names.sort();
    Ok(names)
}

fn get_meta(txn: &impl Tables, key: &str) -> io::Result<Option<u64>> {
    txn.table(META)?
        .map_or(Ok(None), |meta| read_meta(&meta, key))
}

impl ReadTx<'_> {
    pub fn get(&self, keyspace: &str, key: &[u8]) -> io::Result<Option<StoredEntry>> {
        get_entry(&self.txn, keyspace, key)
    }

    pub fn entries(&self, keyspace: &str) -> io::Result<Vec<ListedEntry>> {
        entries(&self.txn, keyspace)
    }

    /// Entries within `[start, end)` and under `prefix`, at most `limit`
    /// (capped at 10000) of them.
    #[allow(clippy::too_many_arguments)]
    pub fn list(
        &self,
        keyspace: &str,
        start: Option<&[u8]>,
        end: Option<&[u8]>,
        prefix: Option<&[u8]>,
        limit: usize,
        reverse: bool,
    ) -> io::Result<Vec<ListedEntry>> {
        // Intersect explicit bounds with the prefix range; the prefix end
        // is the successor exclusive bound (absent when all bytes are 0xff).
        let prefix_end = prefix.and_then(prefix_successor);
        let lower = match (start, prefix) {
            (Some(a), Some(b)) => Bound::Included(a.max(b)),
            (Some(a), None) => Bound::Included(a),
            (None, Some(b)) => Bound::Included(b),
            (None, None) => Bound::Unbounded,
        };
        let upper = match (end, prefix_end.as_deref()) {
            (Some(a), Some(b)) => Bound::Excluded(a.min(b)),
            (Some(a), None) => Bound::Excluded(a),
            (None, Some(b)) => Bound::Excluded(b),
            (None, None) => Bound::Unbounded,
        };
        scan(
            &self.txn,
            keyspace,
            (lower, upper),
            limit.min(MAX_LIST),
            reverse,
        )
    }

    pub fn count(&self, keyspace: &str) -> io::Result<u64> {
        count(&self.txn, keyspace)
    }

    pub fn list_keyspaces(&self) -> io::Result<Vec<String>> {
        list_keyspaces(&self.txn)
    }

    pub fn get_meta(&self, key: &str) -> io::Result<Option<u64>> {
        get_meta(&self.txn, key)
    }
}

impl WriteTx<'_> {
    pub fn get(&self, keyspace: &str, key: &[u8]) -> io::Result<Option<StoredEntry>> {
        get_entry(self.txn, keyspace, key)
    }

    pub fn put(&mut self, keyspace: &str, key: &[u8], value: &[u8]) -> io::Result<String> {
        check_keyspace_name(keyspace)?;
        let counter = {
            let mut meta = self.txn.open_table(META).map_err(io::Error::other)?;
            bump_counter(&mut meta, &mut self.counter)?
        };
        let name = keyspace_table_name(keyspace);
        let mut table = self
            .txn
            .open_table(TableDefinition::<Bytes, Entry>::new(&name))
            .map_err(io::Error::other)?;
        table
            .insert(key, (counter, value))
            .map_err(io::Error::other)?;
        Ok(stamp_string(counter))
    }

    pub fn delete(&mut self, keyspace: &str, key: &[u8]) -> io::Result<bool> {
        let Some(mut table) = keyspace_table(self.txn, keyspace)? else {
            return Ok(false);
        };
        if table.remove(key).map_err(io::Error::other)?.is_none() {
            return Ok(false);
        }
        drop(table);
        let mut meta = self.txn.open_table(META).map_err(io::Error::other)?;
        bump_counter(&mut meta, &mut self.counter)?;
        Ok(true)
    }

    /// Empty the keyspace, materializing it if it did not exist.
    pub fn clear(&mut self, keyspace: &str) -> io::Result<()> {
        check_keyspace_name(keyspace)?;
        let name = keyspace_table_name(keyspace);
        let definition = TableDefinition::<Bytes, Entry>::new(&name);
        self.txn
            .delete_table(definition)
            .map_err(io::Error::other)?;
        self.txn.open_table(definition).map_err(io::Error::other)?;
        Ok(())
    }

    pub fn entries(&self, keyspace: &str) -> io::Result<Vec<ListedEntry>> {
        entries(self.txn, keyspace)
    }

    pub fn count(&self, keyspace: &str) -> io::Result<u64> {
        count(self.txn, keyspace)
    }

    pub fn list_keyspaces(&self) -> io::Result<Vec<String>> {
        list_keyspaces(self.txn)
    }

    pub fn get_meta(&self, key: &str) -> io::Result<Option<u64>> {
        get_meta(self.txn, key)
    }

    pub fn put_meta(&mut self, key: &str, value: u64) -> io::Result<()> {
        check_meta_key(key)?;
        let mut meta = self.txn.open_table(META).map_err(io::Error::other)?;
        meta.insert(key, value).map_err(io::Error::other)?;
        Ok(())
    }

    /// This transaction's commit stamp, if an entry mutation allocated one.
    /// Check-only work leaves it empty; entry writes share exactly one stamp.
    pub fn stamp(&self) -> Option<String> {
        self.counter.map(stamp_string)
    }

    /// Allocate the commit stamp without writing an entry, for commits that
    /// must carry a stamp even though no entry changed.
    pub fn allocate_stamp(&mut self) -> io::Result<String> {
        let mut meta = self.txn.open_table(META).map_err(io::Error::other)?;
        Ok(stamp_string(bump_counter(&mut meta, &mut self.counter)?))
    }
}

impl Store {
    /// Run read-only work in one read transaction.
    pub fn read_tx<T>(&self, work: impl FnOnce(&ReadTx) -> io::Result<T>) -> io::Result<T> {
        let _guard = lock(&self.inner.mutex);
        let mut opened = None;
        let txn = match &self.inner.held {
            Held::Shared { path } => match open_database(path, false) {
                Ok(db) => Some(opened.insert(db).begin_read()),
                // A shared file removed since open reads as empty, like a
                // missing read-only file.
                Err(error) if error.kind() == io::ErrorKind::NotFound => None,
                Err(error) => return Err(error),
            },
            Held::Exclusive { db } => Some(db.begin_read()),
            Held::ReadOnly { db } => db.as_ref().map(ReadableDatabase::begin_read),
        };
        let txn = txn.transpose().map_err(io::Error::other)?;
        work(&ReadTx { txn: txn.as_ref() })
    }

    /// Run read-write work in one write transaction, committing on
    /// [`Decision::Commit`] and dropping everything uncommitted on
    /// [`Decision::Abort`].
    pub fn write_tx<T>(
        &self,
        work: impl FnOnce(&mut WriteTx) -> io::Result<Decision<T>>,
    ) -> io::Result<T> {
        let _guard = lock(&self.inner.mutex);
        let mut opened = None;
        let db = match &self.inner.held {
            Held::Shared { path } => &*opened.insert(open_database(path, false)?),
            Held::Exclusive { db } => db,
            Held::ReadOnly { .. } => {
                return Err(io::Error::new(
                    io::ErrorKind::PermissionDenied,
                    "kv database is open read-only",
                ));
            }
        };
        let txn = db.begin_write().map_err(io::Error::other)?;
        match work(&mut WriteTx { txn: &txn, counter: None })? {
            Decision::Commit(value) => txn.commit().map_err(io::Error::other).map(|()| value),
            Decision::Abort(value) => Ok(value),
        }
    }
}