1use std::fs::File;
5use std::io::{BufReader, Read, Seek, SeekFrom};
6use std::path::{Path, PathBuf};
7use std::sync::atomic::{AtomicU64, Ordering};
8
9use mentedb_core::MemoryNode;
10use mentedb_core::config::StorageConfig;
11use mentedb_core::error::{MenteError, MenteResult};
12use mentedb_core::types::MemoryId;
13
14use parking_lot::{Mutex, RwLock};
15use tracing::{info, warn};
16
17use crate::cache::{BufferStats, RecordCache};
18use crate::legacy::{self, MigrationFault};
19use crate::memory_map::{
20 self, MappedMemory, RecordEntry, RecordTable, Token, entry_stamp, new_token,
21};
22use crate::record::{
23 self, BLOCK_SIZE, DATA_FILE, EXTENT_HEADER_SIZE, ExtentHeader, ExtentKind, FreeSpace, PageId,
24 RecordOp, WAL_FILE,
25};
26use crate::wal::{Wal, WalEntry, WalEntryType};
27
28struct State {
33 wal: Wal,
34 table: RecordTable,
35 free: FreeSpace,
36 end: u64,
38 token: Token,
40 anchored: bool,
42 dirty: bool,
44}
45
46pub struct StorageEngine {
65 dir: PathBuf,
66 file: File,
67 state: RwLock<State>,
68 cache: RecordCache,
69 process_lock: Mutex<Option<File>>,
72 wal_checkpoint_bytes: u64,
73 max_record_bytes: u64,
74 memory_loads: AtomicU64,
76 disk_reads: AtomicU64,
78 bytes_read: AtomicU64,
80}
81
82fn lock_dir(path: &Path) -> MenteResult<File> {
83 let lock_file = std::fs::OpenOptions::new()
84 .create(true)
85 .truncate(false)
86 .write(true)
87 .open(path.join("LOCK"))?;
88 match crate::lock::try_lock_exclusive(&lock_file) {
89 Ok(true) => Ok(lock_file),
90 Ok(false) => Err(MenteError::Storage(format!(
91 "database directory {} is locked by another process",
92 path.display()
93 ))),
94 Err(e) => Err(MenteError::Storage(format!(
95 "failed to lock database directory {}: {e}",
96 path.display()
97 ))),
98 }
99}
100
101fn create_data_file(dir: &Path) -> MenteResult<()> {
103 for stale in [WAL_FILE, memory_map::MEMORY_MAP_FILE] {
104 let p = dir.join(stale);
105 if p.exists() {
106 std::fs::remove_file(&p)?;
107 }
108 }
109 let tmp = dir.join(format!("{DATA_FILE}.tmp"));
110 {
111 let f = File::create(&tmp)?;
112 record::write_at(&f, &record::file_header(), 0)?;
113 f.sync_all()?;
114 }
115 std::fs::rename(&tmp, dir.join(DATA_FILE))?;
116 record::sync_dir(dir)?;
117 info!(dir = %dir.display(), "created record file");
118 Ok(())
119}
120
121impl StorageEngine {
122 pub fn open(path: &Path) -> MenteResult<Self> {
133 Self::open_with_config(path, &StorageConfig::default())
134 }
135
136 pub fn open_with_config(path: &Path, config: &StorageConfig) -> MenteResult<Self> {
143 Self::open_inner(path, config, None)
144 }
145
146 #[doc(hidden)]
149 pub fn open_with_migration_fault(
150 path: &Path,
151 config: &StorageConfig,
152 fault: MigrationFault,
153 ) -> MenteResult<Self> {
154 Self::open_inner(path, config, Some(fault))
155 }
156
157 fn open_inner(
158 path: &Path,
159 config: &StorageConfig,
160 fault: Option<MigrationFault>,
161 ) -> MenteResult<Self> {
162 std::fs::create_dir_all(path)?;
163 let lock_file = lock_dir(path)?;
164
165 if path.join(DATA_FILE).exists() {
166 legacy::remove_leftovers(path)?;
167 } else if legacy::is_legacy(path) {
168 legacy::migrate(path, fault)?;
169 } else {
170 create_data_file(path)?;
171 }
172
173 let file = std::fs::OpenOptions::new()
174 .read(true)
175 .write(true)
176 .open(path.join(DATA_FILE))?;
177 let mut header = [0u8; BLOCK_SIZE];
178 record::read_at(&file, &mut header, 0)?;
179 record::check_file_header(&header)?;
180 let wal = Wal::open_file(&path.join(WAL_FILE))?;
181
182 let engine = Self {
183 dir: path.to_path_buf(),
184 file,
185 state: RwLock::new(State {
186 wal,
187 table: RecordTable::default(),
188 free: FreeSpace::default(),
189 end: 1,
190 token: new_token(),
191 anchored: false,
192 dirty: false,
193 }),
194 cache: RecordCache::new(config.cache_bytes),
195 process_lock: Mutex::new(Some(lock_file)),
196 wal_checkpoint_bytes: config.wal_checkpoint_bytes,
197 max_record_bytes: config.max_record_bytes,
198 memory_loads: AtomicU64::new(0),
199 disk_reads: AtomicU64::new(0),
200 bytes_read: AtomicU64::new(0),
201 };
202
203 let replayed = engine.recover()?;
204 let st = engine.state.read();
205 info!(
206 ?path,
207 replayed,
208 records = st.table.records.len(),
209 bytes = st.end * BLOCK_SIZE as u64,
210 "storage engine opened"
211 );
212 drop(st);
213 Ok(engine)
214 }
215
216 fn recover(&self) -> MenteResult<usize> {
220 let mut st = self.state.write();
221 let entries = st.wal.iterate()?;
222
223 let mut ops: Vec<(u64, RecordOp)> = Vec::new();
227 for e in &entries {
228 let op = match e.entry_type {
229 WalEntryType::RecordWrite => RecordOp::decode(true, e.page_id, &e.data)?,
230 WalEntryType::RecordFree => RecordOp::decode(false, e.page_id, &e.data)?,
231 WalEntryType::Checkpoint | WalEntryType::Commit => continue,
232 WalEntryType::PageWrite | WalEntryType::PageFree => {
233 return Err(MenteError::Storage(
234 "record WAL holds a page file entry".into(),
235 ));
236 }
237 };
238 op.redo(&self.file)?;
239 ops.push((e.lsn, op));
240 }
241 if !ops.is_empty() {
242 self.file.sync_data()?;
243 }
244 let file_end = self
245 .file
246 .metadata()?
247 .len()
248 .div_ceil(BLOCK_SIZE as u64)
249 .max(1);
250
251 let (table, token, anchored) = match self.load_snapshot(&entries, file_end) {
252 Some((table, token)) => (table, token, true),
253 None => (self.scan_chain(file_end)?, new_token(), false),
254 };
255 st.end = file_end.max(table.max_end());
256 st.free = table.free_space(st.end)?;
257 st.table = table;
258 st.token = token;
259 st.anchored = anchored;
260 st.dirty = !ops.is_empty();
261
262 if !anchored || !ops.is_empty() || st.wal.has_torn_tail() {
263 self.checkpoint_locked(&mut st, true)?;
264 }
265 if !ops.is_empty() {
266 info!(count = ops.len(), "WAL recovery replayed entries");
267 }
268 Ok(ops.len())
269 }
270
271 fn load_snapshot(&self, entries: &[WalEntry], file_end: u64) -> Option<(RecordTable, Token)> {
275 let snap = match memory_map::read_snapshot(&self.dir) {
276 Ok(Some(s)) => s,
277 Ok(None) => return None,
278 Err(e) => {
279 warn!(error = %e, "record map snapshot unreadable, open will scan the record file");
280 return None;
281 }
282 };
283 let Some(at) = entries
284 .iter()
285 .rposition(|e| e.entry_type == WalEntryType::Checkpoint && e.data == snap.token)
286 else {
287 info!("record map snapshot predates the WAL, open will scan the record file");
288 return None;
289 };
290 let mut table = snap.table;
291 for e in &entries[at + 1..] {
292 let write = match e.entry_type {
293 WalEntryType::RecordWrite => true,
294 WalEntryType::RecordFree => false,
295 _ => continue,
296 };
297 let op = RecordOp::decode(write, e.page_id, &e.data).ok()?;
298 table.apply(&op, entry_stamp(&snap.token, e.lsn));
299 }
300 if let Err(e) = table.free_space(file_end) {
301 warn!(error = %e, "record map snapshot does not fit the record file, open will scan");
302 return None;
303 }
304 Some((table, snap.token))
305 }
306
307 fn scan_chain(&self, file_end: u64) -> MenteResult<RecordTable> {
311 let mut table = RecordTable::default();
312 let mut reader = BufReader::with_capacity(1 << 20, self.file.try_clone()?);
313 reader.seek(SeekFrom::Start(BLOCK_SIZE as u64))?;
314 let mut at = 1u64;
315 let mut header = [0u8; EXTENT_HEADER_SIZE];
316 while at < file_end {
317 reader.read_exact(&mut header)?;
318 self.disk_reads.fetch_add(1, Ordering::Relaxed);
319 let h = ExtentHeader::decode(&header)
320 .ok_or_else(|| MenteError::Storage(format!("record chain broken at block {at}")))?;
321 let extent = record::Extent {
322 start: at,
323 blocks: h.blocks as u64,
324 };
325 if h.kind == ExtentKind::Live {
326 table.records.insert(
327 h.locator,
328 RecordEntry {
329 id: h.id,
330 stamp: memory_map::random_u64(),
331 extent,
332 },
333 );
334 table.next_locator = table.next_locator.max(h.locator + 1);
335 }
336 at = extent.end();
337 reader.seek_relative((extent.byte_len() - EXTENT_HEADER_SIZE) as i64)?;
338 }
339 self.bytes_read
340 .fetch_add(file_end * BLOCK_SIZE as u64, Ordering::Relaxed);
341 table.next_locator = table.next_locator.max(1);
342 info!(
343 records = table.records.len(),
344 "record map rebuilt by scanning the record file"
345 );
346 Ok(table)
347 }
348
349 pub fn close(&self) -> MenteResult<()> {
352 let mut st = self.state.write();
353 self.file.sync_data()?;
354 st.wal.sync()?;
355 drop(st);
356 if let Some(lock_file) = self.process_lock.lock().take() {
357 let _ = crate::lock::unlock(&lock_file);
358 }
359 info!("storage engine closed");
360 Ok(())
361 }
362
363 #[doc(hidden)]
368 pub fn release_process_lock(&self) {
369 if let Some(lock_file) = self.process_lock.lock().take() {
370 let _ = crate::lock::unlock(&lock_file);
371 }
372 }
373
374 pub fn buffer_stats(&self) -> BufferStats {
378 self.cache.stats()
379 }
380
381 pub fn set_cache_capacity(&self, bytes: u64) {
384 self.cache.set_capacity(bytes);
385 }
386
387 pub fn cache_capacity(&self) -> u64 {
389 self.cache.capacity()
390 }
391
392 pub fn page_count(&self) -> u64 {
394 self.state.read().end
395 }
396
397 pub fn data_bytes(&self) -> u64 {
399 self.state.read().end * BLOCK_SIZE as u64
400 }
401
402 pub fn free_bytes(&self) -> u64 {
404 self.state.read().free.free_blocks() * BLOCK_SIZE as u64
405 }
406
407 pub fn memory_loads(&self) -> u64 {
409 self.memory_loads.load(Ordering::Relaxed)
410 }
411
412 pub fn page_reads(&self) -> u64 {
415 self.disk_reads.load(Ordering::Relaxed)
416 }
417
418 pub fn bytes_read(&self) -> u64 {
420 self.bytes_read.load(Ordering::Relaxed)
421 }
422
423 fn encode(&self, node: &MemoryNode) -> MenteResult<Vec<u8>> {
426 let payload = record::encode_node(node)?;
427 if payload.len() as u64 > self.max_record_bytes {
428 return Err(MenteError::CapacityExceeded(format!(
429 "memory node encoded to {} bytes (max {})",
430 payload.len(),
431 self.max_record_bytes
432 )));
433 }
434 Ok(payload)
435 }
436
437 fn maybe_checkpoint(&self, wal_bytes: u64) {
439 if wal_bytes > self.wal_checkpoint_bytes
440 && let Err(e) = self.checkpoint()
441 {
442 warn!("auto-checkpoint failed: {e}");
443 }
444 }
445
446 pub fn store_memory(&self, node: &MemoryNode) -> MenteResult<PageId> {
463 let ids = self.store_memory_batch(std::slice::from_ref(node))?;
464 Ok(ids[0])
465 }
466
467 pub fn store_memory_batch(&self, nodes: &[MemoryNode]) -> MenteResult<Vec<PageId>> {
469 let payloads = nodes
470 .iter()
471 .map(|n| self.encode(n))
472 .collect::<MenteResult<Vec<_>>>()?;
473
474 let mut st = self.state.write();
475 let mut free = st.free.clone();
477 let mut end = st.end;
478 let first = st.table.next_locator;
479 let mut ops = Vec::with_capacity(nodes.len());
480 for (i, (node, payload)) in nodes.iter().zip(payloads).enumerate() {
481 let place = free.place(record::blocks_for(payload.len()), end);
482 end = free.commit(&place, end);
483 ops.push(RecordOp::Write {
484 locator: first + i as u64,
485 id: node.id,
486 extent: place.extent,
487 payload,
488 remainder: place.remainder,
489 freed: None,
490 });
491 }
492 let mut lsns = Vec::with_capacity(ops.len());
493 for op in &ops {
494 lsns.push(
495 st.wal
496 .append(WalEntryType::RecordWrite, op.locator(), &op.encode())?,
497 );
498 }
499 st.wal.sync()?;
500 st.free = free;
501 st.end = end;
502
503 let mut ids = Vec::with_capacity(ops.len());
504 for ((op, lsn), node) in ops.iter().zip(lsns).zip(nodes) {
505 op.redo(&self.file)?;
506 let stamp = entry_stamp(&st.token, lsn);
507 st.table.apply(op, stamp);
508 if let RecordOp::Write { payload, .. } = op {
509 self.cache
510 .insert(op.locator(), node.clone(), payload.len() as u64);
511 }
512 ids.push(PageId(op.locator()));
513 }
514 st.dirty = true;
515 let wal_bytes = st.wal.file_size();
516 drop(st);
517
518 self.maybe_checkpoint(wal_bytes);
519 Ok(ids)
520 }
521
522 pub fn update_memory(&self, page_id: PageId, node: &MemoryNode) -> MenteResult<()> {
527 let payload = self.encode(node)?;
528 let need = record::blocks_for(payload.len());
529 let len = payload.len() as u64;
530
531 let mut st = self.state.write();
532 let entry =
533 *st.table.records.get(&page_id.0).ok_or_else(|| {
534 MenteError::Storage(format!("page {} holds no memory", page_id.0))
535 })?;
536 let place = (need > entry.extent.blocks).then(|| st.free.place(need, st.end));
537 let op = RecordOp::Write {
538 locator: page_id.0,
539 id: node.id,
540 extent: place.map_or(entry.extent, |p| p.extent),
541 payload,
542 remainder: place.and_then(|p| p.remainder),
543 freed: place.map(|_| entry.extent),
544 };
545 let lsn = st
546 .wal
547 .append(WalEntryType::RecordWrite, page_id.0, &op.encode())?;
548 st.wal.sync()?;
549 if let Some(p) = place {
550 let end = st.end;
551 st.end = st.free.commit(&p, end);
552 st.free.insert(entry.extent);
553 }
554 op.redo(&self.file)?;
555 let stamp = entry_stamp(&st.token, lsn);
556 st.table.apply(&op, stamp);
557 st.dirty = true;
558 self.cache.insert(page_id.0, node.clone(), len);
559 let wal_bytes = st.wal.file_size();
560 drop(st);
561
562 self.maybe_checkpoint(wal_bytes);
563 Ok(())
564 }
565
566 pub fn delete_memory(&self, page_id: PageId) -> MenteResult<()> {
569 let mut st = self.state.write();
570 let entry =
571 *st.table.records.get(&page_id.0).ok_or_else(|| {
572 MenteError::Storage(format!("page {} holds no memory", page_id.0))
573 })?;
574 let op = RecordOp::Free {
575 locator: page_id.0,
576 extent: entry.extent,
577 };
578 let _lsn = st
579 .wal
580 .append(WalEntryType::RecordFree, page_id.0, &op.encode())?;
581 st.wal.sync()?;
582 st.free.insert(entry.extent);
583 op.redo(&self.file)?;
584 st.table.apply(&op, 0);
585 st.dirty = true;
586 self.cache.invalidate(page_id.0);
587 let wal_bytes = st.wal.file_size();
588 drop(st);
589
590 self.maybe_checkpoint(wal_bytes);
591 Ok(())
592 }
593
594 pub fn load_memory(&self, page_id: PageId) -> MenteResult<MemoryNode> {
604 self.memory_loads.fetch_add(1, Ordering::Relaxed);
605 if let Some(node) = self.cache.get(page_id.0) {
606 return Ok(node);
607 }
608 let st = self.state.read();
611 let entry =
612 st.table.records.get(&page_id.0).ok_or_else(|| {
613 MenteError::Storage(format!("page {} holds no memory", page_id.0))
614 })?;
615 let mut buf = vec![0u8; entry.extent.byte_len()];
616 record::read_at(&self.file, &mut buf, entry.extent.offset())?;
617 self.disk_reads.fetch_add(1, Ordering::Relaxed);
618 self.bytes_read
619 .fetch_add(buf.len() as u64, Ordering::Relaxed);
620 let payload = record::live_payload(&buf, page_id.0)?;
621 let node = record::decode_node(payload)?;
622 self.cache
623 .insert(page_id.0, node.clone(), payload.len() as u64);
624 Ok(node)
625 }
626
627 pub fn checkpoint(&self) -> MenteResult<()> {
632 let mut st = self.state.write();
633 self.checkpoint_locked(&mut st, false)
634 }
635
636 fn checkpoint_locked(&self, st: &mut State, force: bool) -> MenteResult<()> {
637 if !force && !st.dirty && st.anchored {
638 return Ok(());
639 }
640 self.file.sync_data()?;
642 let token = new_token();
643 let lsn = st.wal.append(WalEntryType::Checkpoint, 0, &token)?;
644 st.wal.sync()?;
645 match memory_map::write_snapshot(&self.dir, &token, &st.table) {
649 Ok(()) => st.anchored = true,
650 Err(e) => {
651 warn!(error = %e, "record map snapshot write failed, next open will scan");
652 memory_map::remove_snapshot(&self.dir);
653 st.anchored = false;
654 }
655 }
656 st.token = token;
657 st.dirty = false;
658 st.wal.reset_to(lsn, WalEntryType::Checkpoint, 0, &token)?;
659 info!(lsn, records = st.table.records.len(), "checkpoint complete");
660 Ok(())
661 }
662
663 pub fn memory_map(&self) -> MenteResult<Vec<MappedMemory>> {
666 Ok(self.state.read().table.memories())
667 }
668
669 pub fn memory_stamp(&self, page: PageId) -> Option<(MemoryId, u64)> {
671 self.state
672 .read()
673 .table
674 .records
675 .get(&page.0)
676 .map(|e| (e.id, e.stamp))
677 }
678
679 pub fn scan_all_memories(&self) -> Vec<(MemoryId, PageId)> {
683 let st = self.state.read();
684 let end = st.end;
685 let Ok(file) = self.file.try_clone() else {
686 return Vec::new();
687 };
688 let mut reader = BufReader::with_capacity(1 << 20, file);
689 if reader.seek(SeekFrom::Start(BLOCK_SIZE as u64)).is_err() {
690 return Vec::new();
691 }
692 let mut out = Vec::new();
693 let mut at = 1u64;
694 let mut header = [0u8; EXTENT_HEADER_SIZE];
695 while at < end {
696 if reader.read_exact(&mut header).is_err() {
697 break;
698 }
699 self.disk_reads.fetch_add(1, Ordering::Relaxed);
700 let Some(h) = ExtentHeader::decode(&header) else {
701 break;
702 };
703 let mut rest = vec![0u8; h.blocks as usize * BLOCK_SIZE - EXTENT_HEADER_SIZE];
704 if reader.read_exact(&mut rest).is_err() {
705 break;
706 }
707 if h.kind == ExtentKind::Live {
708 let mut buf = header.to_vec();
709 buf.extend_from_slice(&rest);
710 if let Ok(node) =
711 record::live_payload(&buf, h.locator).and_then(record::decode_node)
712 {
713 out.push((node.id, PageId(h.locator)));
714 }
715 }
716 at += h.blocks as u64;
717 }
718 drop(st);
719 out.sort_by_key(|(_, p)| *p);
720 out
721 }
722}
723
724#[cfg(test)]
725mod tests;