embedmind-core 0.1.0

Embedded memory engine for AI agents: single crash-safe file, vector + full-text + graph
//! WAL sidecar management: append, recovery scan, reset.
//!
//! The WAL is a physical page-level redo log (`docs/adr/0001`); its byte
//! layout lives in [`crate::format`] (`docs/FORMAT.md` §8). This module owns
//! the file lifecycle: a transaction is durable iff its commit frame is fully
//! valid on disk, and recovery applies exactly the committed prefix, in
//! order, discarding any torn tail.

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

use crate::error::Result;
use crate::format::{
    WAL_FRAME_HEADER_LEN, WAL_HEADER_LEN, WalFrameHeader, WalHeader, page_size_is_valid,
};
use crate::storage::vfs::VfsFile;

/// Result of scanning an existing WAL file on open.
#[derive(Debug)]
pub struct RecoveredWal {
    /// Page size recorded in the WAL header.
    pub page_size: u32,
    /// Latest committed image of every page present in the valid prefix:
    /// `page_no → byte offset of the page image` inside the WAL file.
    /// Built by applying committed transactions in order, so later
    /// transactions override earlier ones.
    pub committed: BTreeMap<u64, u64>,
}

/// Scans a WAL file and returns its committed page images.
///
/// Returns `Ok(None)` when there is no usable WAL (missing/short/invalid
/// header) — that is *not* an error: a torn WAL header can only exist if no
/// commit was ever acknowledged from that generation, so there is nothing to
/// lose. Frames after the first invalid frame (torn tail) are ignored, as are
/// trailing frames of a transaction with no valid commit frame.
pub fn scan(file: &dyn VfsFile) -> Result<Option<RecoveredWal>> {
    let file_len = file.len()?;
    if file_len < WAL_HEADER_LEN as u64 {
        return Ok(None);
    }
    let mut header_buf = [0u8; WAL_HEADER_LEN];
    file.read_at(&mut header_buf, 0)?;
    let Some(header) = WalHeader::decode(&header_buf) else {
        return Ok(None);
    };

    let page_size = header.page_size as u64;
    let frame_len = WAL_FRAME_HEADER_LEN as u64 + page_size;
    let mut committed = BTreeMap::new();
    let mut pending: Vec<(u64, u64)> = Vec::new();
    let mut pending_txn: Option<u64> = None;
    let mut offset = WAL_HEADER_LEN as u64;
    let mut frame_header = [0u8; WAL_FRAME_HEADER_LEN];
    let mut image = vec![0u8; header.page_size as usize];

    while let Some(end) = offset.checked_add(frame_len) {
        if end > file_len {
            break; // partial frame at the tail
        }
        file.read_at(&mut frame_header, offset)?;
        file.read_at(&mut image, offset + WAL_FRAME_HEADER_LEN as u64)?;
        let Some(frame) = WalFrameHeader::decode(&frame_header, &image, header.salt) else {
            break; // torn/corrupt frame ends the valid prefix
        };
        if pending_txn.is_some_and(|t| t != frame.txn_id) {
            break; // interleaved transactions never happen in a well-formed log
        }
        pending.push((frame.page_no, offset + WAL_FRAME_HEADER_LEN as u64));
        if frame.commit {
            committed.extend(pending.drain(..));
            pending_txn = None;
        } else {
            pending_txn = Some(frame.txn_id);
        }
        offset = end;
    }

    Ok(Some(RecoveredWal {
        page_size: header.page_size,
        committed,
    }))
}

/// An open WAL file, ready for appends. Created over an **empty** file
/// (recovery and reset happen before construction, in the pager).
pub struct Wal {
    file: Box<dyn VfsFile>,
    page_size: u32,
    /// Salt of the current generation; `None` until the first append writes
    /// a fresh WAL header.
    salt: Option<u64>,
    /// Logical append offset (bytes of valid content).
    len: u64,
}

impl Wal {
    /// Wraps an empty WAL file.
    pub fn new(file: Box<dyn VfsFile>, page_size: u32) -> Self {
        debug_assert!(page_size_is_valid(page_size));
        Wal {
            file,
            page_size,
            salt: None,
            len: 0,
        }
    }

    /// Current logical size in bytes (checkpoint-threshold input).
    pub fn size(&self) -> u64 {
        self.len
    }

    /// Appends one transaction: one frame per page, the last one carrying the
    /// commit flag, then fsyncs (`docs/FORMAT.md` §8 commit protocol). The
    /// transaction is durable iff this returns `Ok`. Returns
    /// `(page_no, image offset)` for each page, for the pager's WAL index.
    ///
    /// `pages` must be non-empty and each image exactly one page long.
    pub fn append_txn(&mut self, txn_id: u64, pages: &[(u64, &[u8])]) -> Result<Vec<(u64, u64)>> {
        let salt = match self.salt {
            Some(s) => s,
            None => {
                // First append of a generation: fresh salt, fresh header.
                // Random salt prevents stale frames from a previous, longer
                // generation from ever validating (§8).
                let salt = generate_salt();
                let header = WalHeader {
                    format_version: crate::format::FORMAT_VERSION,
                    page_size: self.page_size,
                    salt,
                };
                self.file.write_at(&header.encode(), 0)?;
                self.salt = Some(salt);
                self.len = WAL_HEADER_LEN as u64;
                salt
            }
        };

        let mut offsets = Vec::with_capacity(pages.len());
        let mut frame = Vec::with_capacity(WAL_FRAME_HEADER_LEN + self.page_size as usize);
        for (i, (page_no, image)) in pages.iter().enumerate() {
            if image.len() != self.page_size as usize {
                return Err(crate::Error::InvalidArgument("WAL image must be one page"));
            }
            let header = WalFrameHeader {
                page_no: *page_no,
                txn_id,
                commit: i + 1 == pages.len(),
            };
            frame.clear();
            frame.extend_from_slice(&header.encode(image, salt));
            frame.extend_from_slice(image);
            self.file.write_at(&frame, self.len)?;
            offsets.push((*page_no, self.len + WAL_FRAME_HEADER_LEN as u64));
            self.len += frame.len() as u64;
        }
        self.file.sync()?;
        Ok(offsets)
    }

    /// Reads a page image previously appended (offset from [`Wal::append_txn`]
    /// or [`scan`]).
    pub fn read_image(&self, offset: u64, buf: &mut [u8]) -> Result<()> {
        self.file.read_at(buf, offset)?;
        Ok(())
    }

    /// Ends the current generation: truncates to zero and fsyncs. The next
    /// append starts a new generation with a new salt.
    pub fn reset(&mut self) -> Result<()> {
        self.file.truncate(0)?;
        self.file.sync()?;
        self.salt = None;
        self.len = 0;
        Ok(())
    }
}

/// Per-generation WAL salt. Does not need to be cryptographically strong —
/// its only job is making frames from a previous generation fail their
/// checksum (anti-replay), so time + pid + a process-wide counter hashed
/// together is plenty, with zero added dependencies.
fn generate_salt() -> u64 {
    static COUNTER: AtomicU64 = AtomicU64::new(0);
    let nanos = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_nanos() as u64)
        .unwrap_or(0);
    let mut hasher = xxhash_rust::xxh3::Xxh3::with_seed(nanos);
    hasher.update(&std::process::id().to_le_bytes());
    hasher.update(&COUNTER.fetch_add(1, Ordering::Relaxed).to_le_bytes());
    hasher.digest()
}

#[cfg(test)]
mod tests {
    #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]

    use super::*;
    use crate::format::stamp_page_checksum;
    use crate::storage::sim::SimVfs;
    use crate::storage::vfs::{OpenMode, Vfs};
    use std::path::Path;

    const PS: u32 = 512;

    fn page(fill: u8) -> Vec<u8> {
        let mut p = vec![fill; PS as usize];
        stamp_page_checksum(&mut p);
        p
    }

    fn open(vfs: &SimVfs) -> Box<dyn VfsFile> {
        vfs.open(Path::new("test-wal"), OpenMode::OpenOrCreate)
            .unwrap()
    }

    #[test]
    fn scan_empty_and_garbage_wal() {
        let vfs = SimVfs::new();
        let f = open(&vfs);
        assert!(scan(&*f).unwrap().is_none()); // empty
        f.write_at(b"garbage that is longer than a wal header....", 0)
            .unwrap();
        assert!(scan(&*f).unwrap().is_none()); // invalid header
    }

    #[test]
    fn append_scan_roundtrip_with_override() {
        let vfs = SimVfs::new();
        let mut wal = Wal::new(open(&vfs), PS);
        let (a, b, c) = (page(1), page(2), page(3));
        wal.append_txn(1, &[(5, &a), (0, &b)]).unwrap();
        wal.append_txn(2, &[(5, &c)]).unwrap(); // overrides page 5

        let rec = scan(&*open(&vfs)).unwrap().unwrap();
        assert_eq!(rec.page_size, PS);
        assert_eq!(rec.committed.len(), 2);
        let mut buf = vec![0u8; PS as usize];
        wal.read_image(rec.committed[&5], &mut buf).unwrap();
        assert_eq!(buf, c);
        wal.read_image(rec.committed[&0], &mut buf).unwrap();
        assert_eq!(buf, b);
    }

    #[test]
    fn torn_tail_is_dropped_committed_prefix_survives() {
        let vfs = SimVfs::new();
        let mut wal = Wal::new(open(&vfs), PS);
        wal.append_txn(1, &[(1, &page(1))]).unwrap();
        wal.append_txn(2, &[(2, &page(2))]).unwrap();

        // Corrupt one byte inside txn 2's frame: prefix = txn 1 only.
        let f = open(&vfs);
        let txn2_frame_start = WAL_HEADER_LEN as u64 + (WAL_FRAME_HEADER_LEN as u64 + PS as u64);
        let mut byte = [0u8; 1];
        f.read_at(&mut byte, txn2_frame_start + 40).unwrap();
        byte[0] ^= 0xff;
        f.write_at(&byte, txn2_frame_start + 40).unwrap();

        let rec = scan(&*f).unwrap().unwrap();
        assert_eq!(rec.committed.keys().copied().collect::<Vec<_>>(), vec![1]);
    }

    #[test]
    fn uncommitted_trailing_frames_are_ignored() {
        let vfs = SimVfs::new();
        let mut wal = Wal::new(open(&vfs), PS);
        wal.append_txn(1, &[(1, &page(1))]).unwrap();
        // Simulate a crash mid-transaction: append a valid frame with no
        // commit flag by writing it manually.
        let img = page(9);
        let fh = WalFrameHeader {
            page_no: 3,
            txn_id: 2,
            commit: false,
        };
        let salt_hdr = {
            let mut b = [0u8; WAL_HEADER_LEN];
            open(&vfs).read_at(&mut b, 0).unwrap();
            WalHeader::decode(&b).unwrap().salt
        };
        let f = open(&vfs);
        let off = f.len().unwrap();
        f.write_at(&fh.encode(&img, salt_hdr), off).unwrap();
        f.write_at(&img, off + WAL_FRAME_HEADER_LEN as u64).unwrap();

        let rec = scan(&*f).unwrap().unwrap();
        assert_eq!(rec.committed.keys().copied().collect::<Vec<_>>(), vec![1]);
    }

    #[test]
    fn reset_starts_a_new_generation_old_frames_never_replay() {
        let vfs = SimVfs::new();
        let mut wal = Wal::new(open(&vfs), PS);
        wal.append_txn(1, &[(1, &page(1))]).unwrap();
        wal.append_txn(2, &[(2, &page(2))]).unwrap();
        let long = open(&vfs).len().unwrap();
        wal.reset().unwrap();

        // New generation writes one shorter txn; then simulate the truncate
        // being lost EXCEPT the new header+frame (worst realistic case is
        // covered by the crash harness; here we check the salt logic itself:
        // stale bytes beyond the new generation's content must not validate).
        wal.append_txn(1, &[(3, &page(3))]).unwrap();
        let f = open(&vfs);
        assert!(f.len().unwrap() < long);
        let rec = scan(&*f).unwrap().unwrap();
        assert_eq!(rec.committed.keys().copied().collect::<Vec<_>>(), vec![3]);
    }
}