Skip to main content

mentedb_storage/
engine.rs

1//! Storage Engine: facade that ties the page manager, WAL, and buffer pool together.
2
3use std::fs::File;
4use std::path::{Path, PathBuf};
5
6use mentedb_core::MemoryNode;
7use mentedb_core::error::{MenteError, MenteResult};
8
9use parking_lot::Mutex;
10use tracing::info;
11
12use crate::buffer::BufferPool;
13use crate::memory_map::{self, MapAnchor, MappedMemory, MemoryTable};
14use crate::page::{PAGE_DATA_SIZE, Page, PageId, PageManager, PageType};
15use crate::wal::{Wal, WalEntryType};
16/// Default number of page frames in the buffer pool.
17const DEFAULT_BUFFER_POOL_SIZE: usize = 1024;
18/// Data file of the record storage format that replaces the page file. Its
19/// presence means the directory was migrated by a newer engine.
20const NEWER_FORMAT_FILE: &str = "records.db";
21
22/// Auto-checkpoint when WAL file exceeds this size (8 MB).
23const WAL_AUTO_CHECKPOINT_BYTES: u64 = 8 * 1024 * 1024;
24
25/// The unified storage engine for MenteDB.
26///
27/// Coordinates page allocation, caching, and write-ahead logging to provide
28/// crash-safe, page-oriented storage for memory nodes.
29///
30/// Concurrency model (inspired by WAL-mode databases):
31/// - **Reads are lock-free**: `read_page` only touches the buffer pool and page
32///   manager — no file locks, no WAL access.
33/// - **Writes are fully serialized** via a blocking `flock(2)` on the WAL file.
34///   The entire write transaction (page allocation + WAL append + page write +
35///   fsync) executes under a single flock, ensuring correctness across multiple
36///   processes sharing the same data directory.
37/// - **State is refreshed from disk** under the flock: page count is re-read
38///   from the file header and LSN is re-read from the WAL tail, so no process
39///   can act on stale in-memory state.
40/// - **One process per database.** Open acquires an exclusive lock on a LOCK
41///   file in the data directory and holds it for the engine's lifetime. The
42///   per operation flock protects individual write transactions, but layers
43///   above the storage engine cache state across operations (buffer pool,
44///   page map, index snapshots), and a second concurrent process flushing
45///   those caches rolls the database back to its own open time snapshot.
46///   A second open therefore fails fast with a "locked" error instead of
47///   silently corrupting state; callers retry until the owner exits.
48pub struct StorageEngine {
49    page_manager: Mutex<PageManager>,
50    buffer_pool: BufferPool,
51    wal: Mutex<Wal>,
52    /// Held exclusively for the lifetime of this engine; released on close
53    /// or process exit. Guards against concurrent multi process opens.
54    process_lock: Mutex<Option<File>>,
55    /// Database directory (holds the memory map snapshot).
56    dir: PathBuf,
57    /// The memory map as of the checkpoint entry carrying its token, when
58    /// known. None until a valid snapshot is loaded or a full scan builds one;
59    /// every checkpoint rolls it forward and persists it. Lock order: after
60    /// the WAL and page manager mutexes.
61    map_anchor: Mutex<Option<MapAnchor>>,
62    /// Memory loads served since open, cached or not (metrics and tests).
63    memory_loads: std::sync::atomic::AtomicU64,
64}
65
66impl StorageEngine {
67    /// Open (or create) a storage engine rooted at `path`.
68    ///
69    /// `path` must be a directory; it will be created if it does not exist.
70    /// After opening, any uncommitted WAL entries are replayed for crash recovery.
71    ///
72    /// # Example
73    ///
74    /// ```no_run
75    /// use mentedb_storage::StorageEngine;
76    ///
77    /// let engine = StorageEngine::open("/tmp/mentedb".as_ref())?;
78    /// // engine is ready — WAL recovery already ran if needed
79    /// # Ok::<(), mentedb_core::error::MenteError>(())
80    /// ```
81    pub fn open(path: &Path) -> MenteResult<Self> {
82        std::fs::create_dir_all(path)?;
83
84        // Exclusive process lock, held until close or process exit. A second
85        // process opening the same directory would interleave stale cached
86        // state with ours and roll the database back, so refuse it up front.
87        let lock_path = path.join("LOCK");
88        let lock_file = std::fs::OpenOptions::new()
89            .create(true)
90            .truncate(false)
91            .write(true)
92            .open(&lock_path)?;
93        // Cross-host lock: fcntl OFD lock on Linux (enforced across hosts by the
94        // NFSv4 lock manager), flock elsewhere. This makes single-writer safe
95        // when the directory lives on a network filesystem mounted by more than
96        // one host.
97        match crate::lock::try_lock_exclusive(&lock_file) {
98            Ok(true) => {}
99            Ok(false) => {
100                return Err(MenteError::Storage(format!(
101                    "database directory {} is locked by another process",
102                    path.display()
103                )));
104            }
105            Err(e) => {
106                return Err(MenteError::Storage(format!(
107                    "failed to lock database directory {}: {e}",
108                    path.display()
109                )));
110            }
111        }
112
113        // A newer engine migrates this layout into a record file and moves the
114        // page file aside. Opening such a directory here would find no page
115        // file and create an empty database, silently hiding every memory and
116        // taking writes that the newer engine never sees. Refuse instead, so a
117        // rollback or a mixed version fleet fails loudly and loses nothing.
118        if path.join(NEWER_FORMAT_FILE).exists() {
119            return Err(MenteError::Storage(format!(
120                "database directory {} uses a newer storage format ({NEWER_FORMAT_FILE}); \
121                 open it with a newer version of mentedb",
122                path.display()
123            )));
124        }
125
126        let page_manager = PageManager::open(path)?;
127        let buffer_pool = BufferPool::new(DEFAULT_BUFFER_POOL_SIZE);
128        let wal = Wal::open(path)?;
129
130        let engine = Self {
131            page_manager: Mutex::new(page_manager),
132            buffer_pool,
133            wal: Mutex::new(wal),
134            process_lock: Mutex::new(Some(lock_file)),
135            dir: path.to_path_buf(),
136            map_anchor: Mutex::new(None),
137            memory_loads: std::sync::atomic::AtomicU64::new(0),
138        };
139
140        let recovered = engine.recover()?;
141        if recovered > 0 {
142            info!(recovered, ?path, "storage engine opened with WAL recovery");
143        } else {
144            info!(?path, "storage engine opened");
145        }
146
147        Ok(engine)
148    }
149
150    /// Replay WAL entries to recover writes that were not checkpointed.
151    ///
152    /// For each `PageWrite` entry the serialized data is written back to its page.
153    /// After replay the WAL is truncated. Returns the number of entries replayed.
154    pub fn recover(&self) -> MenteResult<usize> {
155        let mut wal = self.wal.lock();
156        wal.lock_exclusive()?;
157        let result = self.recover_locked(&mut wal);
158        let unlocked = wal.unlock();
159        let count = result?;
160        unlocked?;
161        Ok(count)
162    }
163
164    /// Load the persisted memory map and roll it forward over `entries`, or
165    /// None when it is missing, unreadable, or does not belong to this WAL
166    /// and page file. Never an error: the fallback is a full scan.
167    fn load_anchor(&self, entries: &[crate::wal::WalEntry], page_count: u64) -> Option<MapAnchor> {
168        let mut anchor = match memory_map::read_snapshot(&self.dir) {
169            Ok(Some(a)) => a,
170            Ok(None) => return None,
171            Err(e) => {
172                tracing::warn!(error = %e, "memory map snapshot unreadable, open will scan pages");
173                return None;
174            }
175        };
176        if !anchor.fits(page_count) {
177            tracing::warn!(
178                snapshot_pages = anchor.page_count,
179                page_count,
180                "memory map snapshot does not fit the page file, open will scan pages"
181            );
182            return None;
183        }
184        if anchor.advance(entries).is_none() {
185            info!("memory map snapshot predates the WAL, open will scan pages");
186            return None;
187        }
188        Some(anchor)
189    }
190
191    fn recover_locked(&self, wal: &mut Wal) -> MenteResult<usize> {
192        let entries = wal.iterate()?;
193        let mut count = 0usize;
194        let mut pm = self.page_manager.lock();
195
196        // Refresh page count from disk — another process may have written pages.
197        pm.reload_header()?;
198
199        // Roll the persisted memory map forward before replay may truncate
200        // the WAL entries it is anchored on.
201        let anchor = self.load_anchor(&entries, pm.page_count());
202
203        // Replay with last-op-wins per page: every PageWrite carries a full page
204        // image and PageFree discards the page, so only the final entry for each
205        // page matters. This is what makes free-then-reuse sequences safe: a
206        // PageFree followed by a later PageWrite for the same page must not
207        // leave the page on the free list.
208        let mut last_op: std::collections::HashMap<u64, &crate::wal::WalEntry> = Default::default();
209        let mut order: Vec<u64> = Vec::new();
210        for entry in &entries {
211            match entry.entry_type {
212                WalEntryType::PageWrite | WalEntryType::PageFree => {
213                    if !last_op.contains_key(&entry.page_id) {
214                        order.push(entry.page_id);
215                    }
216                    last_op.insert(entry.page_id, entry);
217                }
218                WalEntryType::Checkpoint | WalEntryType::Commit => {}
219            }
220        }
221
222        for page_id_raw in order {
223            let entry = last_op[&page_id_raw];
224            let page_id = PageId(entry.page_id);
225            match entry.entry_type {
226                WalEntryType::PageWrite => {
227                    while pm.page_count() <= entry.page_id {
228                        pm.allocate_page()?;
229                    }
230
231                    let mut page = pm.read_page(page_id)?;
232                    let copy_len = entry.data.len().min(PAGE_DATA_SIZE);
233                    page.data[..copy_len].copy_from_slice(&entry.data[..copy_len]);
234                    if copy_len < PAGE_DATA_SIZE {
235                        page.data[copy_len..].fill(0);
236                    }
237                    page.header.page_id = entry.page_id;
238                    page.header.lsn = entry.lsn;
239                    page.header.page_type = PageType::Data as u8;
240                    page.header.free_space = (PAGE_DATA_SIZE - copy_len) as u16;
241                    page.header.checksum = page.compute_checksum();
242
243                    pm.write_page(page_id, &page)?;
244                    count += 1;
245                }
246                WalEntryType::PageFree => {
247                    // Mark the page Free without touching the free list; the
248                    // list is rebuilt from page types after replay so it stays
249                    // consistent with WAL-derived page states.
250                    if entry.page_id < pm.page_count() {
251                        let mut page = Page::zeroed();
252                        page.header.page_id = entry.page_id;
253                        page.header.page_type = PageType::Free as u8;
254                        pm.write_page(page_id, &page)?;
255                        self.buffer_pool.invalidate(page_id);
256                        count += 1;
257                    }
258                }
259                WalEntryType::Checkpoint | WalEntryType::Commit => {}
260            }
261        }
262
263        if count > 0 {
264            // Header updates (free list head, page count) are not fsynced on
265            // every write, so after a crash the free list can disagree with the
266            // replayed page states. Rebuild it from page types: the WAL is the
267            // sole source of truth for which pages are Free vs Data. With a
268            // memory map the free pages are known, so only those are read
269            // (and verified); otherwise every page is scanned.
270            let rebuilt = match &anchor {
271                Some(a) => pm.rebuild_free_list_from(a.table.free_pages())?,
272                None => false,
273            };
274            if !rebuilt {
275                pm.rebuild_free_list()?;
276            }
277            pm.sync()?;
278            match anchor {
279                // The replayed pages are now durable, which makes this a
280                // checkpoint: anchor the rolled forward map on a fresh entry
281                // before truncating the entries it was derived from.
282                Some(anchor) => {
283                    self.anchor_checkpoint(wal, &pm, anchor.table)?;
284                }
285                None => {
286                    let next_lsn = wal.next_lsn();
287                    wal.truncate(next_lsn)?;
288                }
289            }
290            info!(count, "WAL recovery replayed entries");
291        } else if let Some(anchor) = anchor {
292            // Nothing to replay, so no entry after the anchor changed the map.
293            *self.map_anchor.lock() = Some(anchor);
294        }
295
296        Ok(count)
297    }
298
299    /// Append a checkpoint entry with a fresh token, persist `table` as the map
300    /// as of that entry, and truncate the WAL to it. The page file must already
301    /// be durable. When the snapshot cannot be written the anchor is dropped
302    /// (the next open scans) but the checkpoint still completes.
303    fn anchor_checkpoint(
304        &self,
305        wal: &mut Wal,
306        pm: &PageManager,
307        table: MemoryTable,
308    ) -> MenteResult<u64> {
309        let token = memory_map::new_token();
310        let lsn = wal.append(WalEntryType::Checkpoint, 0, &token)?;
311        wal.sync()?;
312        let anchor = MapAnchor {
313            token,
314            page_count: pm.page_count(),
315            table,
316            pending: Default::default(),
317        };
318        // The file is written before the slot is locked, so stamp lookups on
319        // the recall path never wait on snapshot I/O.
320        let written = memory_map::write_snapshot(&self.dir, &anchor);
321        match written {
322            Ok(()) => *self.map_anchor.lock() = Some(anchor),
323            Err(e) => {
324                tracing::warn!(error = %e, "memory map snapshot write failed, next open will scan pages");
325                memory_map::remove_snapshot(&self.dir);
326                *self.map_anchor.lock() = None;
327            }
328        }
329        wal.truncate(lsn)?;
330        Ok(lsn)
331    }
332
333    /// Gracefully shut down: flush dirty pages, sync files.
334    ///
335    /// # Example
336    ///
337    /// ```no_run
338    /// # use mentedb_storage::StorageEngine;
339    /// # let engine = StorageEngine::open("/tmp/mentedb".as_ref())?;
340    /// engine.close()?;
341    /// # Ok::<(), mentedb_core::error::MenteError>(())
342    /// ```
343    pub fn close(&self) -> MenteResult<()> {
344        let mut pm = self.page_manager.lock();
345        self.buffer_pool.flush_all(&mut pm)?;
346        pm.sync()?;
347        self.wal.lock().sync()?;
348        // Release the process lock so the directory can be reopened, by this
349        // process or another, without waiting for us to exit.
350        if let Some(lock_file) = self.process_lock.lock().take() {
351            let _ = crate::lock::unlock(&lock_file);
352        }
353        info!("storage engine closed");
354        Ok(())
355    }
356
357    /// Release the process lock without flushing anything.
358    ///
359    /// Test support for simulating a process crash: a real crash releases
360    /// the OS file lock but persists nothing beyond what was already synced.
361    #[doc(hidden)]
362    pub fn release_process_lock(&self) {
363        if let Some(lock_file) = self.process_lock.lock().take() {
364            let _ = crate::lock::unlock(&lock_file);
365        }
366    }
367
368    // ---- low-level page operations ----
369
370    /// Allocate a fresh page (for internal/test use).
371    ///
372    /// **WARNING**: In multi-process scenarios, prefer `store_memory` which
373    /// allocates under the WAL flock. This method does NOT acquire the flock.
374    pub fn allocate_page(&self) -> MenteResult<PageId> {
375        self.page_manager.lock().allocate_page()
376    }
377
378    /// Read a page through the buffer pool (lock-free — no WAL access).
379    pub fn read_page(&self, page_id: PageId) -> MenteResult<Box<Page>> {
380        self.buffer_pool
381            .fetch_page(page_id, &mut self.page_manager.lock())
382    }
383
384    /// Buffer-pool cache activity, for metrics (hit ratio, evictions, residency).
385    pub fn buffer_stats(&self) -> crate::buffer::BufferStats {
386        self.buffer_pool.stats()
387    }
388
389    /// Total pages in the store; multiply by `PAGE_SIZE` for on-disk bytes.
390    pub fn page_count(&self) -> u64 {
391        self.page_manager.lock().page_count()
392    }
393
394    /// Write data into an already-allocated page with WAL protection.
395    ///
396    /// Acquires the WAL flock for the duration of the write transaction.
397    /// For new pages, prefer `store_memory` which allocates + writes atomically.
398    pub fn write_page(&self, page_id: PageId, data: &[u8]) -> MenteResult<()> {
399        let id = memory_map::decode_memory(data).ok().map(|n| n.id);
400        self.write_page_as(page_id, data, id)
401    }
402
403    /// `write_page` for data known to decode as memory `id` (or as none).
404    fn write_page_as(
405        &self,
406        page_id: PageId,
407        data: &[u8],
408        id: Option<mentedb_core::types::MemoryId>,
409    ) -> MenteResult<()> {
410        let lsn = {
411            let mut wal = self.wal.lock();
412            wal.lock_exclusive()?;
413            wal.reload_lsn()?;
414            let lsn = wal.append(WalEntryType::PageWrite, page_id.0, data)?;
415            wal.sync()?;
416            self.note_map_write(page_id, id, lsn);
417            wal.unlock()?;
418            lsn
419        };
420
421        self.apply_page_write(page_id, data, lsn)
422    }
423
424    /// Apply a page write to the buffer pool and page manager (after WAL).
425    fn apply_page_write(&self, page_id: PageId, data: &[u8], lsn: u64) -> MenteResult<()> {
426        let mut pm = self.page_manager.lock();
427        let mut page = self.buffer_pool.fetch_page(page_id, &mut pm)?;
428        drop(pm);
429
430        let copy_len = data.len().min(PAGE_DATA_SIZE);
431        page.data[..copy_len].copy_from_slice(&data[..copy_len]);
432        if copy_len < PAGE_DATA_SIZE {
433            page.data[copy_len..].fill(0);
434        }
435        page.header.lsn = lsn;
436        page.header.page_type = PageType::Data as u8;
437        page.header.free_space = (PAGE_DATA_SIZE - copy_len) as u16;
438        page.header.checksum = page.compute_checksum();
439
440        if self.buffer_pool.update_page(page_id, &page).is_err() {
441            self.page_manager.lock().write_page(page_id, &page)?;
442        }
443        self.buffer_pool.unpin_page(page_id, true).ok();
444
445        Ok(())
446    }
447
448    // ---- high-level memory operations ----
449
450    /// Serialize and store a [`MemoryNode`] into a single page.
451    ///
452    /// The entire operation — page allocation, WAL append, page write — executes
453    /// under a single WAL flock, making it safe across multiple processes.
454    ///
455    /// # Example
456    ///
457    /// ```no_run
458    /// use mentedb_storage::StorageEngine;
459    /// use mentedb_core::{MemoryNode, memory::MemoryType, types::AgentId};
460    ///
461    /// let engine = StorageEngine::open("/tmp/mentedb".as_ref())?;
462    /// let node = MemoryNode::new(
463    ///     AgentId::new(),
464    ///     MemoryType::Semantic,
465    ///     "User likes dark mode".to_string(),
466    ///     vec![0.1, 0.2],
467    /// );
468    /// let page_id = engine.store_memory(&node)?;
469    /// # Ok::<(), mentedb_core::error::MenteError>(())
470    /// ```
471    pub fn store_memory(&self, node: &MemoryNode) -> MenteResult<PageId> {
472        let serialized =
473            serde_json::to_vec(node).map_err(|e| MenteError::Serialization(e.to_string()))?;
474
475        if serialized.len() + 4 > PAGE_DATA_SIZE {
476            return Err(MenteError::CapacityExceeded(format!(
477                "memory node serialized to {} bytes (max {})",
478                serialized.len(),
479                PAGE_DATA_SIZE - 4,
480            )));
481        }
482
483        let mut buf = Vec::with_capacity(4 + serialized.len());
484        buf.extend_from_slice(&(serialized.len() as u32).to_le_bytes());
485        buf.extend_from_slice(&serialized);
486
487        // Atomic write transaction: allocate + WAL + page write under one flock
488        let (page_id, lsn) = {
489            let mut wal = self.wal.lock();
490            let mut pm = self.page_manager.lock();
491
492            // Acquire flock and refresh state from disk
493            wal.lock_exclusive()?;
494            pm.reload_header()?;
495            wal.reload_lsn()?;
496
497            // Allocate page (using fresh page_count from disk)
498            let page_id = pm.allocate_page()?;
499
500            // WAL append + sync (WAL fsync guarantees durability;
501            // page data is written but not fsynced — checkpoint handles that)
502            let lsn = wal.append(WalEntryType::PageWrite, page_id.0, &buf)?;
503            wal.sync()?;
504            self.note_map_write(page_id, Some(node.id), lsn);
505
506            // Write page data to disk (no fsync — recoverable from WAL)
507            let mut page = Page::zeroed();
508            page.header.page_id = page_id.0;
509            let copy_len = buf.len().min(PAGE_DATA_SIZE);
510            page.data[..copy_len].copy_from_slice(&buf[..copy_len]);
511            page.header.lsn = lsn;
512            page.header.page_type = PageType::Data as u8;
513            page.header.free_space = (PAGE_DATA_SIZE - copy_len) as u16;
514            page.header.checksum = page.compute_checksum();
515            pm.write_page(page_id, &page)?;
516
517            // Release flock — other processes can now write
518            wal.unlock()?;
519
520            (page_id, lsn)
521        };
522
523        // Drop any stale cached copy (the page may have been reused from the
524        // free list); the next read loads the fresh data from disk.
525        let _ = lsn;
526        self.buffer_pool.invalidate(page_id);
527
528        // Auto-checkpoint when WAL exceeds threshold to prevent unbounded growth.
529        // This keeps reload_lsn() fast for subsequent writes.
530        if self.wal.lock().file_size() > WAL_AUTO_CHECKPOINT_BYTES
531            && let Err(e) = self.checkpoint()
532        {
533            tracing::warn!("auto-checkpoint failed: {e}");
534        }
535
536        info!(
537            page_id = page_id.0,
538            bytes = serialized.len(),
539            "stored memory node"
540        );
541        Ok(page_id)
542    }
543
544    /// Store multiple [`MemoryNode`]s in a single locked transaction.
545    ///
546    /// Acquires the WAL flock once, writes all nodes, then releases. This avoids
547    /// the per-write overhead of `reload_header` / `reload_lsn` for bulk inserts.
548    /// Auto-checkpoints after the batch if the WAL exceeds the threshold.
549    pub fn store_memory_batch(&self, nodes: &[MemoryNode]) -> MenteResult<Vec<PageId>> {
550        // Phase 1: serialize all nodes upfront (no locks held)
551        let mut bufs = Vec::with_capacity(nodes.len());
552        for node in nodes {
553            let serialized =
554                serde_json::to_vec(node).map_err(|e| MenteError::Serialization(e.to_string()))?;
555            if serialized.len() + 4 > PAGE_DATA_SIZE {
556                return Err(MenteError::CapacityExceeded(format!(
557                    "memory node serialized to {} bytes (max {})",
558                    serialized.len(),
559                    PAGE_DATA_SIZE - 4,
560                )));
561            }
562            let mut buf = Vec::with_capacity(4 + serialized.len());
563            buf.extend_from_slice(&(serialized.len() as u32).to_le_bytes());
564            buf.extend_from_slice(&serialized);
565            bufs.push(buf);
566        }
567
568        // Phase 2: single locked transaction for all writes
569        let page_ids = {
570            let mut wal = self.wal.lock();
571            let mut pm = self.page_manager.lock();
572
573            wal.lock_exclusive()?;
574            pm.reload_header()?;
575            wal.reload_lsn()?;
576
577            let mut ids = Vec::with_capacity(bufs.len());
578            for (buf, node) in bufs.iter().zip(nodes) {
579                let page_id = pm.allocate_page()?;
580                let lsn = wal.append(WalEntryType::PageWrite, page_id.0, buf)?;
581                self.note_map_write(page_id, Some(node.id), lsn);
582
583                let mut page = Page::zeroed();
584                page.header.page_id = page_id.0;
585                let copy_len = buf.len().min(PAGE_DATA_SIZE);
586                page.data[..copy_len].copy_from_slice(&buf[..copy_len]);
587                page.header.lsn = lsn;
588                page.header.page_type = PageType::Data as u8;
589                page.header.free_space = (PAGE_DATA_SIZE - copy_len) as u16;
590                page.header.checksum = page.compute_checksum();
591                pm.write_page(page_id, &page)?;
592
593                ids.push(page_id);
594            }
595
596            // WAL fsync only — page data is recoverable from WAL on crash.
597            // Checkpoint handles page file fsync.
598            wal.sync()?;
599            wal.unlock()?;
600
601            ids
602        };
603
604        // Drop stale cached copies for pages reused from the free list.
605        for page_id in &page_ids {
606            self.buffer_pool.invalidate(*page_id);
607        }
608
609        // Auto-checkpoint if WAL grew too large
610        if self.wal.lock().file_size() > WAL_AUTO_CHECKPOINT_BYTES
611            && let Err(e) = self.checkpoint()
612        {
613            tracing::warn!("auto-checkpoint failed: {e}");
614        }
615
616        info!(count = page_ids.len(), "stored memory batch");
617        Ok(page_ids)
618    }
619
620    /// Update a [`MemoryNode`] in place on its existing page.
621    ///
622    /// The write goes through the WAL-protected `write_page` path, so it is
623    /// crash-durable and keeps the buffer pool coherent. Unlike storing to a
624    /// fresh page, this never orphans the old copy.
625    pub fn update_memory(&self, page_id: PageId, node: &MemoryNode) -> MenteResult<()> {
626        let serialized =
627            serde_json::to_vec(node).map_err(|e| MenteError::Serialization(e.to_string()))?;
628
629        if serialized.len() + 4 > PAGE_DATA_SIZE {
630            return Err(MenteError::CapacityExceeded(format!(
631                "memory node serialized to {} bytes (max {})",
632                serialized.len(),
633                PAGE_DATA_SIZE - 4,
634            )));
635        }
636
637        let mut buf = Vec::with_capacity(4 + serialized.len());
638        buf.extend_from_slice(&(serialized.len() as u32).to_le_bytes());
639        buf.extend_from_slice(&serialized);
640
641        self.write_page_as(page_id, &buf, Some(node.id))
642    }
643
644    /// Delete a memory by returning its page to the free list.
645    ///
646    /// The deletion is WAL-logged before the page is freed, so it survives a
647    /// crash: recovery replays the `PageFree` entry and the memory does not
648    /// resurrect on reopen. The freed page is reused by later allocations.
649    pub fn delete_memory(&self, page_id: PageId) -> MenteResult<()> {
650        {
651            let mut wal = self.wal.lock();
652            let mut pm = self.page_manager.lock();
653
654            wal.lock_exclusive()?;
655            pm.reload_header()?;
656            wal.reload_lsn()?;
657
658            // WAL fsync guarantees the deletion is durable before the page
659            // is touched; the free-list update itself is recoverable.
660            let lsn = wal.append(WalEntryType::PageFree, page_id.0, &[])?;
661            wal.sync()?;
662            self.note_map_write(page_id, None, lsn);
663
664            pm.free_page(page_id)?;
665            wal.unlock()?;
666        }
667
668        // A stale cached copy must never be served or flushed back.
669        self.buffer_pool.invalidate(page_id);
670
671        info!(page_id = page_id.0, "deleted memory node");
672        Ok(())
673    }
674
675    /// Load and deserialize a [`MemoryNode`] from the given page.
676    ///
677    /// # Example
678    ///
679    /// ```no_run
680    /// # use mentedb_storage::{StorageEngine, PageId};
681    /// # let engine = StorageEngine::open("/tmp/mentedb".as_ref())?;
682    /// let node = engine.load_memory(PageId(1))?;
683    /// println!("memory: {}", node.content);
684    /// # Ok::<(), mentedb_core::error::MenteError>(())
685    /// ```
686    pub fn load_memory(&self, page_id: PageId) -> MenteResult<MemoryNode> {
687        self.memory_loads
688            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
689        let page = self.read_page(page_id)?;
690        self.buffer_pool.unpin_page(page_id, false).ok();
691
692        if PageType::from(page.header.page_type) != PageType::Data {
693            return Err(MenteError::Storage(format!(
694                "page {} is not a data page",
695                page_id.0
696            )));
697        }
698
699        memory_map::decode_memory(&page.data)
700    }
701
702    // ---- durability ----
703
704    /// Checkpoint: flush all dirty pages, sync to disk, and truncate the WAL.
705    ///
706    /// # Example
707    ///
708    /// ```no_run
709    /// # use mentedb_storage::StorageEngine;
710    /// # let engine = StorageEngine::open("/tmp/mentedb".as_ref())?;
711    /// // After a batch of writes, checkpoint to reclaim WAL space
712    /// engine.checkpoint()?;
713    /// # Ok::<(), mentedb_core::error::MenteError>(())
714    /// ```
715    pub fn checkpoint(&self) -> MenteResult<()> {
716        let mut wal = self.wal.lock();
717        let mut pm = self.page_manager.lock();
718
719        wal.lock_exclusive()?;
720        let result = self.checkpoint_locked(&mut wal, &mut pm, None);
721        let unlocked = wal.unlock();
722        let lsn = result?;
723        unlocked?;
724
725        info!(lsn, "checkpoint complete");
726        Ok(())
727    }
728
729    /// Checkpoint under the WAL flock. `fresh` is a memory table known to be
730    /// current (a full scan just built it); otherwise the anchored table is
731    /// rolled forward over the WAL. The snapshot is rewritten only when the
732    /// map changed; an unchanged map keeps its token, so the new checkpoint
733    /// entry still anchors the existing snapshot.
734    fn checkpoint_locked(
735        &self,
736        wal: &mut Wal,
737        pm: &mut PageManager,
738        fresh: Option<MemoryTable>,
739    ) -> MenteResult<u64> {
740        wal.reload_lsn()?;
741
742        self.buffer_pool.flush_all(pm)?;
743        pm.sync()?;
744
745        let table = match fresh {
746            Some(t) => Some(t),
747            None => {
748                let current = self.map_anchor.lock().take();
749                match current {
750                    Some(mut anchor) => {
751                        let entries = wal.iterate()?;
752                        match anchor.advance(&entries) {
753                            Some(true) => Some(anchor.table),
754                            Some(false) => {
755                                // Unchanged: re-anchor on the same token, the
756                                // snapshot on disk already holds this map.
757                                let lsn = wal.append(WalEntryType::Checkpoint, 0, &anchor.token)?;
758                                wal.sync()?;
759                                *self.map_anchor.lock() = Some(anchor);
760                                wal.truncate(lsn)?;
761                                return Ok(lsn);
762                            }
763                            None => {
764                                tracing::warn!(
765                                    "memory map anchor left the WAL, next open will scan pages"
766                                );
767                                None
768                            }
769                        }
770                    }
771                    None => None,
772                }
773            }
774        };
775
776        match table {
777            Some(table) => self.anchor_checkpoint(wal, pm, table),
778            None => {
779                // No map to persist. A fresh token keeps any old snapshot from
780                // ever matching this WAL.
781                let lsn = wal.append(WalEntryType::Checkpoint, 0, &memory_map::new_token())?;
782                wal.sync()?;
783                wal.truncate(lsn)?;
784                Ok(lsn)
785            }
786        }
787    }
788
789    /// Every memory in storage with its page and write stamp, in page order:
790    /// the same memories, in the same order, as [`Self::scan_all_memories`].
791    ///
792    /// Served from the persisted memory map rolled forward over the WAL, which
793    /// reads no pages. When there is no usable map this falls back to a full
794    /// page scan and then checkpoints, persisting the map so the next call
795    /// (and the next open) is fast. Holds the WAL lock throughout, so writers
796    /// wait; it is meant for open, before the database serves traffic.
797    pub fn memory_map(&self) -> MenteResult<Vec<MappedMemory>> {
798        let mut wal = self.wal.lock();
799        wal.lock_exclusive()?;
800        let result = self.memory_map_locked(&mut wal);
801        let unlocked = wal.unlock();
802        let map = result?;
803        unlocked?;
804        Ok(map)
805    }
806
807    fn memory_map_locked(&self, wal: &mut Wal) -> MenteResult<Vec<MappedMemory>> {
808        let entries = wal.iterate()?;
809        {
810            let mut slot = self.map_anchor.lock();
811            if let Some(anchor) = slot.as_ref() {
812                let mut current = anchor.clone();
813                if current.advance(&entries).is_some() {
814                    return Ok(current.table.memories());
815                }
816                *slot = None;
817            }
818        }
819
820        let mut pm = self.page_manager.lock();
821        let table = self.scan_table(&mut pm)?;
822        let memories = table.memories();
823        info!(
824            memories = memories.len(),
825            "memory map rebuilt by full page scan"
826        );
827        self.checkpoint_locked(wal, &mut pm, Some(table))?;
828        Ok(memories)
829    }
830
831    /// Build the memory table by reading every page.
832    fn scan_table(&self, pm: &mut PageManager) -> MenteResult<MemoryTable> {
833        pm.reload_header()?;
834        let count = pm.page_count();
835        let mut table = MemoryTable::default();
836        for i in 1..count {
837            let page_id = PageId(i);
838            let Ok(page) = self.buffer_pool.fetch_page(page_id, pm) else {
839                continue;
840            };
841            self.buffer_pool.unpin_page(page_id, false).ok();
842            match PageType::from(page.header.page_type) {
843                PageType::Data => {
844                    if let Ok(node) = memory_map::decode_memory(&page.data) {
845                        table.set_scanned(i, node.id);
846                    }
847                }
848                PageType::Free => table.set_free(i),
849                _ => {}
850            }
851        }
852        Ok(table)
853    }
854
855    /// The memory and write stamp the persisted map records for `page`, as of
856    /// the last checkpoint. None when the map is not anchored or the page holds
857    /// no memory.
858    pub fn memory_stamp(&self, page: PageId) -> Option<(mentedb_core::types::MemoryId, u64)> {
859        self.map_anchor.lock().as_ref()?.current(page.0)
860    }
861
862    /// Record a logged write in the anchored map's pending set. Called with
863    /// the WAL lock held, right after the append, so it is ordered with the
864    /// checkpoint that later rolls the same entry into the table.
865    fn note_map_write(&self, page: PageId, id: Option<mentedb_core::types::MemoryId>, lsn: u64) {
866        if let Some(anchor) = self.map_anchor.lock().as_mut() {
867            anchor.note_write(page.0, id, lsn);
868        }
869    }
870
871    /// Memory loads served since open, whether the page came from the buffer
872    /// pool or from disk. Counts what a caller asked for, so a test can bound
873    /// the work of an operation independent of cache state.
874    pub fn memory_loads(&self) -> u64 {
875        self.memory_loads.load(std::sync::atomic::Ordering::Relaxed)
876    }
877
878    /// Pages read from disk since open (buffer pool misses and direct reads).
879    pub fn page_reads(&self) -> u64 {
880        self.page_manager.lock().reads()
881    }
882
883    /// Scan all pages and return (MemoryId, PageId) pairs for every valid memory node.
884    ///
885    /// Refreshes the page count from disk before scanning so pages written by
886    /// other processes are included. Used to rebuild the page map on startup.
887    ///
888    /// # Example
889    ///
890    /// ```no_run
891    /// # use mentedb_storage::StorageEngine;
892    /// # let engine = StorageEngine::open("/tmp/mentedb".as_ref())?;
893    /// let memories = engine.scan_all_memories();
894    /// for (memory_id, page_id) in &memories {
895    ///     println!("{memory_id} -> page {}", page_id.0);
896    /// }
897    /// # Ok::<(), mentedb_core::error::MenteError>(())
898    /// ```
899    pub fn scan_all_memories(&self) -> Vec<(mentedb_core::types::MemoryId, PageId)> {
900        let mut pm = self.page_manager.lock();
901        // Refresh from disk to see pages written by other processes
902        let _ = pm.reload_header();
903        let count = pm.page_count();
904        drop(pm);
905
906        let mut results = Vec::new();
907        for i in 1..count {
908            let page_id = PageId(i);
909            if let Ok(node) = self.load_memory(page_id) {
910                results.push((node.id, page_id));
911            }
912        }
913        results
914    }
915}
916
917#[cfg(test)]
918mod tests {
919
920    #[test]
921    fn refuses_a_directory_migrated_to_the_newer_format() {
922        let dir = tempfile::tempdir().unwrap();
923        std::fs::write(dir.path().join(NEWER_FORMAT_FILE), b"record database").unwrap();
924        let err = match StorageEngine::open(dir.path()) {
925            Err(e) => e.to_string(),
926            Ok(_) => panic!("a migrated directory must not open as an empty page database"),
927        };
928        assert!(err.contains("newer storage format"), "{err}");
929        assert!(
930            !dir.path().join("pages.db").exists(),
931            "no empty page file may be created beside the record database"
932        );
933    }
934
935    use super::*;
936    use crate::memory_map::MappedMemory;
937    use mentedb_core::memory::MemoryType;
938    use mentedb_core::types::AgentId;
939
940    fn setup() -> (tempfile::TempDir, StorageEngine) {
941        let dir = tempfile::tempdir().unwrap();
942        let engine = StorageEngine::open(dir.path()).unwrap();
943        (dir, engine)
944    }
945
946    #[test]
947    fn test_allocate_write_read() {
948        let (_dir, engine) = setup();
949
950        let pid = engine.allocate_page().unwrap();
951        engine.write_page(pid, b"hello storage engine").unwrap();
952
953        let page = engine.read_page(pid).unwrap();
954        assert_eq!(&page.data[..20], b"hello storage engine");
955        engine.buffer_pool.unpin_page(pid, false).ok();
956    }
957
958    #[test]
959    fn test_store_and_load_memory() {
960        let (_dir, engine) = setup();
961
962        let node = MemoryNode::new(
963            AgentId::new(),
964            MemoryType::Episodic,
965            "The user prefers Rust over Go".to_string(),
966            vec![0.1, 0.2, 0.3, 0.4],
967        );
968
969        let page_id = engine.store_memory(&node).unwrap();
970        let loaded = engine.load_memory(page_id).unwrap();
971
972        assert_eq!(node.id, loaded.id);
973        assert_eq!(node.content, loaded.content);
974        assert_eq!(node.embedding, loaded.embedding);
975        assert_eq!(node.memory_type, loaded.memory_type);
976    }
977
978    #[test]
979    fn test_checkpoint() {
980        let (_dir, engine) = setup();
981
982        let node = MemoryNode::new(
983            AgentId::new(),
984            MemoryType::Semantic,
985            "checkpoint test".to_string(),
986            vec![1.0, 2.0],
987        );
988
989        let pid = engine.store_memory(&node).unwrap();
990        engine.checkpoint().unwrap();
991
992        let loaded = engine.load_memory(pid).unwrap();
993        assert_eq!(loaded.content, "checkpoint test");
994    }
995
996    #[test]
997    fn test_close_and_reopen() {
998        let dir = tempfile::tempdir().unwrap();
999        let pid;
1000        {
1001            let engine = StorageEngine::open(dir.path()).unwrap();
1002            let node = MemoryNode::new(
1003                AgentId::new(),
1004                MemoryType::Procedural,
1005                "persist across close".to_string(),
1006                vec![0.5],
1007            );
1008            pid = engine.store_memory(&node).unwrap();
1009            engine.close().unwrap();
1010        }
1011        {
1012            let engine = StorageEngine::open(dir.path()).unwrap();
1013            let loaded = engine.load_memory(pid).unwrap();
1014            assert_eq!(loaded.content, "persist across close");
1015        }
1016    }
1017
1018    #[test]
1019    fn test_crash_recovery() {
1020        let dir = tempfile::tempdir().unwrap();
1021        let mut ids = Vec::new();
1022        let mut contents = Vec::new();
1023        {
1024            let engine = StorageEngine::open(dir.path()).unwrap();
1025            for i in 0..3 {
1026                let content = format!("crash-recovery-{i}");
1027                let node = MemoryNode::new(
1028                    AgentId::new(),
1029                    MemoryType::Episodic,
1030                    content.clone(),
1031                    vec![i as f32],
1032                );
1033                let pid = engine.store_memory(&node).unwrap();
1034                ids.push(pid);
1035                contents.push(content);
1036            }
1037            // Simulate crash: sync the WAL but do NOT call close/checkpoint.
1038            engine.wal.lock().sync().unwrap();
1039        }
1040        {
1041            let engine = StorageEngine::open(dir.path()).unwrap();
1042            for (pid, expected) in ids.iter().zip(contents.iter()) {
1043                let loaded = engine.load_memory(*pid).unwrap();
1044                assert_eq!(&loaded.content, expected);
1045            }
1046        }
1047    }
1048
1049    #[test]
1050    fn test_recovery_idempotent() {
1051        let dir = tempfile::tempdir().unwrap();
1052        let pid;
1053        let content = "idempotent-check".to_string();
1054        {
1055            let engine = StorageEngine::open(dir.path()).unwrap();
1056            let node = MemoryNode::new(
1057                AgentId::new(),
1058                MemoryType::Semantic,
1059                content.clone(),
1060                vec![1.0, 2.0],
1061            );
1062            pid = engine.store_memory(&node).unwrap();
1063            engine.checkpoint().unwrap();
1064            engine.close().unwrap();
1065        }
1066        {
1067            let engine = StorageEngine::open(dir.path()).unwrap();
1068            let loaded = engine.load_memory(pid).unwrap();
1069            assert_eq!(loaded.content, content);
1070        }
1071    }
1072
1073    #[test]
1074    fn test_partial_write_recovery() {
1075        let dir = tempfile::tempdir().unwrap();
1076        let mut ids = Vec::new();
1077        let mut contents = Vec::new();
1078        {
1079            let engine = StorageEngine::open(dir.path()).unwrap();
1080            for i in 0..3 {
1081                let content = format!("checkpointed-{i}");
1082                let node = MemoryNode::new(
1083                    AgentId::new(),
1084                    MemoryType::Semantic,
1085                    content.clone(),
1086                    vec![i as f32],
1087                );
1088                let pid = engine.store_memory(&node).unwrap();
1089                ids.push(pid);
1090                contents.push(content);
1091            }
1092            engine.checkpoint().unwrap();
1093
1094            for i in 3..5 {
1095                let content = format!("unckeckpointed-{i}");
1096                let node = MemoryNode::new(
1097                    AgentId::new(),
1098                    MemoryType::Episodic,
1099                    content.clone(),
1100                    vec![i as f32],
1101                );
1102                let pid = engine.store_memory(&node).unwrap();
1103                ids.push(pid);
1104                contents.push(content);
1105            }
1106            // Simulate crash — sync WAL but don't close.
1107            engine.wal.lock().sync().unwrap();
1108        }
1109        {
1110            let engine = StorageEngine::open(dir.path()).unwrap();
1111            for (pid, expected) in ids.iter().zip(contents.iter()) {
1112                let loaded = engine.load_memory(*pid).unwrap();
1113                assert_eq!(&loaded.content, expected);
1114            }
1115        }
1116    }
1117
1118    #[test]
1119    fn test_delete_memory_durable() {
1120        let dir = tempfile::tempdir().unwrap();
1121        let pid;
1122        {
1123            let engine = StorageEngine::open(dir.path()).unwrap();
1124            let node = MemoryNode::new(
1125                AgentId::new(),
1126                MemoryType::Semantic,
1127                "to be deleted".to_string(),
1128                vec![1.0],
1129            );
1130            pid = engine.store_memory(&node).unwrap();
1131            engine.delete_memory(pid).unwrap();
1132            assert!(engine.load_memory(pid).is_err());
1133            assert!(engine.scan_all_memories().is_empty());
1134            engine.close().unwrap();
1135        }
1136        {
1137            let engine = StorageEngine::open(dir.path()).unwrap();
1138            assert!(
1139                engine.load_memory(pid).is_err(),
1140                "deleted memory must not resurrect on reopen"
1141            );
1142            assert!(engine.scan_all_memories().is_empty());
1143        }
1144    }
1145
1146    #[test]
1147    fn test_delete_survives_crash() {
1148        let dir = tempfile::tempdir().unwrap();
1149        let pid;
1150        {
1151            let engine = StorageEngine::open(dir.path()).unwrap();
1152            let node = MemoryNode::new(
1153                AgentId::new(),
1154                MemoryType::Semantic,
1155                "crash delete".to_string(),
1156                vec![1.0],
1157            );
1158            pid = engine.store_memory(&node).unwrap();
1159            engine.delete_memory(pid).unwrap();
1160            // Simulate crash: no close, no checkpoint.
1161        }
1162        {
1163            let engine = StorageEngine::open(dir.path()).unwrap();
1164            assert!(
1165                engine.load_memory(pid).is_err(),
1166                "deletion must survive a crash via WAL replay"
1167            );
1168            assert!(engine.scan_all_memories().is_empty());
1169        }
1170    }
1171
1172    #[test]
1173    fn test_deleted_page_reused() {
1174        let (_dir, engine) = setup();
1175
1176        let a = MemoryNode::new(AgentId::new(), MemoryType::Semantic, "a".into(), vec![1.0]);
1177        let pid_a = engine.store_memory(&a).unwrap();
1178        engine.delete_memory(pid_a).unwrap();
1179
1180        let b = MemoryNode::new(AgentId::new(), MemoryType::Semantic, "b".into(), vec![2.0]);
1181        let pid_b = engine.store_memory(&b).unwrap();
1182        assert_eq!(pid_a, pid_b, "freed page should be reused");
1183
1184        let loaded = engine.load_memory(pid_b).unwrap();
1185        assert_eq!(loaded.content, "b");
1186    }
1187
1188    #[test]
1189    fn test_delete_reuse_crash_recovery() {
1190        let dir = tempfile::tempdir().unwrap();
1191        let pid;
1192        let b_id;
1193        {
1194            let engine = StorageEngine::open(dir.path()).unwrap();
1195            let a = MemoryNode::new(AgentId::new(), MemoryType::Semantic, "a".into(), vec![1.0]);
1196            pid = engine.store_memory(&a).unwrap();
1197            engine.delete_memory(pid).unwrap();
1198            let b = MemoryNode::new(AgentId::new(), MemoryType::Semantic, "b".into(), vec![2.0]);
1199            let pid_b = engine.store_memory(&b).unwrap();
1200            assert_eq!(pid, pid_b);
1201            b_id = b.id;
1202            // Simulate crash: the WAL now holds PageFree(p) then PageWrite(p).
1203        }
1204        {
1205            let engine = StorageEngine::open(dir.path()).unwrap();
1206            let loaded = engine.load_memory(pid).unwrap();
1207            assert_eq!(loaded.content, "b", "later write must win over the free");
1208            assert_eq!(loaded.id, b_id);
1209            // The page must NOT be on the free list: a fresh allocation must
1210            // not clobber b.
1211            let c = MemoryNode::new(AgentId::new(), MemoryType::Semantic, "c".into(), vec![3.0]);
1212            let pid_c = engine.store_memory(&c).unwrap();
1213            assert_ne!(pid_c, pid, "recovered free list must exclude reused page");
1214            assert_eq!(engine.load_memory(pid).unwrap().content, "b");
1215        }
1216    }
1217
1218    #[test]
1219    fn test_update_memory_in_place() {
1220        let dir = tempfile::tempdir().unwrap();
1221        let pid;
1222        let id;
1223        {
1224            let engine = StorageEngine::open(dir.path()).unwrap();
1225            let mut node = MemoryNode::new(
1226                AgentId::new(),
1227                MemoryType::Semantic,
1228                "original".to_string(),
1229                vec![1.0],
1230            );
1231            pid = engine.store_memory(&node).unwrap();
1232            id = node.id;
1233
1234            node.content = "updated".to_string();
1235            engine.update_memory(pid, &node).unwrap();
1236
1237            let loaded = engine.load_memory(pid).unwrap();
1238            assert_eq!(loaded.content, "updated");
1239            // No orphan copy: exactly one entry in a full scan.
1240            let scanned = engine.scan_all_memories();
1241            assert_eq!(scanned.len(), 1);
1242            engine.close().unwrap();
1243        }
1244        {
1245            let engine = StorageEngine::open(dir.path()).unwrap();
1246            let loaded = engine.load_memory(pid).unwrap();
1247            assert_eq!(loaded.content, "updated");
1248            assert_eq!(loaded.id, id);
1249            assert_eq!(engine.scan_all_memories().len(), 1);
1250        }
1251    }
1252
1253    #[test]
1254    fn test_concurrent_open_is_rejected() {
1255        let dir = tempfile::tempdir().unwrap();
1256
1257        // The per operation flock serializes individual write transactions,
1258        // but layers above the storage engine cache state across operations,
1259        // so a second concurrent engine rolls the database back to its own
1260        // open time snapshot. The process lock must refuse it.
1261        let engine1 = StorageEngine::open(dir.path()).unwrap();
1262        let second = StorageEngine::open(dir.path());
1263        assert!(second.is_err(), "second concurrent open must fail");
1264        let msg = second.err().unwrap().to_string();
1265        assert!(msg.contains("locked"), "error names the lock: {msg}");
1266
1267        // Close releases the lock; the directory can be opened again.
1268        engine1.close().unwrap();
1269        let engine2 = StorageEngine::open(dir.path()).unwrap();
1270        engine2.close().unwrap();
1271    }
1272
1273    #[test]
1274    fn test_concurrent_writes_from_threads() {
1275        use std::sync::Arc;
1276        let dir = tempfile::tempdir().unwrap();
1277        let engine = Arc::new(StorageEngine::open(dir.path()).unwrap());
1278
1279        let handles: Vec<_> = (0..10)
1280            .map(|i| {
1281                let eng = Arc::clone(&engine);
1282                std::thread::spawn(move || {
1283                    let node = MemoryNode::new(
1284                        AgentId::new(),
1285                        MemoryType::Episodic,
1286                        format!("thread-{i}"),
1287                        vec![i as f32],
1288                    );
1289                    eng.store_memory(&node).unwrap()
1290                })
1291            })
1292            .collect();
1293
1294        let pids: Vec<PageId> = handles.into_iter().map(|h| h.join().unwrap()).collect();
1295
1296        // All 10 writes succeeded and are readable
1297        for (i, pid) in pids.iter().enumerate() {
1298            let loaded = engine.load_memory(*pid).unwrap();
1299            assert_eq!(loaded.content, format!("thread-{i}"));
1300        }
1301    }
1302
1303    // ---- persisted memory map ----
1304
1305    fn node(content: &str) -> MemoryNode {
1306        MemoryNode::new(
1307            AgentId::new(),
1308            MemoryType::Semantic,
1309            content.to_string(),
1310            vec![1.0, 2.0],
1311        )
1312    }
1313
1314    /// (memory id, page) pairs, the part of the map a full scan can confirm.
1315    fn ids(map: &[MappedMemory]) -> Vec<(mentedb_core::types::MemoryId, PageId)> {
1316        map.iter().map(|m| (m.id, m.page)).collect()
1317    }
1318
1319    /// Open a second, independent engine over the same files after `engine`
1320    /// crashed, and return what a full page scan sees there.
1321    fn scan_truth(dir: &Path) -> Vec<(mentedb_core::types::MemoryId, PageId)> {
1322        let engine = StorageEngine::open(dir).unwrap();
1323        let scanned = engine.scan_all_memories();
1324        engine.release_process_lock();
1325        scanned
1326    }
1327
1328    #[test]
1329    fn open_after_checkpoint_reads_no_pages() {
1330        let dir = tempfile::tempdir().unwrap();
1331        let n = 300;
1332        let mut stored = Vec::new();
1333        {
1334            let engine = StorageEngine::open(dir.path()).unwrap();
1335            for i in 0..n {
1336                let m = node(&format!("m{i}"));
1337                let pid = engine.store_memory(&m).unwrap();
1338                stored.push((m.id, pid));
1339            }
1340            // First map request scans once and persists the map.
1341            assert_eq!(ids(&engine.memory_map().unwrap()), stored);
1342            engine.checkpoint().unwrap();
1343            engine.close().unwrap();
1344        }
1345        let engine = StorageEngine::open(dir.path()).unwrap();
1346        let map = engine.memory_map().unwrap();
1347        assert_eq!(ids(&map), stored);
1348        assert_eq!(
1349            engine.page_reads(),
1350            0,
1351            "open must serve the map from the snapshot, not by reading {n} pages"
1352        );
1353    }
1354
1355    #[test]
1356    fn open_replays_only_wal_pages_after_crash() {
1357        let dir = tempfile::tempdir().unwrap();
1358        let mut stored = Vec::new();
1359        {
1360            let engine = StorageEngine::open(dir.path()).unwrap();
1361            for i in 0..200 {
1362                let m = node(&format!("base{i}"));
1363                stored.push((m.id, engine.store_memory(&m).unwrap()));
1364            }
1365            engine.memory_map().unwrap();
1366            engine.checkpoint().unwrap();
1367            for i in 0..5 {
1368                let m = node(&format!("after{i}"));
1369                stored.push((m.id, engine.store_memory(&m).unwrap()));
1370            }
1371            // Crash: nothing flushed or checkpointed beyond the WAL fsync.
1372            engine.release_process_lock();
1373        }
1374        let engine = StorageEngine::open(dir.path()).unwrap();
1375        let map = engine.memory_map().unwrap();
1376        assert_eq!(ids(&map), stored, "uncheckpointed stores survive the crash");
1377        assert_eq!(
1378            engine.page_reads(),
1379            5,
1380            "replay reads the 5 WAL pages only, not the 200 checkpointed ones"
1381        );
1382        engine.release_process_lock();
1383        assert_eq!(scan_truth(dir.path()), stored);
1384    }
1385
1386    #[test]
1387    fn deletes_and_updates_after_checkpoint_are_replayed() {
1388        let dir = tempfile::tempdir().unwrap();
1389        let (keep, gone, updated);
1390        let updated_pid;
1391        {
1392            let engine = StorageEngine::open(dir.path()).unwrap();
1393            let k = node("keep");
1394            let g = node("gone");
1395            let mut u = node("old");
1396            let kp = engine.store_memory(&k).unwrap();
1397            let gp = engine.store_memory(&g).unwrap();
1398            updated_pid = engine.store_memory(&u).unwrap();
1399            engine.memory_map().unwrap();
1400            engine.checkpoint().unwrap();
1401            engine.delete_memory(gp).unwrap();
1402            u.content = "new".into();
1403            engine.update_memory(updated_pid, &u).unwrap();
1404            keep = (k.id, kp);
1405            gone = g.id;
1406            updated = u.id;
1407            engine.release_process_lock();
1408        }
1409        let engine = StorageEngine::open(dir.path()).unwrap();
1410        let map = engine.memory_map().unwrap();
1411        assert!(
1412            map.iter().all(|m| m.id != gone),
1413            "deleted memory stays deleted"
1414        );
1415        assert!(ids(&map).contains(&keep));
1416        assert!(ids(&map).contains(&(updated, updated_pid)));
1417        assert_eq!(engine.load_memory(updated_pid).unwrap().content, "new");
1418        assert_eq!(map.len(), 2);
1419        engine.release_process_lock();
1420        assert_eq!(scan_truth(dir.path()), ids(&map));
1421    }
1422
1423    #[test]
1424    fn free_list_reuse_after_checkpoint_maps_the_new_memory() {
1425        let dir = tempfile::tempdir().unwrap();
1426        let (a_id, b_id, pid);
1427        {
1428            let engine = StorageEngine::open(dir.path()).unwrap();
1429            let a = node("a");
1430            pid = engine.store_memory(&a).unwrap();
1431            engine.store_memory(&node("other")).unwrap();
1432            engine.memory_map().unwrap();
1433            engine.checkpoint().unwrap();
1434            engine.delete_memory(pid).unwrap();
1435            let b = node("b");
1436            assert_eq!(engine.store_memory(&b).unwrap(), pid, "freed page reused");
1437            a_id = a.id;
1438            b_id = b.id;
1439            engine.release_process_lock();
1440        }
1441        let engine = StorageEngine::open(dir.path()).unwrap();
1442        let map = engine.memory_map().unwrap();
1443        assert!(
1444            map.iter().all(|m| m.id != a_id),
1445            "no stale id for the reused page"
1446        );
1447        assert!(ids(&map).contains(&(b_id, pid)));
1448        // The reused page is not on the rebuilt free list.
1449        let c = node("c");
1450        assert_ne!(engine.store_memory(&c).unwrap(), pid);
1451        assert_eq!(engine.load_memory(pid).unwrap().id, b_id);
1452        engine.release_process_lock();
1453        let truth = scan_truth(dir.path());
1454        assert!(truth.contains(&(b_id, pid)));
1455        assert!(truth.iter().all(|(id, _)| *id != a_id));
1456    }
1457
1458    #[test]
1459    fn free_list_rebuilt_from_map_after_crash() {
1460        let dir = tempfile::tempdir().unwrap();
1461        let mut freed = Vec::new();
1462        {
1463            let engine = StorageEngine::open(dir.path()).unwrap();
1464            let mut pids = Vec::new();
1465            for i in 0..50 {
1466                pids.push(engine.store_memory(&node(&format!("m{i}"))).unwrap());
1467            }
1468            engine.delete_memory(pids[3]).unwrap();
1469            freed.push(pids[3]);
1470            engine.memory_map().unwrap();
1471            engine.checkpoint().unwrap();
1472            engine.delete_memory(pids[7]).unwrap();
1473            freed.push(pids[7]);
1474            engine.release_process_lock();
1475        }
1476        let engine = StorageEngine::open(dir.path()).unwrap();
1477        // One replayed free plus verification reads of the two free pages,
1478        // not a scan of all 50.
1479        assert!(engine.page_reads() <= 3, "reads {}", engine.page_reads());
1480        let a = engine.store_memory(&node("x")).unwrap();
1481        let b = engine.store_memory(&node("y")).unwrap();
1482        let mut got = vec![a, b];
1483        got.sort_by_key(|p| p.0);
1484        assert_eq!(got, freed, "both freed pages are reusable after recovery");
1485        let c = engine.store_memory(&node("z")).unwrap();
1486        assert_eq!(c, PageId(51), "then allocation grows the file");
1487    }
1488
1489    #[test]
1490    fn missing_corrupt_or_stale_snapshot_falls_back_to_scan() {
1491        let dir = tempfile::tempdir().unwrap();
1492        let mut stored = Vec::new();
1493        {
1494            let engine = StorageEngine::open(dir.path()).unwrap();
1495            for i in 0..40 {
1496                let m = node(&format!("m{i}"));
1497                stored.push((m.id, engine.store_memory(&m).unwrap()));
1498            }
1499            engine.memory_map().unwrap();
1500            engine.checkpoint().unwrap();
1501            engine.close().unwrap();
1502        }
1503        let snap = dir.path().join(memory_map::MEMORY_MAP_FILE);
1504        let good = std::fs::read(&snap).unwrap();
1505
1506        let check = |label: &str| {
1507            let engine = StorageEngine::open(dir.path()).unwrap();
1508            let map = engine.memory_map().unwrap();
1509            assert_eq!(ids(&map), stored, "{label}: same result as a scan");
1510            assert!(
1511                engine.page_reads() >= 40,
1512                "{label}: must have scanned, read {}",
1513                engine.page_reads()
1514            );
1515            engine.close().unwrap();
1516            // The scan re-persisted a valid map.
1517            let engine = StorageEngine::open(dir.path()).unwrap();
1518            assert_eq!(ids(&engine.memory_map().unwrap()), stored);
1519            assert_eq!(engine.page_reads(), 0, "{label}: next open is fast again");
1520            engine.close().unwrap();
1521        };
1522
1523        std::fs::remove_file(&snap).unwrap();
1524        check("missing");
1525
1526        let mut bad = std::fs::read(&snap).unwrap();
1527        let mid = bad.len() / 2;
1528        bad[mid] ^= 0x55;
1529        std::fs::write(&snap, &bad).unwrap();
1530        check("corrupt");
1531
1532        // Stale: a snapshot from an earlier checkpoint whose token has left
1533        // the WAL, even though the file itself is intact.
1534        std::fs::write(&snap, &good).unwrap();
1535        check("stale");
1536    }
1537
1538    #[test]
1539    fn stale_snapshot_with_unchanged_token_reuse_is_never_matched() {
1540        // A checkpoint whose snapshot write is skipped (no map) must not let an
1541        // older snapshot match afterwards: deletes after it would resurrect.
1542        let dir = tempfile::tempdir().unwrap();
1543        let (a, b);
1544        {
1545            let engine = StorageEngine::open(dir.path()).unwrap();
1546            let na = node("a");
1547            let nb = node("b");
1548            a = engine.store_memory(&na).unwrap();
1549            b = (nb.id, engine.store_memory(&nb).unwrap());
1550            engine.memory_map().unwrap();
1551            engine.checkpoint().unwrap();
1552            // Drop the anchor as if the map were never loaded, then change the
1553            // map and checkpoint without persisting it.
1554            *engine.map_anchor.lock() = None;
1555            engine.delete_memory(a).unwrap();
1556            engine.checkpoint().unwrap();
1557            engine.close().unwrap();
1558        }
1559        let engine = StorageEngine::open(dir.path()).unwrap();
1560        assert_eq!(ids(&engine.memory_map().unwrap()), vec![b]);
1561    }
1562
1563    #[test]
1564    fn handoff_between_handles_sees_writes_and_checkpoints() {
1565        let dir = tempfile::tempdir().unwrap();
1566        let mut stored = Vec::new();
1567        // Handle 1 builds the map and exits cleanly.
1568        {
1569            let e1 = StorageEngine::open(dir.path()).unwrap();
1570            for i in 0..20 {
1571                let m = node(&format!("h1-{i}"));
1572                stored.push((m.id, e1.store_memory(&m).unwrap()));
1573            }
1574            e1.memory_map().unwrap();
1575            e1.checkpoint().unwrap();
1576            e1.close().unwrap();
1577        }
1578        // Handle 2 (another process taking over) writes, checkpoints, writes
1579        // more, then dies.
1580        {
1581            let e2 = StorageEngine::open(dir.path()).unwrap();
1582            e2.memory_map().unwrap();
1583            for i in 0..10 {
1584                let m = node(&format!("h2-{i}"));
1585                stored.push((m.id, e2.store_memory(&m).unwrap()));
1586            }
1587            e2.checkpoint().unwrap();
1588            e2.delete_memory(stored[0].1).unwrap();
1589            stored.remove(0);
1590            let m = node("h2-last");
1591            stored.push((m.id, e2.store_memory(&m).unwrap()));
1592            e2.release_process_lock();
1593        }
1594        // Handle 3 reopens and sees all of it.
1595        let e3 = StorageEngine::open(dir.path()).unwrap();
1596        let mut got = ids(&e3.memory_map().unwrap());
1597        got.sort_by_key(|(_, p)| p.0);
1598        stored.sort_by_key(|(_, p)| p.0);
1599        assert_eq!(got, stored);
1600        assert!(e3.page_reads() < 10, "reads {}", e3.page_reads());
1601        e3.release_process_lock();
1602        assert_eq!(scan_truth(dir.path()), stored);
1603    }
1604
1605    #[test]
1606    fn live_stamps_survive_checkpoint_and_crash_replay() {
1607        let dir = tempfile::tempdir().unwrap();
1608        let (b_page, b_live);
1609        {
1610            let engine = StorageEngine::open(dir.path()).unwrap();
1611            engine.store_memory(&node("a")).unwrap();
1612            engine.memory_map().unwrap();
1613            engine.checkpoint().unwrap();
1614            // A write after the checkpoint has its stamp before any checkpoint.
1615            let c = engine.store_memory(&node("c")).unwrap();
1616            let c_live = engine.memory_stamp(c).unwrap();
1617            engine.checkpoint().unwrap();
1618            assert_eq!(
1619                engine.memory_stamp(c).unwrap(),
1620                c_live,
1621                "checkpoint keeps it"
1622            );
1623            b_page = engine.store_memory(&node("b")).unwrap();
1624            b_live = engine.memory_stamp(b_page).unwrap();
1625            engine.release_process_lock();
1626        }
1627        let engine = StorageEngine::open(dir.path()).unwrap();
1628        let map = engine.memory_map().unwrap();
1629        let b = map.iter().find(|m| m.page == b_page).unwrap();
1630        assert_eq!(
1631            (b.id, b.stamp),
1632            b_live,
1633            "WAL replay reproduces the live stamp"
1634        );
1635    }
1636
1637    #[test]
1638    fn unchanged_checkpoints_keep_the_snapshot_and_stamps() {
1639        let dir = tempfile::tempdir().unwrap();
1640        let engine = StorageEngine::open(dir.path()).unwrap();
1641        let pid = engine.store_memory(&node("a")).unwrap();
1642        engine.memory_map().unwrap();
1643        engine.checkpoint().unwrap();
1644        let snap = dir.path().join(memory_map::MEMORY_MAP_FILE);
1645        let before = std::fs::metadata(&snap).unwrap().modified().unwrap();
1646        let stamp = engine.memory_stamp(pid).unwrap();
1647        std::thread::sleep(std::time::Duration::from_millis(20));
1648        engine.checkpoint().unwrap();
1649        engine.checkpoint().unwrap();
1650        assert_eq!(
1651            std::fs::metadata(&snap).unwrap().modified().unwrap(),
1652            before
1653        );
1654        assert_eq!(engine.memory_stamp(pid).unwrap(), stamp);
1655        // A write to the page changes its stamp at the next checkpoint.
1656        let mut n = engine.load_memory(pid).unwrap();
1657        n.content = "b".into();
1658        engine.update_memory(pid, &n).unwrap();
1659        engine.checkpoint().unwrap();
1660        assert_ne!(engine.memory_stamp(pid).unwrap().1, stamp.1);
1661        engine.close().unwrap();
1662        let engine = StorageEngine::open(dir.path()).unwrap();
1663        let map = engine.memory_map().unwrap();
1664        assert_eq!(map[0].stamp, engine.memory_stamp(pid).unwrap().1);
1665    }
1666}