mmdb 4.0.4

The storage engine behind vsdb — a pure-Rust LSM-Tree key-value store
Documentation
//! Table cache: caches SST table readers to avoid repeated file open/parse.

use std::{
    path::{Path, PathBuf},
    sync::Arc,
};

use crate::cache::block_cache::BlockCache;
use crate::error::{Result, ResultExt};
use crate::sst::table_reader::TableReader;
use crate::stats::DbStats;

/// Cache for open TableReader instances.
pub struct TableCache {
    db_path: PathBuf,
    inner: moka::sync::Cache<u64, Arc<TableReader>>,
    block_cache: Option<Arc<BlockCache>>,
    stats: Option<Arc<DbStats>>,
}

impl TableCache {
    /// Create a new table cache.
    /// `max_open_files` is the maximum number of TableReader entries retained by
    /// this cache; live Versions and iterators can pin additional readers.
    #[cfg(test)]
    pub fn new(db_path: &Path, max_open_files: u64, block_cache: Option<Arc<BlockCache>>) -> Self {
        Self::new_with_stats(db_path, max_open_files, block_cache, None)
    }

    /// Create a new table cache with optional stats.
    pub fn new_with_stats(
        db_path: &Path,
        max_open_files: u64,
        block_cache: Option<Arc<BlockCache>>,
        stats: Option<Arc<DbStats>>,
    ) -> Self {
        Self {
            db_path: db_path.to_path_buf(),
            inner: moka::sync::Cache::builder()
                .max_capacity(max_open_files)
                .build(),
            block_cache,
            stats,
        }
    }

    /// Get or open a table reader for the given file number.
    /// Uses moka's `try_get_with` to coalesce concurrent loads for the same file.
    /// `Error: Clone` lets every coalesced waiter receive the full typed error
    /// chain instead of a stringified copy.
    pub fn get_reader(&self, file_number: u64) -> Result<Arc<TableReader>> {
        let db_path = self.db_path.clone();
        let block_cache = self.block_cache.clone();
        let stats = self.stats.clone();
        self.inner
            .try_get_with(file_number, || {
                let path = db_path.join(format!("{:06}.sst", file_number));
                TableReader::open_with_all(&path, file_number, block_cache, stats).map(Arc::new)
            })
            .map_err(|e: Arc<crate::error::Error>| (*e).clone())
            .with_ctx(|| format!("table cache load failed for file {:06}", file_number))
    }

    /// Evict a file from the cache.
    pub fn evict(&self, file_number: u64) {
        self.inner.invalidate(&file_number);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::sst::table_builder::{TableBuildOptions, TableBuilder};

    #[test]
    fn test_table_cache() {
        let dir = tempfile::tempdir().unwrap();

        // Create an SST file
        let path = dir.path().join("000001.sst");
        let mut builder = TableBuilder::new(&path, TableBuildOptions::default()).unwrap();
        builder.add(b"key", b"value").unwrap();
        builder.finish().unwrap();

        let cache = TableCache::new(dir.path(), 100, None);

        // First access opens the file
        let reader = cache.get_reader(1).unwrap();
        let val = reader.get(b"key").unwrap();
        assert_eq!(val, Some(b"value".to_vec()));

        // Second access should hit cache
        let reader2 = cache.get_reader(1).unwrap();
        assert!(Arc::ptr_eq(&reader, &reader2));

        // Evict and re-open
        cache.evict(1);
        cache.inner.run_pending_tasks();
        let reader3 = cache.get_reader(1).unwrap();
        assert!(!Arc::ptr_eq(&reader, &reader3));
    }
}