nodedb_wal/

mmap_reader.rs

// SPDX-License-Identifier: Apache-2.0

//! Memory-mapped WAL segment reader for Event Plane catchup.
//!
//! Unlike the standard `WalReader` (which uses sequential `read_exact`),
//! this reader maps sealed WAL segments into the process address space via
//! `mmap`. The kernel manages the page cache — no slab allocator memory is
//! pinned, and mmap reads from page cache don't contend with the Data Plane's
//! O_DIRECT WAL append path (O_DIRECT bypasses page cache entirely).
//!
//! **Tier progression:**
//! 1. In-memory Arc slabs (hot, zero-copy from ring buffer)
//! 2. Mmap WAL segment reads (warm, kernel-managed pages)
//! 3. Shed consumer + cold WAL replay (last resort)
//!
//! This reader is used in tier 2: when the Event Plane enters WAL Catchup
//! Mode, it mmap's the relevant sealed segments and iterates records.

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

#[cfg(target_os = "linux")]
use std::os::fd::AsRawFd;

use memmap2::Mmap;

use crate::error::{Result, WalError};
use crate::record::{HEADER_SIZE, RecordHeader, RecordType, WAL_MAGIC, WalRecord};

/// Module-scoped atomic counters for observing mmap and fadvise behaviour in
/// production. These counters are incremented by the live code paths (open,
/// madvise, fadvise) and may be read from tests or from a metrics scrape.
pub mod observability {
    use super::{AtomicU64, Ordering};
    pub(super) static SEGMENTS_OPENED: AtomicU64 = AtomicU64::new(0);
    pub(super) static FADV_DONTNEED_COUNT: AtomicU64 = AtomicU64::new(0);
    pub(super) static MADV_SEQUENTIAL_COUNT: AtomicU64 = AtomicU64::new(0);

    pub fn segments_opened() -> u64 {
        SEGMENTS_OPENED.load(Ordering::Relaxed)
    }
    pub fn fadv_dontneed_count() -> u64 {
        FADV_DONTNEED_COUNT.load(Ordering::Relaxed)
    }
    pub fn madv_sequential_count() -> u64 {
        MADV_SEQUENTIAL_COUNT.load(Ordering::Relaxed)
    }
}

/// Call `posix_fadvise(POSIX_FADV_DONTNEED)` on an open WAL segment fd.
///
/// Once a segment has been iterated end-to-end during catchup, we don't
/// need its pages in cache any longer. Release them back to the kernel so
/// replay doesn't pin GiBs of page cache.
fn fadv_dontneed(fd: &std::fs::File, len: usize, path: &Path) {
    if len == 0 {
        return;
    }
    #[cfg(target_os = "linux")]
    {
        let rc = unsafe {
            libc::posix_fadvise(
                fd.as_raw_fd(),
                0,
                len as libc::off_t,
                libc::POSIX_FADV_DONTNEED,
            )
        };
        if rc == 0 {
            observability::FADV_DONTNEED_COUNT.fetch_add(1, Ordering::Relaxed);
        } else {
            tracing::warn!(
                path = %path.display(),
                errno = rc,
                "posix_fadvise(DONTNEED) failed on exhausted WAL segment",
            );
        }
    }
    #[cfg(not(target_os = "linux"))]
    {
        let _ = (fd, path);
    }
}

/// Memory-mapped WAL segment reader.
///
/// Opens a sealed WAL segment file via mmap and provides zero-copy
/// iteration over records. The mmap'd region is read-only and the
/// kernel manages page residency — no application-level memory pinning.
pub struct MmapWalReader {
    mmap: Mmap,
    offset: usize,
    file: std::fs::File,
    path: std::path::PathBuf,
    madvise_state: Option<libc::c_int>,
}

impl MmapWalReader {
    /// Open a WAL segment file for mmap'd reading.
    pub fn open(path: &Path) -> Result<Self> {
        observability::SEGMENTS_OPENED.fetch_add(1, Ordering::Relaxed);
        let file = std::fs::File::open(path)?;
        // SAFETY: The file is a sealed WAL segment (not being written to).
        // The Data Plane writes to the ACTIVE segment via O_DIRECT; sealed
        // segments are immutable after rollover.
        let mmap = unsafe { Mmap::map(&file)? };

        // Catchup iterates forward through a segment. MADV_SEQUENTIAL
        // doubles readahead and drops already-consumed pages eagerly so
        // replay doesn't grow buff/cache by the full WAL size.
        let mut madvise_state = None;
        if !mmap.is_empty() {
            let rc = unsafe {
                libc::madvise(
                    mmap.as_ptr() as *mut libc::c_void,
                    mmap.len(),
                    libc::MADV_SEQUENTIAL,
                )
            };
            if rc == 0 {
                madvise_state = Some(libc::MADV_SEQUENTIAL);
                observability::MADV_SEQUENTIAL_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_SEQUENTIAL) failed on WAL segment; continuing",
                );
            }
        }

        Ok(Self {
            mmap,
            offset: 0,
            file,
            path: path.to_path_buf(),
            madvise_state,
        })
    }

    /// The madvise hint applied to the mapped segment (if any).
    pub fn madvise_state(&self) -> Option<libc::c_int> {
        self.madvise_state
    }

    /// Hint to the kernel that pages for this segment can be dropped from
    /// cache. Call this after a segment has been iterated end-to-end.
    pub fn release_pages(&self) {
        fadv_dontneed(&self.file, self.mmap.len(), &self.path);
    }

    /// Read the next record from the mmap'd region.
    ///
    /// Returns `None` at EOF or at the first corruption point.
    /// Zero-copy: payload bytes reference the mmap'd region directly.
    pub fn next_record(&mut self) -> Result<Option<WalRecord>> {
        let data = &self.mmap[..];

        loop {
            // Check if we have enough bytes for a header.
            if self.offset + HEADER_SIZE > data.len() {
                return Ok(None);
            }

            // Parse header.
            let header_bytes: &[u8; HEADER_SIZE] = data[self.offset..self.offset + HEADER_SIZE]
                .try_into()
                .map_err(|_| {
                    WalError::Io(std::io::Error::new(
                        std::io::ErrorKind::InvalidData,
                        "header slice conversion failed",
                    ))
                })?;
            let header = RecordHeader::from_bytes(header_bytes);

            // Validate magic — corruption or end of valid data.
            if header.magic != WAL_MAGIC {
                return Ok(None);
            }

            // Validate version.
            if header.validate(self.offset as u64).is_err() {
                return Ok(None);
            }

            let payload_len = header.payload_len as usize;
            let record_end = self.offset + HEADER_SIZE + payload_len;

            // Check if payload is fully within the mmap'd region.
            if record_end > data.len() {
                return Ok(None); // Torn write at segment end.
            }

            // Extract payload (copies from mmap to owned Vec).
            let payload = data[self.offset + HEADER_SIZE..record_end].to_vec();
            self.offset = record_end;

            let record = WalRecord { header, payload };

            // Verify checksum.
            if record.verify_checksum().is_err() {
                return Ok(None); // Corruption — end of committed prefix.
            }

            // Check record type.
            let logical_type = record.logical_record_type();
            if RecordType::from_raw(logical_type).is_none() {
                if RecordType::is_required(logical_type) {
                    return Err(WalError::UnknownRequiredRecordType {
                        record_type: header.record_type,
                        lsn: header.lsn,
                    });
                }
                // Unknown optional record — skip and continue loop.
                continue;
            }

            return Ok(Some(record));
        }
    }

    /// Iterator over all valid records in the mmap'd segment.
    pub fn records(self) -> MmapRecordIter {
        MmapRecordIter { reader: self }
    }

    /// Current read offset.
    pub fn offset(&self) -> usize {
        self.offset
    }

    /// Total size of the mmap'd region.
    pub fn len(&self) -> usize {
        self.mmap.len()
    }

    /// Whether the mmap'd region is empty.
    pub fn is_empty(&self) -> bool {
        self.mmap.is_empty()
    }
}

/// Iterator over records in a mmap'd WAL segment.
pub struct MmapRecordIter {
    reader: MmapWalReader,
}

impl Iterator for MmapRecordIter {
    type Item = Result<WalRecord>;

    fn next(&mut self) -> Option<Self::Item> {
        match self.reader.next_record() {
            Ok(Some(record)) => Some(Ok(record)),
            Ok(None) => None,
            Err(e) => Some(Err(e)),
        }
    }
}

/// Minimum number of segments to justify parallel replay overhead.
const PARALLEL_SEGMENT_THRESHOLD: usize = 4;

/// Replay WAL segments from a directory using mmap, starting from `from_lsn`.
///
/// Discovers all sealed segments, mmap's each, and returns records with
/// LSN >= `from_lsn`. This is the Event Plane's tier-2 catchup path.
///
/// When 4+ segments need scanning, uses `std::thread::scope` to read
/// segments in parallel (one thread per segment). Each thread mmap's its
/// segment and filters records independently; results are merged in
/// segment order (already LSN-sorted since segments are monotonic).
pub fn replay_segments_mmap(wal_dir: &Path, from_lsn: u64) -> Result<Vec<WalRecord>> {
    let segments = crate::segment::discover_segments(wal_dir)?;
    let live = filter_segments_by_lsn(&segments, from_lsn);

    if live.len() < PARALLEL_SEGMENT_THRESHOLD {
        return replay_segments_sequential(live, from_lsn);
    }

    replay_segments_parallel(live, from_lsn)
}

/// Return the slice of `segments` whose LSN range may contain records with
/// lsn >= `from_lsn`. A segment at index `i` is skippable iff the next
/// segment's `first_lsn` is `<= from_lsn` — meaning segment `i`'s entire
/// range is strictly below the cutoff. The last segment is never skipped
/// on this criterion because its upper bound is unknown.
fn filter_segments_by_lsn(
    segments: &[crate::segment::SegmentMeta],
    from_lsn: u64,
) -> &[crate::segment::SegmentMeta] {
    // Find the first segment whose next-segment first_lsn > from_lsn, OR
    // the last segment (always live). Since segments are LSN-sorted, the
    // live tail starts at the largest i such that segments[i].first_lsn
    // <= from_lsn.
    let mut start = 0;
    for i in 0..segments.len() {
        // Segment i covers [first_lsn_i, first_lsn_{i+1}).
        let upper = segments.get(i + 1).map(|s| s.first_lsn).unwrap_or(u64::MAX);
        if upper > from_lsn {
            start = i;
            break;
        }
        start = i + 1;
    }
    if start >= segments.len() {
        // All segments strictly below from_lsn; nothing to replay.
        return &[];
    }
    &segments[start..]
}

/// Sequential segment replay (used for small segment counts).
fn replay_segments_sequential(
    segments: &[crate::segment::SegmentMeta],
    from_lsn: u64,
) -> Result<Vec<WalRecord>> {
    let mut records = Vec::new();
    for seg in segments {
        let mut reader = MmapWalReader::open(&seg.path)?;
        while let Some(record) = reader.next_record()? {
            if record.header.lsn >= from_lsn {
                records.push(record);
            }
        }
        reader.release_pages();
    }
    Ok(records)
}

/// Parallel segment replay using scoped threads.
///
/// Each segment is read in its own thread via mmap. Since segments are
/// monotonically ordered by LSN, concatenating per-segment results in
/// segment order produces a globally LSN-ordered result.
fn replay_segments_parallel(
    segments: &[crate::segment::SegmentMeta],
    from_lsn: u64,
) -> Result<Vec<WalRecord>> {
    // Collect per-segment results. Index corresponds to segment order.
    let mut per_segment: Vec<Result<Vec<WalRecord>>> = Vec::with_capacity(segments.len());

    std::thread::scope(|scope| {
        let handles: Vec<_> = segments
            .iter()
            .map(|seg| {
                scope.spawn(move || -> Result<Vec<WalRecord>> {
                    let mut reader = MmapWalReader::open(&seg.path)?;
                    let mut seg_records = Vec::new();
                    while let Some(record) = reader.next_record()? {
                        if record.header.lsn >= from_lsn {
                            seg_records.push(record);
                        }
                    }
                    reader.release_pages();
                    Ok(seg_records)
                })
            })
            .collect();

        for handle in handles {
            per_segment.push(handle.join().unwrap_or_else(|_| {
                Err(WalError::Io(std::io::Error::other(
                    "segment replay thread panicked",
                )))
            }));
        }
    });

    // Merge in segment order (preserves LSN ordering).
    let total_estimate: usize = per_segment
        .iter()
        .map(|r| r.as_ref().map(|v| v.len()).unwrap_or(0))
        .sum();
    let mut records = Vec::with_capacity(total_estimate);
    for seg_result in per_segment {
        records.extend(seg_result?);
    }

    Ok(records)
}

/// Paginated mmap replay: reads at most `max_records` from `from_lsn`.
///
/// Returns `(records, has_more)` where `has_more` is `true` if the limit
/// was reached before all segments were exhausted. This bounds memory
/// usage per catch-up cycle to O(max_records) instead of O(all WAL data).
///
/// Always uses sequential reading (no parallel threads) since the bounded
/// record count makes parallel overhead unnecessary.
pub fn replay_segments_mmap_limit(
    wal_dir: &Path,
    from_lsn: u64,
    max_records: usize,
) -> Result<(Vec<WalRecord>, bool)> {
    let segments = crate::segment::discover_segments(wal_dir)?;
    let live = filter_segments_by_lsn(&segments, from_lsn);
    let mut records = Vec::with_capacity(max_records.min(4096));

    for seg in live {
        let mut reader = MmapWalReader::open(&seg.path)?;
        while let Some(record) = reader.next_record()? {
            if record.header.lsn >= from_lsn {
                records.push(record);
                if records.len() >= max_records {
                    // Partial scan — don't release pages for a segment
                    // we'll likely re-open on the next catchup cycle.
                    return Ok((records, true));
                }
            }
        }
        reader.release_pages();
    }

    Ok((records, false))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::record::RecordType;
    use crate::writer::{WalWriter, WalWriterConfig};

    fn test_writer(path: &Path) -> WalWriter {
        let config = WalWriterConfig {
            use_direct_io: false, // Tests run without O_DIRECT.
            ..Default::default()
        };
        WalWriter::open(path, config).unwrap()
    }

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

        // Write some records with the standard writer.
        {
            let mut writer = test_writer(&path);
            writer
                .append(RecordType::Put as u32, 1, 0, 0, b"hello")
                .unwrap();
            writer
                .append(RecordType::Put as u32, 1, 0, 0, b"world")
                .unwrap();
            writer.sync().unwrap();
        }

        // Read back with mmap reader.
        let reader = MmapWalReader::open(&path).unwrap();
        let records: Vec<WalRecord> = reader.records().collect::<Result<Vec<_>>>().unwrap();

        assert_eq!(records.len(), 2);
        assert_eq!(records[0].payload, b"hello");
        assert_eq!(records[1].payload, b"world");
    }

    #[test]
    fn mmap_reader_empty_file() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("empty.wal");
        std::fs::write(&path, []).unwrap();

        let reader = MmapWalReader::open(&path).unwrap();
        let records: Vec<WalRecord> = reader.records().collect::<Result<Vec<_>>>().unwrap();
        assert!(records.is_empty());
    }

    #[test]
    fn mmap_reader_truncated_header() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("truncated.wal");
        // Write 10 bytes — not enough for a header (30 bytes).
        std::fs::write(&path, [0u8; 10]).unwrap();

        let reader = MmapWalReader::open(&path).unwrap();
        let records: Vec<WalRecord> = reader.records().collect::<Result<Vec<_>>>().unwrap();
        assert!(records.is_empty());
    }

    #[test]
    fn replay_mmap_from_lsn() {
        let dir = tempfile::tempdir().unwrap();
        let wal_dir = dir.path().join("wal");
        std::fs::create_dir_all(&wal_dir).unwrap();

        let config = crate::segmented::SegmentedWalConfig::for_testing(wal_dir.clone());
        let mut wal = crate::segmented::SegmentedWal::open(config).unwrap();

        let lsn1 = wal.append(RecordType::Put as u32, 1, 0, 0, b"a").unwrap();
        let lsn2 = wal.append(RecordType::Put as u32, 1, 0, 0, b"b").unwrap();
        let lsn3 = wal.append(RecordType::Put as u32, 1, 0, 0, b"c").unwrap();
        wal.sync().unwrap();

        // Replay from lsn2 — should get records b and c.
        let records = replay_segments_mmap(&wal_dir, lsn2).unwrap();
        assert_eq!(records.len(), 2);
        assert_eq!(records[0].header.lsn, lsn2);
        assert_eq!(records[1].header.lsn, lsn3);

        // Replay from lsn1 — all 3.
        let all = replay_segments_mmap(&wal_dir, lsn1).unwrap();
        assert_eq!(all.len(), 3);
    }
}