mentedb_storage/
memory_map.rs

1//! Persisted memory map: which page holds which memory.
2//!
3//! Opening a database used to rebuild the memory id to page id map by loading
4//! every page, which on network storage costs minutes for large databases. The
5//! map is instead persisted at checkpoint time and replayed forward from the
6//! WAL on open, so a cold open reads only the pages the WAL touched.
7//!
8//! # Anchoring
9//!
10//! Each checkpoint appends a `Checkpoint` WAL entry whose payload is a random
11//! 16 byte token, then writes the snapshot stamped with the same token, then
12//! truncates the WAL to that entry. The snapshot is the map as of its
13//! checkpoint entry, so the map now is the snapshot plus every WAL entry after
14//! the entry carrying its token. A snapshot whose token is not in the WAL
15//! (the WAL was truncated by a checkpoint or recovery that did not write a
16//! snapshot, the snapshot write failed, or the file is from another database)
17//! can never be matched again, because tokens are never reused for a changed
18//! map. On any doubt (missing, corrupt, unmatched, inconsistent page count)
19//! the caller falls back to a full page scan.
20//!
21//! The table is keyed by page, mirroring exactly what a full scan observes:
22//! a page holds a memory when its last write decodes as one, and a freed page
23//! holds none. Each entry also carries a stamp, a random value refreshed on
24//! every write applied to that page, which lets derived caches (mode probes)
25//! persist alongside and detect, without loading the page, that the memory
26//! changed since they were captured.
27
28use std::collections::{BTreeMap, BTreeSet};
29use std::io::Write;
30use std::path::Path;
31
32use mentedb_core::MemoryNode;
33use mentedb_core::error::{MenteError, MenteResult};
34use mentedb_core::types::MemoryId;
35
36use crate::page::{PAGE_DATA_SIZE, PageId};
37use crate::wal::{WalEntry, WalEntryType};
38
39/// File name of the persisted memory map inside the database directory.
40pub const MEMORY_MAP_FILE: &str = "memory_map.bin";
41
42const MAGIC: &[u8; 8] = b"MDBMMAP1";
43const FORMAT_VERSION: u32 = 1;
44
45/// Checkpoint token carried in the `Checkpoint` WAL entry payload.
46pub(crate) type Token = [u8; 16];
47
48/// A memory the map knows about: its id, its page, and the page's write stamp.
49#[derive(Debug, Clone, Copy, PartialEq, Eq)]
50pub struct MappedMemory {
51    pub id: MemoryId,
52    pub page: PageId,
53    /// Changes whenever the page is written. Equal stamps for the same page
54    /// mean no write was applied to it in between.
55    pub stamp: u64,
56}
57
58/// Decode a memory node from a page's data region, with exactly the rules
59/// [`crate::StorageEngine::load_memory`] applies after the page type check.
60/// `data` may be shorter than a page; the missing tail reads as zeros.
61pub(crate) fn decode_memory(data: &[u8]) -> MenteResult<MemoryNode> {
62    let data = &data[..data.len().min(PAGE_DATA_SIZE)];
63    let mut prefix = [0u8; 4];
64    let n = data.len().min(4);
65    prefix[..n].copy_from_slice(&data[..n]);
66    let len = u32::from_le_bytes(prefix) as usize;
67    if len == 0 || len + 4 > PAGE_DATA_SIZE {
68        return Err(MenteError::Storage(format!(
69            "invalid memory node length prefix: {len}"
70        )));
71    }
72    if 4 + len > data.len() {
73        // The page tail is zero filled, and JSON never ends in NUL bytes, so
74        // the page would not decode either.
75        return Err(MenteError::Serialization(
76            "memory node runs past the written data".into(),
77        ));
78    }
79    serde_json::from_slice(&data[4..4 + len]).map_err(|e| MenteError::Serialization(e.to_string()))
80}
81
82/// A fresh random stamp or token half.
83fn random_u64() -> u64 {
84    let (a, b) = uuid::Uuid::new_v4().as_u64_pair();
85    a ^ b.rotate_left(17)
86}
87
88pub(crate) fn new_token() -> Token {
89    *uuid::Uuid::new_v4().as_bytes()
90}
91
92/// The stamp of the WAL entry at `lsn` written while `token` anchored the
93/// map. Deterministic, so the stamp a write gets live, the stamp a checkpoint
94/// gives it rolling the map forward, and the stamp a reopen gives it replaying
95/// the WAL are the same value. LSNs only grow while a token anchors the map,
96/// so (token, lsn) never repeats.
97pub(crate) fn entry_stamp(token: &Token, lsn: u64) -> u64 {
98    fn mix(mut z: u64) -> u64 {
99        z = z.wrapping_add(0x9E37_79B9_7F4A_7C15);
100        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
101        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
102        z ^ (z >> 31)
103    }
104    let mut a = [0u8; 8];
105    let mut b = [0u8; 8];
106    a.copy_from_slice(&token[..8]);
107    b.copy_from_slice(&token[8..]);
108    mix(u64::from_le_bytes(a) ^ mix(u64::from_le_bytes(b) ^ mix(lsn)))
109}
110
111/// The page level view of which memories live where.
112#[derive(Debug, Clone, Default, PartialEq, Eq)]
113pub(crate) struct MemoryTable {
114    /// page -> (memory id, write stamp), for pages holding a decodable memory.
115    pages: BTreeMap<u64, (MemoryId, u64)>,
116    /// Pages whose last logged operation freed them.
117    free: BTreeSet<u64>,
118}
119
120impl MemoryTable {
121    /// Record that `page` now holds `id` (or no memory when `None`), written
122    /// with `stamp`.
123    pub(crate) fn set_page(&mut self, page: u64, id: Option<MemoryId>, stamp: u64) {
124        self.free.remove(&page);
125        match id {
126            Some(id) => {
127                self.pages.insert(page, (id, stamp));
128            }
129            None => {
130                self.pages.remove(&page);
131            }
132        }
133    }
134
135    /// Record that `page` was freed.
136    pub(crate) fn set_free(&mut self, page: u64) {
137        self.pages.remove(&page);
138        self.free.insert(page);
139    }
140
141    /// Record a page found by a full scan. There is no WAL entry to derive a
142    /// stamp from, so it gets a fresh random one.
143    pub(crate) fn set_scanned(&mut self, page: u64, id: MemoryId) {
144        self.set_page(page, Some(id), random_u64());
145    }
146
147    /// Apply one WAL entry written under `token`. Returns true when the entry
148    /// can change the map.
149    pub(crate) fn apply(&mut self, entry: &WalEntry, token: &Token) -> bool {
150        match entry.entry_type {
151            WalEntryType::PageWrite => {
152                let id = decode_memory(&entry.data).ok().map(|n| n.id);
153                self.set_page(entry.page_id, id, entry_stamp(token, entry.lsn));
154                true
155            }
156            WalEntryType::PageFree => {
157                self.set_free(entry.page_id);
158                true
159            }
160            WalEntryType::Checkpoint | WalEntryType::Commit => false,
161        }
162    }
163
164    /// Pages whose last logged operation freed them.
165    pub(crate) fn free_pages(&self) -> &BTreeSet<u64> {
166        &self.free
167    }
168
169    /// Highest page id referenced, if any.
170    fn max_page(&self) -> Option<u64> {
171        let a = self.pages.keys().next_back().copied();
172        let b = self.free.iter().next_back().copied();
173        a.max(b)
174    }
175
176    /// Every mapped memory in page order, the order a full scan returns.
177    pub(crate) fn memories(&self) -> Vec<MappedMemory> {
178        self.pages
179            .iter()
180            .map(|(&page, &(id, stamp))| MappedMemory {
181                id,
182                page: PageId(page),
183                stamp,
184            })
185            .collect()
186    }
187
188    /// The memory and stamp recorded for `page`.
189    pub(crate) fn get(&self, page: u64) -> Option<(MemoryId, u64)> {
190        self.pages.get(&page).copied()
191    }
192}
193
194/// The memory table as of the checkpoint entry carrying `token`.
195#[derive(Debug, Clone)]
196pub(crate) struct MapAnchor {
197    pub(crate) token: Token,
198    /// Page count when the snapshot was taken; the page file never shrinks,
199    /// so a smaller count on open means the files do not belong together.
200    pub(crate) page_count: u64,
201    pub(crate) table: MemoryTable,
202    /// Writes since the anchor's checkpoint, recorded as they happen (page ->
203    /// memory and stamp, None when freed or not a memory), so stamp lookups
204    /// see the current write. Not persisted: the WAL holds the same writes and
205    /// rolling forward reproduces these exact stamps. A write path that does
206    /// not record here only makes lookups older, which callers treat as stale.
207    pub(crate) pending: std::collections::HashMap<u64, Option<(MemoryId, u64)>>,
208}
209
210impl MapAnchor {
211    /// Index of the last `Checkpoint` entry carrying this anchor's token.
212    fn position(&self, entries: &[WalEntry]) -> Option<usize> {
213        entries
214            .iter()
215            .rposition(|e| e.entry_type == WalEntryType::Checkpoint && e.data == self.token)
216    }
217
218    /// Roll the table forward over the entries after this anchor's checkpoint.
219    /// Returns None when the anchor's checkpoint is not in `entries` (the map
220    /// cannot be derived), else whether any entry changed the table.
221    pub(crate) fn advance(&mut self, entries: &[WalEntry]) -> Option<bool> {
222        let pos = self.position(entries)?;
223        let mut changed = false;
224        for e in &entries[pos + 1..] {
225            changed |= self.table.apply(e, &self.token);
226        }
227        self.pending.clear();
228        Some(changed)
229    }
230
231    /// Record a write logged at `lsn` under this anchor.
232    pub(crate) fn note_write(&mut self, page: u64, id: Option<MemoryId>, lsn: u64) {
233        let stamp = entry_stamp(&self.token, lsn);
234        self.pending.insert(page, id.map(|i| (i, stamp)));
235    }
236
237    /// The memory and stamp `page` holds now.
238    pub(crate) fn current(&self, page: u64) -> Option<(MemoryId, u64)> {
239        match self.pending.get(&page) {
240            Some(v) => *v,
241            None => self.table.get(page),
242        }
243    }
244
245    /// Whether the snapshot is internally consistent with a page file of
246    /// `page_count` pages.
247    pub(crate) fn fits(&self, page_count: u64) -> bool {
248        self.page_count <= page_count
249            && self.table.max_page().is_none_or(|m| m < self.page_count)
250            && !self.table.pages.contains_key(&0)
251            && !self.table.free.contains(&0)
252    }
253}
254
255fn put_u64(buf: &mut Vec<u8>, v: u64) {
256    buf.extend_from_slice(&v.to_le_bytes());
257}
258
259fn encode(anchor: &MapAnchor) -> Vec<u8> {
260    let t = &anchor.table;
261    let mut buf = Vec::with_capacity(64 + t.pages.len() * 32 + t.free.len() * 8);
262    buf.extend_from_slice(MAGIC);
263    buf.extend_from_slice(&FORMAT_VERSION.to_le_bytes());
264    buf.extend_from_slice(&anchor.token);
265    put_u64(&mut buf, anchor.page_count);
266    put_u64(&mut buf, t.pages.len() as u64);
267    for (&page, &(id, stamp)) in &t.pages {
268        put_u64(&mut buf, page);
269        buf.extend_from_slice(id.0.as_bytes());
270        put_u64(&mut buf, stamp);
271    }
272    put_u64(&mut buf, t.free.len() as u64);
273    for &page in &t.free {
274        put_u64(&mut buf, page);
275    }
276    let crc = crc32fast::hash(&buf);
277    buf.extend_from_slice(&crc.to_le_bytes());
278    buf
279}
280
281struct Reader<'a> {
282    buf: &'a [u8],
283    pos: usize,
284}
285
286impl Reader<'_> {
287    fn take(&mut self, n: usize) -> MenteResult<&[u8]> {
288        let end = self
289            .pos
290            .checked_add(n)
291            .filter(|&e| e <= self.buf.len())
292            .ok_or_else(|| MenteError::Storage("memory map snapshot truncated".into()))?;
293        let s = &self.buf[self.pos..end];
294        self.pos = end;
295        Ok(s)
296    }
297
298    fn u64(&mut self) -> MenteResult<u64> {
299        let mut b = [0u8; 8];
300        b.copy_from_slice(self.take(8)?);
301        Ok(u64::from_le_bytes(b))
302    }
303}
304
305fn decode(bytes: &[u8]) -> MenteResult<MapAnchor> {
306    let bad = |what: &str| MenteError::Storage(format!("memory map snapshot {what}"));
307    if bytes.len() < MAGIC.len() + 4 + 16 + 8 + 8 + 8 + 4 {
308        return Err(bad("too short"));
309    }
310    let (body, crc_bytes) = bytes.split_at(bytes.len() - 4);
311    let mut crc = [0u8; 4];
312    crc.copy_from_slice(crc_bytes);
313    if crc32fast::hash(body) != u32::from_le_bytes(crc) {
314        return Err(bad("checksum mismatch"));
315    }
316    let mut r = Reader { buf: body, pos: 0 };
317    if r.take(8)? != MAGIC {
318        return Err(bad("magic mismatch"));
319    }
320    let mut v = [0u8; 4];
321    v.copy_from_slice(r.take(4)?);
322    if u32::from_le_bytes(v) != FORMAT_VERSION {
323        return Err(bad("version mismatch"));
324    }
325    let mut token = [0u8; 16];
326    token.copy_from_slice(r.take(16)?);
327    let page_count = r.u64()?;
328    let n_pages = r.u64()?;
329    let mut table = MemoryTable::default();
330    for _ in 0..n_pages {
331        let page = r.u64()?;
332        let id = uuid::Uuid::from_slice(r.take(16)?).map_err(|e| bad(&e.to_string()))?;
333        let stamp = r.u64()?;
334        table.pages.insert(page, (MemoryId(id), stamp));
335    }
336    let n_free = r.u64()?;
337    for _ in 0..n_free {
338        table.free.insert(r.u64()?);
339    }
340    if r.pos != body.len() {
341        return Err(bad("has trailing bytes"));
342    }
343    Ok(MapAnchor {
344        token,
345        page_count,
346        table,
347        pending: Default::default(),
348    })
349}
350
351/// Read the snapshot in `dir`. `Ok(None)` when there is none.
352pub(crate) fn read_snapshot(dir: &Path) -> MenteResult<Option<MapAnchor>> {
353    match std::fs::read(dir.join(MEMORY_MAP_FILE)) {
354        Ok(bytes) => decode(&bytes).map(Some),
355        Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
356        Err(e) => Err(e.into()),
357    }
358}
359
360/// Write the snapshot atomically: a reader sees the old file or the complete
361/// new one, never a mix (temp file, fsync, rename).
362pub(crate) fn write_snapshot(dir: &Path, anchor: &MapAnchor) -> MenteResult<()> {
363    let path = dir.join(MEMORY_MAP_FILE);
364    let tmp = dir.join(format!("{MEMORY_MAP_FILE}.tmp"));
365    {
366        let mut file = std::fs::File::create(&tmp)?;
367        file.write_all(&encode(anchor))?;
368        file.sync_all()?;
369    }
370    std::fs::rename(&tmp, &path)?;
371    Ok(())
372}
373
374/// Remove the snapshot, ignoring a missing file.
375pub(crate) fn remove_snapshot(dir: &Path) {
376    let _ = std::fs::remove_file(dir.join(MEMORY_MAP_FILE));
377}
378
379#[cfg(test)]
380mod tests {
381    use super::*;
382
383    fn anchor() -> MapAnchor {
384        let mut table = MemoryTable::default();
385        table.set_scanned(1, MemoryId::new());
386        table.set_scanned(3, MemoryId::new());
387        table.set_free(2);
388        MapAnchor {
389            token: new_token(),
390            page_count: 4,
391            table,
392            pending: Default::default(),
393        }
394    }
395
396    #[test]
397    fn snapshot_roundtrips() {
398        let dir = tempfile::tempdir().unwrap();
399        let a = anchor();
400        write_snapshot(dir.path(), &a).unwrap();
401        let b = read_snapshot(dir.path()).unwrap().unwrap();
402        assert_eq!(a.token, b.token);
403        assert_eq!(a.page_count, b.page_count);
404        assert_eq!(a.table, b.table);
405        assert!(!dir.path().join(format!("{MEMORY_MAP_FILE}.tmp")).exists());
406    }
407
408    #[test]
409    fn corrupt_snapshot_is_rejected() {
410        let dir = tempfile::tempdir().unwrap();
411        write_snapshot(dir.path(), &anchor()).unwrap();
412        let path = dir.path().join(MEMORY_MAP_FILE);
413        let mut bytes = std::fs::read(&path).unwrap();
414        let mid = bytes.len() / 2;
415        bytes[mid] ^= 0xff;
416        std::fs::write(&path, &bytes).unwrap();
417        assert!(read_snapshot(dir.path()).is_err());
418        std::fs::write(&path, &bytes[..10]).unwrap();
419        assert!(read_snapshot(dir.path()).is_err());
420        std::fs::remove_file(&path).unwrap();
421        assert!(read_snapshot(dir.path()).unwrap().is_none());
422    }
423
424    #[test]
425    fn fits_rejects_pages_past_the_count() {
426        let a = anchor();
427        assert!(a.fits(4));
428        assert!(a.fits(10));
429        assert!(!a.fits(3), "page file smaller than the snapshot");
430        let mut b = anchor();
431        b.page_count = 3;
432        assert!(!b.fits(10), "page 3 recorded beyond its own page count");
433    }
434}
mentedb_storage/memory_map.rs

mentedb_storage/
memory_map.rs
