armdb 0.12.0

sharded bitcask key-value storage optimized for NVMe
Documentation
use std::path::Path;
use std::time::Duration;

use armour_core::Typ;

use crate::{Config, DbResult, FixedConfig};

/// Child-process crash injection for the format-migration durable state machine.
///
/// This is deliberately exposed only by `test-faults`: production builds have
/// neither the atomics nor the exit seam.
#[cfg(any(test, feature = "test-faults"))]
pub mod failpoints {
    use std::path::{Path, PathBuf};
    use std::sync::Mutex;
    use std::sync::atomic::{AtomicU8, Ordering};

    pub const NONE: u8 = 0;
    pub const AFTER_FORMAT_VERSION: u8 = 1;
    pub const BEFORE_SOURCE_RENAME: u8 = 2;
    pub const AFTER_SOURCE_RENAME: u8 = 3;
    pub const BEFORE_TARGET_RENAME: u8 = 4;
    pub const AFTER_TARGET_RENAME: u8 = 5;
    pub const BEFORE_OLD_REMOVE: u8 = 6;
    pub const AFTER_OLD_REMOVE: u8 = 7;
    pub const BEFORE_INFO_COMMIT: u8 = 8;
    pub const AFTER_INFO_COMMIT: u8 = 9;
    pub const BEFORE_MARKER_REMOVE: u8 = 10;
    pub const DURING_STAGE_FINALIZE: u8 = 11;

    /// Ordinary-error cleanup injection modes. Unlike the numbered durable
    /// boundary stages these return an I/O error instead of terminating the
    /// process, so tests can verify that cleanup failures reach the caller.
    pub const CLEANUP_NONE: u8 = 0;
    pub const CLEANUP_REMOVE_ERROR: u8 = 1;
    pub const CLEANUP_SYNC_ERROR: u8 = 2;

    pub static STAGE: AtomicU8 = AtomicU8::new(NONE);
    pub static CLEANUP_ERROR: AtomicU8 = AtomicU8::new(CLEANUP_NONE);
    static READ_DIR_ENTRY_ERROR: Mutex<Option<PathBuf>> = Mutex::new(None);

    /// Inject one individual `read_dir` entry error for the named directory.
    pub fn fail_next_read_dir_entry(path: impl Into<PathBuf>) {
        *READ_DIR_ENTRY_ERROR
            .lock()
            .expect("read-dir failpoint poisoned") = Some(path.into());
    }

    pub(crate) fn take_read_dir_entry_error(path: &Path) -> bool {
        let mut configured = READ_DIR_ENTRY_ERROR
            .lock()
            .expect("read-dir failpoint poisoned");
        if configured.as_deref() == Some(path) {
            configured.take();
            true
        } else {
            false
        }
    }

    pub fn hit(stage: u8) {
        if STAGE.load(Ordering::SeqCst) == stage {
            // A hard exit models process loss after the preceding durable
            // action without exercising destructors that could clean it up.
            std::process::exit(9);
        }
    }
}

/// Physical storage parameters required to rewrite a collection before it is
/// opened and its values are decoded.
#[derive(Clone, Copy)]
pub(crate) enum StorageDescriptor<'a> {
    Bitcask {
        config: &'a Config,
        key_len: usize,
        claims: bool,
    },
    Fixed {
        config: &'a FixedConfig,
        key_len: usize,
        value_len: u16,
        claims: bool,
    },
    #[allow(dead_code)]
    Lsm,
}

impl StorageDescriptor<'_> {
    pub(crate) fn validate(&self) -> DbResult<()> {
        match self {
            Self::Bitcask { config, .. } => config.validate(),
            Self::Fixed { config, .. } => config.validate(),
            Self::Lsm => Ok(()),
        }
    }

    pub(crate) fn policy(&self) -> FormatMigrationPolicy {
        match self {
            Self::Bitcask { config, .. } => config.format_migration,
            Self::Fixed { config, .. } => config.format_migration,
            Self::Lsm => FormatMigrationPolicy::Auto,
        }
    }

    /// Validate immutable on-disk backend parameters before a format pass can
    /// create markers, dirty metadata, or rewrite any shard.
    pub(crate) fn validate_persisted(&self, collection_dir: &Path) -> DbResult<()> {
        let meta_path = collection_dir.join("db.meta");
        if !meta_path.exists() || matches!(self, Self::Lsm) {
            return Ok(());
        }
        let meta = crate::meta::DbMeta::decode(&std::fs::read(meta_path)?)?;
        match self {
            Self::Bitcask { config, .. } => {
                if meta.backend != crate::meta::Backend::Bitcask {
                    return Err(crate::DbError::FormatMismatch(format!(
                        "db.meta: expected Bitcask backend, found {:?}",
                        meta.backend
                    )));
                }
                if meta.shard_count as usize != config.shard_count {
                    return Err(crate::DbError::FormatMismatch(format!(
                        "shard_count mismatch: db has {}, config has {}",
                        meta.shard_count, config.shard_count
                    )));
                }
                if meta.shard_prefix_bits as usize != config.shard_prefix_bits {
                    return Err(crate::DbError::FormatMismatch(format!(
                        "shard_prefix_bits mismatch: db has {}, config has {}",
                        meta.shard_prefix_bits, config.shard_prefix_bits
                    )));
                }
                let was_encrypted = meta.flags & 1 != 0;
                #[cfg(feature = "encryption")]
                if was_encrypted != config.encryption_key.is_some() {
                    return Err(crate::DbError::Config(if was_encrypted {
                        "database is encrypted but no encryption_key provided"
                    } else {
                        "database is not encrypted but encryption_key was provided"
                    }));
                }
                #[cfg(not(feature = "encryption"))]
                if was_encrypted {
                    return Err(crate::DbError::FormatMismatch(
                        "database is encrypted but binary lacks the encryption feature".into(),
                    ));
                }
            }
            Self::Fixed { config, .. } => {
                if meta.backend != crate::meta::Backend::FixedStore {
                    return Err(crate::DbError::FormatMismatch(format!(
                        "db.meta: expected FixedStore backend, found {:?}",
                        meta.backend
                    )));
                }
                if meta.shard_count as usize != config.shard_count {
                    return Err(crate::DbError::FormatMismatch(format!(
                        "db.meta: shard_count mismatch: stored {}, config {}",
                        meta.shard_count, config.shard_count
                    )));
                }
                if meta.shard_prefix_bits as usize != config.shard_prefix_bits {
                    return Err(crate::DbError::FormatMismatch(format!(
                        "db.meta: shard_prefix_bits mismatch: stored {}, config {}",
                        meta.shard_prefix_bits, config.shard_prefix_bits
                    )));
                }
            }
            Self::Lsm => {}
        }
        Ok(())
    }

    pub(crate) fn claims(&self) -> bool {
        match self {
            Self::Bitcask { claims, .. } | Self::Fixed { claims, .. } => *claims,
            Self::Lsm => false,
        }
    }

    pub(crate) fn shard_count(&self) -> Option<usize> {
        match self {
            Self::Bitcask { config, .. } => Some(config.shard_count),
            Self::Fixed { config, .. } => Some(config.shard_count),
            Self::Lsm => None,
        }
    }
}

#[cfg(feature = "armour")]
#[allow(dead_code)] // Consumed by the subsequent opener/rewrite tasks.
#[doc(hidden)]
pub mod plan;

#[cfg(feature = "armour")]
#[allow(dead_code)] // Consumed by the subsequent opener/rewrite tasks.
pub(crate) mod marker;

#[cfg(feature = "armour")]
pub(crate) mod bitcask;

#[cfg(feature = "armour")]
pub(crate) mod fixed;

/// Policy controlling whether an opener runs a pending value-format migration.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[cfg_attr(
    any(feature = "armour", feature = "postcard-codec"),
    derive(serde::Deserialize)
)]
pub enum FormatMigrationPolicy {
    /// Run declared format migrations automatically while opening the collection.
    #[default]
    Auto,
    /// Refuse an open that requires a format rewrite.
    Refuse,
}

/// Decision returned by one value-format migration step.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FormatAction {
    Write,
    Keep,
    Drop,
}

/// Byte-to-byte callback used by a format migration step.
pub type FormatStepFn = fn(key: &[u8], old: &[u8], out: &mut Vec<u8>) -> DbResult<FormatAction>;

/// One declared value-format transition.
#[derive(Clone, Copy)]
pub struct FormatMigration {
    pub from_version: u16,
    pub from_typ: Option<&'static Typ>,
    pub to_typ: Option<&'static Typ>,
    pub step: FormatStepFn,
}

impl FormatMigration {
    pub const fn raw(from_version: u16, step: FormatStepFn) -> Self {
        Self {
            from_version,
            from_typ: None,
            to_typ: None,
            step,
        }
    }

    pub const fn raw_checked(
        from_version: u16,
        from_typ: &'static Typ,
        to_typ: &'static Typ,
        step: FormatStepFn,
    ) -> Self {
        Self {
            from_version,
            from_typ: Some(from_typ),
            to_typ: Some(to_typ),
            step,
        }
    }
}

/// Summary returned by an explicit format migration run.
#[derive(Debug, Clone)]
pub struct FormatMigrationReport {
    pub collection: &'static str,
    pub from: u16,
    pub to: u16,
    pub shards: usize,
    pub rows_migrated: u64,
    pub rows_dropped: u64,
    pub bytes_before: u64,
    pub bytes_after: u64,
    pub elapsed: Duration,
    pub warnings: Vec<String>,
}

impl FormatMigrationReport {
    pub(crate) fn empty(collection: &'static str, version: u16) -> Self {
        Self {
            collection,
            from: version,
            to: version,
            shards: 0,
            rows_migrated: 0,
            rows_dropped: 0,
            bytes_before: 0,
            bytes_after: 0,
            elapsed: Duration::ZERO,
            warnings: Vec::new(),
        }
    }
}

#[cfg(feature = "armour")]
pub(crate) fn refuse_if_required(
    policy: FormatMigrationPolicy,
    collection: &'static str,
    plan: &plan::MigrationPlan,
) -> DbResult<()> {
    if policy == FormatMigrationPolicy::Refuse {
        return Err(crate::DbError::FormatMigrationRequired {
            collection,
            from: plan.stored_version,
            to: plan.format_to_version,
        });
    }
    Ok(())
}