mentedb-storage 0.45.0

//! Buffer Pool: in-memory page cache with CLOCK eviction.
//!
//! Pages are loaded from disk into fixed-capacity frame slots. Pin counts
//! prevent eviction of pages currently in use. The CLOCK algorithm sweeps
//! frames looking for an unpinned, unreferenced victim when the pool is full.

use std::sync::atomic::{AtomicU64, Ordering};

use ahash::AHashMap;
use parking_lot::Mutex;

use super::page::{Page, PageId, PageManager};
use mentedb_core::error::{MenteError, MenteResult};
use tracing::{debug, trace};

type FrameId = usize;

/// A snapshot of buffer-pool cache activity, for metrics. `hits / (hits + misses)`
/// is the page cache hit ratio, the primary read-performance signal.
#[derive(Debug, Clone, Copy, Default)]
pub struct BufferStats {
    pub hits: u64,
    pub misses: u64,
    pub evictions: u64,
    /// Frames currently holding a page.
    pub resident_pages: u64,
}

/// A single frame in the buffer pool.
struct Frame {
    page: Box<Page>,
    page_id: Option<PageId>,
    pin_count: u32,
    dirty: bool,
    /// CLOCK reference bit: set on access, cleared by the sweep hand.
    reference: bool,
}

impl Frame {
    fn new() -> Self {
        Self {
            page: Box::new(Page::zeroed()),
            page_id: None,
            pin_count: 0,
            dirty: false,
            reference: false,
        }
    }
}

struct BufferPoolInner {
    frames: Vec<Frame>,
    page_table: AHashMap<PageId, FrameId>,
    clock_hand: usize,
    capacity: usize,
}

/// Thread-safe buffer pool with CLOCK eviction.
pub struct BufferPool {
    inner: Mutex<BufferPoolInner>,
    hits: AtomicU64,
    misses: AtomicU64,
    evictions: AtomicU64,
}

impl BufferPool {
    /// Create a buffer pool that may grow to `capacity` frame slots.
    ///
    /// Frames are allocated on demand rather than up front: each frame owns
    /// a full page buffer, so eager allocation costs `capacity * 64KB` per
    /// open database regardless of use. A process holding many open
    /// databases (the platform gateway keeps an LRU of per user databases)
    /// pays only for pages actually touched.
    pub fn new(capacity: usize) -> Self {
        assert!(capacity > 0, "buffer pool capacity must be > 0");
        Self {
            inner: Mutex::new(BufferPoolInner {
                frames: Vec::new(),
                page_table: AHashMap::new(),
                clock_hand: 0,
                capacity,
            }),
            hits: AtomicU64::new(0),
            misses: AtomicU64::new(0),
            evictions: AtomicU64::new(0),
        }
    }

    /// Cache-activity counters for metrics. Lock-free reads.
    pub fn stats(&self) -> BufferStats {
        let resident_pages = self.inner.lock().page_table.len() as u64;
        BufferStats {
            hits: self.hits.load(Ordering::Relaxed),
            misses: self.misses.load(Ordering::Relaxed),
            evictions: self.evictions.load(Ordering::Relaxed),
            resident_pages,
        }
    }

    /// Fetch a page into the pool (loading from disk if necessary).
    ///
    /// The page is automatically pinned (pin_count incremented).
    /// Caller must call `unpin_page` when done.
    pub fn fetch_page(&self, page_id: PageId, pm: &mut PageManager) -> MenteResult<Box<Page>> {
        let mut inner = self.inner.lock();

        // Cache hit
        if let Some(&frame_id) = inner.page_table.get(&page_id) {
            let frame = &mut inner.frames[frame_id];
            frame.pin_count += 1;
            frame.reference = true;
            self.hits.fetch_add(1, Ordering::Relaxed);
            trace!(page_id = page_id.0, frame_id, "buffer pool hit");
            return Ok(frame.page.clone());
        }

        // Cache miss — find a victim frame.
        self.misses.fetch_add(1, Ordering::Relaxed);
        let frame_id = Self::find_victim(&mut inner)?;

        // Flush dirty victim if needed.
        if inner.frames[frame_id].dirty
            && let Some(old_pid) = inner.frames[frame_id].page_id
        {
            pm.write_page(old_pid, &inner.frames[frame_id].page)?;
            debug!(page_id = old_pid.0, frame_id, "flushed dirty victim");
        }

        // Remove old mapping. A displaced page is a real eviction (the reuse and
        // grow paths in find_victim leave page_id None, so they are not counted).
        if let Some(old_pid) = inner.frames[frame_id].page_id {
            inner.page_table.remove(&old_pid);
            self.evictions.fetch_add(1, Ordering::Relaxed);
        }

        // Load the requested page from disk.
        let page = pm.read_page(page_id)?;
        {
            let frame = &mut inner.frames[frame_id];
            *frame.page = *page;
            frame.page_id = Some(page_id);
            frame.pin_count = 1;
            frame.dirty = false;
            frame.reference = true;
        }

        inner.page_table.insert(page_id, frame_id);
        trace!(
            page_id = page_id.0,
            frame_id, "loaded page into buffer pool"
        );

        Ok(inner.frames[frame_id].page.clone())
    }

    /// Increment the pin count of a page already in the pool.
    pub fn pin_page(&self, page_id: PageId) -> MenteResult<()> {
        let mut inner = self.inner.lock();
        match inner.page_table.get(&page_id) {
            Some(&fid) => {
                inner.frames[fid].pin_count += 1;
                Ok(())
            }
            None => Err(MenteError::Storage(format!(
                "page {} not in buffer pool",
                page_id.0
            ))),
        }
    }

    /// Decrement pin count and optionally mark the page dirty.
    pub fn unpin_page(&self, page_id: PageId, dirty: bool) -> MenteResult<()> {
        let mut inner = self.inner.lock();
        match inner.page_table.get(&page_id) {
            Some(&fid) => {
                let frame = &mut inner.frames[fid];
                if frame.pin_count > 0 {
                    frame.pin_count -= 1;
                }
                if dirty {
                    frame.dirty = true;
                }
                Ok(())
            }
            None => Err(MenteError::Storage(format!(
                "page {} not in buffer pool",
                page_id.0
            ))),
        }
    }

    /// Drop a page from the pool without flushing, discarding any dirty state.
    ///
    /// Used when a page is freed: the on-disk state is already authoritative,
    /// so a stale cached copy must never be served or flushed back. A no-op if
    /// the page is not cached.
    pub fn invalidate(&self, page_id: PageId) {
        let mut inner = self.inner.lock();
        if let Some(fid) = inner.page_table.remove(&page_id) {
            let frame = &mut inner.frames[fid];
            frame.page_id = None;
            frame.pin_count = 0;
            frame.dirty = false;
            frame.reference = false;
            debug!(page_id = page_id.0, "invalidated cached page");
        }
    }

    /// Replace the cached copy of a page and mark it dirty.
    pub fn update_page(&self, page_id: PageId, page: &Page) -> MenteResult<()> {
        let mut inner = self.inner.lock();
        match inner.page_table.get(&page_id) {
            Some(&fid) => {
                let frame = &mut inner.frames[fid];
                *frame.page = page.clone();
                frame.dirty = true;
                Ok(())
            }
            None => Err(MenteError::Storage(format!(
                "page {} not in buffer pool",
                page_id.0
            ))),
        }
    }

    /// Flush a single dirty page to disk.
    pub fn flush_page(&self, page_id: PageId, pm: &mut PageManager) -> MenteResult<()> {
        let mut inner = self.inner.lock();
        match inner.page_table.get(&page_id) {
            Some(&fid) => {
                let frame = &mut inner.frames[fid];
                if frame.dirty {
                    pm.write_page(page_id, &frame.page)?;
                    frame.dirty = false;
                    debug!(page_id = page_id.0, "flushed page");
                }
                Ok(())
            }
            None => Err(MenteError::Storage(format!(
                "page {} not in buffer pool",
                page_id.0
            ))),
        }
    }

    /// Flush all dirty pages to disk.
    pub fn flush_all(&self, pm: &mut PageManager) -> MenteResult<()> {
        let mut inner = self.inner.lock();
        for frame in &mut inner.frames {
            if frame.dirty
                && let Some(pid) = frame.page_id
            {
                pm.write_page(pid, &frame.page)?;
                frame.dirty = false;
            }
        }
        debug!("flushed all dirty pages");
        Ok(())
    }

    /// CLOCK eviction: find an unpinned, unreferenced frame.
    fn find_victim(inner: &mut BufferPoolInner) -> MenteResult<FrameId> {
        // Reuse an invalidated frame first.
        for i in 0..inner.frames.len() {
            if inner.frames[i].page_id.is_none() {
                return Ok(i);
            }
        }

        // Grow on demand up to capacity.
        if inner.frames.len() < inner.capacity {
            inner.frames.push(Frame::new());
            return Ok(inner.frames.len() - 1);
        }

        // Pool is at capacity: CLOCK sweep, at most 2 full rotations.
        let cap = inner.capacity;
        let max_sweeps = cap * 2;
        for _ in 0..max_sweeps {
            let idx = inner.clock_hand;
            inner.clock_hand = (inner.clock_hand + 1) % cap;

            let frame = &mut inner.frames[idx];
            if frame.pin_count == 0 {
                if !frame.reference {
                    return Ok(idx);
                }
                frame.reference = false;
            }
        }

        Err(MenteError::Storage(
            "buffer pool full: all pages are pinned".into(),
        ))
    }
}