nodedb-vector 0.0.6

//! Memory-mapped vector segment for L1 NVMe tiering.
//!
//! Stores FP32 vectors contiguously in a file, memory-mapped for read access.
//! Layout: `[dim:u32][count:u32][v0_f0..v0_fD][v1_f0..v1_fD]...]`

use std::os::fd::AsRawFd;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};

/// Drop-time page-cache policy for a vector segment.
///
/// HNSW traversal touches a small fraction of a segment's pages. When the
/// segment is dropped we hint the kernel that the residual pages can be
/// evicted so they don't crowd hotter engines' working sets.
#[derive(Debug, Clone, Copy)]
pub struct VectorSegmentDropPolicy {
    dontneed_on_drop: bool,
}

impl VectorSegmentDropPolicy {
    pub const fn new(dontneed_on_drop: bool) -> Self {
        Self { dontneed_on_drop }
    }
    pub const fn keep_resident() -> Self {
        Self {
            dontneed_on_drop: false,
        }
    }
    pub const fn dontneed_on_drop(self) -> bool {
        self.dontneed_on_drop
    }
}

impl Default for VectorSegmentDropPolicy {
    fn default() -> Self {
        Self {
            dontneed_on_drop: true,
        }
    }
}

/// Module-scoped counters for observing madvise behaviour in tests.
///
/// These are lightweight atomics, always compiled — the same counters are
/// useful to Event-Plane metrics, not just tests.
pub mod test_hooks {
    use super::{AtomicU64, Ordering};
    pub(super) static DONTNEED_COUNT: AtomicU64 = AtomicU64::new(0);
    pub(super) static RANDOM_COUNT: AtomicU64 = AtomicU64::new(0);

    pub fn dontneed_count() -> u64 {
        DONTNEED_COUNT.load(Ordering::Relaxed)
    }
    pub fn random_count() -> u64 {
        RANDOM_COUNT.load(Ordering::Relaxed)
    }
}

/// Memory-mapped vector segment file.
///
/// Not `Send` or `Sync` — owned by a single Data Plane core.
pub struct MmapVectorSegment {
    path: PathBuf,
    _fd: std::fs::File,
    base: *const u8,
    mmap_size: usize,
    dim: usize,
    count: usize,
    data_offset: usize,
    drop_policy: VectorSegmentDropPolicy,
    madvise_state: Option<libc::c_int>,
}

const HEADER_SIZE: usize = 8;

impl MmapVectorSegment {
    /// Create a new segment file and write vectors to it.
    pub fn create(path: &Path, dim: usize, vectors: &[&[f32]]) -> std::io::Result<Self> {
        Self::create_with_policy(path, dim, vectors, VectorSegmentDropPolicy::default())
    }

    /// Create a new segment file with an explicit drop policy.
    pub fn create_with_policy(
        path: &Path,
        dim: usize,
        vectors: &[&[f32]],
        policy: VectorSegmentDropPolicy,
    ) -> std::io::Result<Self> {
        use std::io::Write;

        if let Some(parent) = path.parent() {
            std::fs::create_dir_all(parent)?;
        }

        let count = vectors.len();

        let mut fd = std::fs::OpenOptions::new()
            .read(true)
            .write(true)
            .create(true)
            .truncate(true)
            .open(path)?;

        fd.write_all(&(dim as u32).to_le_bytes())?;
        fd.write_all(&(count as u32).to_le_bytes())?;

        for v in vectors {
            debug_assert_eq!(v.len(), dim);
            let bytes: &[u8] =
                unsafe { std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 4) };
            fd.write_all(bytes)?;
        }
        fd.sync_all()?;

        drop(fd);
        Self::open_with_policy(path, policy)
    }

    /// Open an existing segment file and memory-map it.
    pub fn open(path: &Path) -> std::io::Result<Self> {
        Self::open_with_policy(path, VectorSegmentDropPolicy::default())
    }

    /// Open an existing segment with an explicit drop policy.
    pub fn open_with_policy(path: &Path, policy: VectorSegmentDropPolicy) -> std::io::Result<Self> {
        let fd = std::fs::OpenOptions::new().read(true).open(path)?;

        let file_size = fd.metadata()?.len() as usize;
        if file_size < HEADER_SIZE {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                "mmap vector segment too small for header",
            ));
        }

        let base = unsafe {
            libc::mmap(
                std::ptr::null_mut(),
                file_size,
                libc::PROT_READ,
                libc::MAP_PRIVATE,
                fd.as_raw_fd(),
                0,
            )
        };

        if base == libc::MAP_FAILED {
            return Err(std::io::Error::last_os_error());
        }

        let base = base as *const u8;

        let dim = unsafe {
            let ptr = base as *const u32;
            u32::from_le(*ptr) as usize
        };
        let count = unsafe {
            let ptr = base.add(4) as *const u32;
            u32::from_le(*ptr) as usize
        };

        // Reject dim=0 with nonzero count: get_vector would compute offset=HEADER_SIZE
        // for every ID, aliasing header bytes as vector data.
        if dim == 0 && count > 0 {
            unsafe {
                libc::munmap(base as *mut libc::c_void, file_size);
            }
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                "mmap segment has dim=0 with nonzero count",
            ));
        }

        // Use checked arithmetic to prevent usize overflow on crafted headers.
        let data_bytes = dim
            .checked_mul(count)
            .and_then(|dc| dc.checked_mul(4))
            .and_then(|bytes| bytes.checked_add(HEADER_SIZE));
        let expected = match data_bytes {
            Some(v) => v,
            None => {
                unsafe {
                    libc::munmap(base as *mut libc::c_void, file_size);
                }
                return Err(std::io::Error::new(
                    std::io::ErrorKind::InvalidData,
                    format!("mmap segment header overflow: dim={dim}, count={count}"),
                ));
            }
        };
        if file_size < expected {
            unsafe {
                libc::munmap(base as *mut libc::c_void, file_size);
            }
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                format!("mmap segment truncated: expected {expected} bytes, got {file_size}"),
            ));
        }

        // Advise MADV_RANDOM: HNSW graph traversal touches non-adjacent
        // vector IDs. Default MADV_NORMAL readahead wastes NVMe bandwidth
        // on neighbouring pages evicted before the walk reaches them.
        //
        // Skip advising on header-only files (dim=0 or count=0): madvise
        // on a zero-data-range region is allowed but meaningless.
        let mut madvise_state = None;
        let data_bytes = file_size.saturating_sub(HEADER_SIZE);
        if data_bytes > 0 {
            let rc =
                unsafe { libc::madvise(base as *mut libc::c_void, file_size, libc::MADV_RANDOM) };
            if rc == 0 {
                madvise_state = Some(libc::MADV_RANDOM);
                test_hooks::RANDOM_COUNT.fetch_add(1, Ordering::Relaxed);
            } else {
                tracing::warn!(
                    path = %path.display(),
                    errno = std::io::Error::last_os_error().raw_os_error().unwrap_or(0),
                    "madvise(MADV_RANDOM) failed on vector segment; continuing with kernel default",
                );
            }
        }

        Ok(Self {
            path: path.to_path_buf(),
            _fd: fd,
            base,
            mmap_size: file_size,
            dim,
            count,
            data_offset: HEADER_SIZE,
            drop_policy: policy,
            madvise_state,
        })
    }

    /// The madvise hint set on this segment (if any).
    pub fn madvise_state(&self) -> Option<libc::c_int> {
        self.madvise_state
    }

    /// Get a vector by ID. Returns a slice into the mmap'd region.
    #[inline]
    pub fn get_vector(&self, id: u32) -> Option<&[f32]> {
        let idx = id as usize;
        if idx >= self.count {
            return None;
        }
        let byte_len = self.dim.checked_mul(4)?;
        let offset = self.data_offset.checked_add(idx.checked_mul(byte_len)?)?;
        let end = offset.checked_add(byte_len)?;
        if end > self.mmap_size {
            return None;
        }
        unsafe {
            let ptr = self.base.add(offset) as *const f32;
            Some(std::slice::from_raw_parts(ptr, self.dim))
        }
    }

    /// Prefetch a vector's page into memory via `madvise(MADV_WILLNEED)`.
    pub fn prefetch(&self, id: u32) {
        let idx = id as usize;
        if idx >= self.count {
            return;
        }
        let byte_len = match self.dim.checked_mul(4) {
            Some(v) => v,
            None => return,
        };
        let Some(idx_bytes) = idx.checked_mul(byte_len) else {
            return;
        };
        let Some(offset) = self.data_offset.checked_add(idx_bytes) else {
            return;
        };
        if offset
            .checked_add(byte_len)
            .is_none_or(|e| e > self.mmap_size)
        {
            return;
        }
        let page_start = offset & !(4095);
        let len = (byte_len + 4095) & !(4095);
        unsafe {
            libc::madvise(
                self.base.add(page_start) as *mut libc::c_void,
                len,
                libc::MADV_WILLNEED,
            );
        }
    }

    /// Prefetch a batch of vector IDs.
    pub fn prefetch_batch(&self, ids: &[u32]) {
        for &id in ids {
            self.prefetch(id);
        }
    }

    pub fn dim(&self) -> usize {
        self.dim
    }

    pub fn count(&self) -> usize {
        self.count
    }

    pub fn path(&self) -> &Path {
        &self.path
    }

    pub fn mmap_bytes(&self) -> usize {
        self.mmap_size
    }

    pub fn file_size(&self) -> usize {
        self.mmap_size
    }
}

impl Drop for MmapVectorSegment {
    fn drop(&mut self) {
        if !self.base.is_null() && self.mmap_size > 0 {
            if self.drop_policy.dontneed_on_drop() {
                let data_bytes = self.mmap_size.saturating_sub(HEADER_SIZE);
                if data_bytes > 0 {
                    unsafe {
                        libc::madvise(
                            self.base as *mut libc::c_void,
                            self.mmap_size,
                            libc::MADV_DONTNEED,
                        );
                    }
                    test_hooks::DONTNEED_COUNT.fetch_add(1, Ordering::Relaxed);
                }
            }
            unsafe {
                libc::munmap(self.base as *mut libc::c_void, self.mmap_size);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn create_and_read() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("test.vseg");

        let v0: Vec<f32> = vec![1.0, 2.0, 3.0];
        let v1: Vec<f32> = vec![4.0, 5.0, 6.0];
        let v2: Vec<f32> = vec![7.0, 8.0, 9.0];

        let seg = MmapVectorSegment::create(&path, 3, &[&v0, &v1, &v2]).unwrap();

        assert_eq!(seg.dim(), 3);
        assert_eq!(seg.count(), 3);

        assert_eq!(seg.get_vector(0).unwrap(), &[1.0, 2.0, 3.0]);
        assert_eq!(seg.get_vector(1).unwrap(), &[4.0, 5.0, 6.0]);
        assert_eq!(seg.get_vector(2).unwrap(), &[7.0, 8.0, 9.0]);
        assert!(seg.get_vector(3).is_none());
    }

    #[test]
    fn reopen_roundtrip() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("reopen.vseg");

        let vectors: Vec<Vec<f32>> = (0..100)
            .map(|i| vec![i as f32, (i as f32).sin(), (i as f32).cos()])
            .collect();
        let refs: Vec<&[f32]> = vectors.iter().map(|v| v.as_slice()).collect();

        MmapVectorSegment::create(&path, 3, &refs).unwrap();

        let seg = MmapVectorSegment::open(&path).unwrap();
        assert_eq!(seg.count(), 100);
        for (i, v) in vectors.iter().enumerate() {
            let loaded = seg.get_vector(i as u32).unwrap();
            assert_eq!(loaded, v.as_slice());
        }
    }

    #[test]
    fn prefetch_does_not_crash() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("prefetch.vseg");

        let v: Vec<f32> = vec![1.0; 768];
        let seg = MmapVectorSegment::create(&path, 768, &[&v]).unwrap();

        seg.prefetch(0);
        seg.prefetch(999);
    }

    #[test]
    fn empty_segment() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("empty.vseg");

        let seg = MmapVectorSegment::create(&path, 3, &[]).unwrap();
        assert_eq!(seg.count(), 0);
        assert!(seg.get_vector(0).is_none());
    }

    #[test]
    fn overflow_dim_count_rejected() {
        use std::io::Write;

        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("overflow.vseg");

        // dim=0x40000001, count=0x40000001: count * dim * 4 overflows usize on 64-bit
        // (0x40000001 * 0x40000001 * 4 = 0x4000000280000004, which wraps to a small value).
        let dim: u32 = 0x40000001;
        let count: u32 = 0x40000001;

        let mut f = std::fs::OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(&path)
            .unwrap();
        f.write_all(&dim.to_le_bytes()).unwrap();
        f.write_all(&count.to_le_bytes()).unwrap();
        // No actual vector data — just a 8-byte header.
        drop(f);

        let result = MmapVectorSegment::open(&path);
        assert!(
            result.is_err(),
            "expected Err for overflow-inducing dim/count, got Ok"
        );
    }

    #[test]
    fn truncated_file_rejected() {
        use std::io::Write;

        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("truncated.vseg");

        // Header claims dim=3, count=100 but only 8 bytes of actual data.
        let dim: u32 = 3;
        let count: u32 = 100;

        let mut f = std::fs::OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(&path)
            .unwrap();
        f.write_all(&dim.to_le_bytes()).unwrap();
        f.write_all(&count.to_le_bytes()).unwrap();
        drop(f);

        let result = MmapVectorSegment::open(&path);
        match result {
            Err(e) => assert_eq!(
                e.kind(),
                std::io::ErrorKind::InvalidData,
                "expected InvalidData, got {:?}",
                e.kind()
            ),
            Ok(_) => panic!("expected Err for truncated file, got Ok"),
        }
    }

    #[test]
    fn zero_dim_with_nonzero_count_rejected() {
        use std::io::Write;

        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("zerodim.vseg");

        // dim=0, count=1000: expected size = HEADER_SIZE + 0 = 8, so the size
        // check passes, but get_vector would read header bytes as vector data.
        // dim=0 must be rejected outright.
        let dim: u32 = 0;
        let count: u32 = 1000;

        let mut f = std::fs::OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(&path)
            .unwrap();
        f.write_all(&dim.to_le_bytes()).unwrap();
        f.write_all(&count.to_le_bytes()).unwrap();
        // Write enough padding so the file passes a naive size check.
        f.write_all(&[0u8; 64]).unwrap();
        drop(f);

        let result = MmapVectorSegment::open(&path);
        assert!(
            result.is_err(),
            "expected Err for dim=0 with nonzero count, got Ok"
        );
    }
}