armdb 0.7.0

sharded bitcask key-value storage optimized for NVMe
Documentation
use std::fs::File;
use std::sync::Arc;

use crate::error::{DbError, DbResult};
use crate::io::direct;
use crate::shard::ImmutableFile;

#[cfg(feature = "encryption")]
use crate::crypto::PageCipher;
#[cfg(feature = "encryption")]
use crate::io::tags::TagFile;

#[derive(Clone)]
enum CatchupDataFile {
    Immutable(Arc<ImmutableFile>),
    Active(Arc<File>),
}

#[derive(Clone)]
pub(crate) struct CatchupFile {
    pub(crate) file_id: u32,
    pub(crate) readable_end: u64,
    data: CatchupDataFile,
    #[cfg(feature = "encryption")]
    tag_file: Option<Arc<TagFile>>,
    #[cfg(feature = "encryption")]
    cipher: Option<Arc<PageCipher>>,
}

impl CatchupFile {
    pub(crate) fn immutable(
        file: Arc<ImmutableFile>,
        #[cfg(feature = "encryption")] cipher: Option<Arc<PageCipher>>,
    ) -> Self {
        Self {
            file_id: file.file_id,
            readable_end: file.total_bytes,
            #[cfg(feature = "encryption")]
            tag_file: file.tag_file.clone(),
            #[cfg(feature = "encryption")]
            cipher,
            data: CatchupDataFile::Immutable(file),
        }
    }

    pub(crate) fn active(
        file_id: u32,
        readable_end: u64,
        file: Arc<File>,
        #[cfg(feature = "encryption")] tag_file: Option<Arc<TagFile>>,
        #[cfg(feature = "encryption")] cipher: Option<Arc<PageCipher>>,
    ) -> Self {
        Self {
            file_id,
            readable_end,
            data: CatchupDataFile::Active(file),
            #[cfg(feature = "encryption")]
            tag_file,
            #[cfg(feature = "encryption")]
            cipher,
        }
    }

    fn file(&self) -> &File {
        match &self.data {
            CatchupDataFile::Immutable(file) => &file.file,
            CatchupDataFile::Active(file) => file,
        }
    }

    pub(crate) fn read_exact_at(&self, offset: u64, len: usize) -> DbResult<Vec<u8>> {
        let end = offset
            .checked_add(len as u64)
            .ok_or(DbError::CorruptedEntry { offset })?;
        if end > self.readable_end {
            return Err(DbError::CorruptedEntry { offset });
        }
        #[cfg(feature = "encryption")]
        if let Some(cipher) = &self.cipher {
            let tags = self
                .tag_file
                .as_ref()
                .ok_or(DbError::CorruptedEntry { offset })?;
            return direct::pread_value_encrypted(
                self.file(),
                tags,
                cipher,
                self.file_id,
                offset,
                len,
            );
        }
        direct::pread_value(self.file(), offset, len)
    }
}

pub(crate) struct ShardCatchupSnapshot {
    pub(crate) shard_id: u8,
    pub(crate) max_gsn: u64,
    pub(crate) overflow_baseline: u64,
    pub(crate) files: Vec<CatchupFile>,
    pinned_bytes: u64,
}

impl ShardCatchupSnapshot {
    pub(crate) fn new(
        shard_id: u8,
        max_gsn: u64,
        overflow_baseline: u64,
        files: Vec<CatchupFile>,
    ) -> Self {
        let pinned_bytes = files.iter().map(|file| file.readable_end).sum();
        metrics::gauge!(
            "armdb.replication.snapshot_pinned_bytes",
            "shard" => shard_id.to_string()
        )
        .set(pinned_bytes as f64);
        Self {
            shard_id,
            max_gsn,
            overflow_baseline,
            files,
            pinned_bytes,
        }
    }
}

impl Drop for ShardCatchupSnapshot {
    fn drop(&mut self) {
        let _ = self.pinned_bytes;
        metrics::gauge!(
            "armdb.replication.snapshot_pinned_bytes",
            "shard" => self.shard_id.to_string()
        )
        .set(0.0);
    }
}