reddb_server/storage/wal/
checkpoint.rs1use std::collections::{HashMap, HashSet};
22use std::io;
23use std::path::Path;
24
25use super::reader::WalReader;
26use super::record::WalRecord;
27use super::writer::WalWriter;
28use crate::storage::engine::{Page, Pager, PAGE_SIZE};
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub enum CheckpointMode {
33 Passive,
35 Full,
37 Restart,
39 Truncate,
41}
42
43#[derive(Debug, Clone, Default)]
45pub struct CheckpointResult {
46 pub transactions_processed: u64,
48 pub max_transaction_id: u64,
50 pub pages_checkpointed: u64,
52 pub records_processed: u64,
54 pub checkpoint_lsn: u64,
56 pub wal_truncated: bool,
58}
59
60#[derive(Debug)]
62pub enum CheckpointError {
63 Io(io::Error),
65 Pager(String),
67 CorruptedWal(String),
69 NoWal,
71}
72
73impl std::fmt::Display for CheckpointError {
74 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
75 match self {
76 Self::Io(e) => write!(f, "I/O error: {}", e),
77 Self::Pager(msg) => write!(f, "Pager error: {}", msg),
78 Self::CorruptedWal(msg) => write!(f, "Corrupted WAL: {}", msg),
79 Self::NoWal => write!(f, "No WAL file found"),
80 }
81 }
82}
83
84impl std::error::Error for CheckpointError {}
85
86impl From<io::Error> for CheckpointError {
87 fn from(e: io::Error) -> Self {
88 Self::Io(e)
89 }
90}
91
92#[derive(Debug, Clone, Copy, PartialEq, Eq)]
94enum TxState {
95 Active,
96 Committed,
97 Aborted,
98}
99
100#[derive(Debug)]
102struct PendingWrite {
103 tx_id: u64,
104 page_id: u32,
105 data: Vec<u8>,
106 lsn: u64,
107}
108
109fn result_observe_tx_id(max_transaction_id: &mut u64, tx_id: u64) {
110 *max_transaction_id = (*max_transaction_id).max(tx_id);
111}
112
113pub struct Checkpointer {
117 mode: CheckpointMode,
119}
120
121impl Checkpointer {
122 pub fn new(mode: CheckpointMode) -> Self {
124 Self { mode }
125 }
126
127 pub fn default_mode() -> Self {
129 Self::new(CheckpointMode::Full)
130 }
131
132 pub fn checkpoint(
145 &self,
146 pager: &Pager,
147 wal_path: &Path,
148 ) -> Result<CheckpointResult, CheckpointError> {
149 let wal_reader = match WalReader::open(wal_path) {
151 Ok(r) => r,
152 Err(e) if e.kind() == io::ErrorKind::NotFound => {
153 return Ok(CheckpointResult::default());
155 }
156 Err(e) => return Err(CheckpointError::Io(e)),
157 };
158
159 let mut tx_states: HashMap<u64, TxState> = HashMap::new();
161 let mut pending_writes: Vec<PendingWrite> = Vec::new();
162 let mut records_processed: u64 = 0;
163 let mut last_lsn: u64 = 0;
164 let mut max_transaction_id: u64 = 0;
165
166 for record_result in wal_reader.iter() {
167 let (lsn, record) = record_result.map_err(CheckpointError::Io)?;
168 records_processed += 1;
169 last_lsn = lsn;
170
171 match record {
172 WalRecord::Begin { tx_id } => {
173 result_observe_tx_id(&mut max_transaction_id, tx_id);
174 tx_states.insert(tx_id, TxState::Active);
175 }
176 WalRecord::Commit { tx_id } => {
177 result_observe_tx_id(&mut max_transaction_id, tx_id);
178 tx_states.insert(tx_id, TxState::Committed);
179 }
180 WalRecord::Rollback { tx_id } => {
181 result_observe_tx_id(&mut max_transaction_id, tx_id);
182 tx_states.insert(tx_id, TxState::Aborted);
183 }
184 WalRecord::PageWrite {
185 tx_id,
186 page_id,
187 data,
188 } => {
189 result_observe_tx_id(&mut max_transaction_id, tx_id);
190 pending_writes.push(PendingWrite {
191 tx_id,
192 page_id,
193 data,
194 lsn,
195 });
196 }
197 WalRecord::Checkpoint {
198 lsn: _checkpoint_lsn,
199 } => {
200 }
203 WalRecord::TxCommitBatch { tx_id, .. } => {
204 result_observe_tx_id(&mut max_transaction_id, tx_id);
205 }
208 WalRecord::VectorInsert { .. } => {
209 crate::runtime::turbo_crash_inject::fire(
213 crate::runtime::turbo_crash_inject::InjectionPoint::MidCheckpoint,
214 );
215 }
216 WalRecord::ProbabilisticDelta { .. } => {
217 }
220 WalRecord::FullPageImage { tx_id, .. } => {
221 result_observe_tx_id(&mut max_transaction_id, tx_id);
222 }
226 }
227 }
228
229 let committed_txs: HashSet<u64> = tx_states
231 .iter()
232 .filter(|(_, state)| **state == TxState::Committed)
233 .map(|(tx_id, _)| *tx_id)
234 .collect();
235
236 let mut latest_writes: HashMap<u32, Vec<u8>> = HashMap::new();
239
240 for write in pending_writes {
241 if committed_txs.contains(&write.tx_id) {
242 latest_writes.insert(write.page_id, write.data);
244 }
245 }
246
247 if !latest_writes.is_empty() {
249 pager
250 .set_checkpoint_in_progress(true, last_lsn)
251 .map_err(|e| CheckpointError::Pager(e.to_string()))?;
252 }
253
254 let mut pages_checkpointed: u64 = 0;
256
257 for (page_id, data) in &latest_writes {
258 if data.len() != PAGE_SIZE {
260 return Err(CheckpointError::CorruptedWal(format!(
261 "Page {} has wrong size: {} (expected {})",
262 page_id,
263 data.len(),
264 PAGE_SIZE
265 )));
266 }
267
268 let mut page_data = [0u8; PAGE_SIZE];
269 page_data.copy_from_slice(data);
270 let page = Page::from_bytes(page_data);
271
272 pager
274 .write_page(*page_id, page)
275 .map_err(|e| CheckpointError::Pager(e.to_string()))?;
276
277 pages_checkpointed += 1;
278 }
279
280 pager
282 .sync()
283 .map_err(|e| CheckpointError::Pager(e.to_string()))?;
284
285 if !latest_writes.is_empty() {
287 pager
288 .complete_checkpoint(last_lsn)
289 .map_err(|e| CheckpointError::Pager(e.to_string()))?;
290 }
291
292 let wal_truncated = matches!(
294 self.mode,
295 CheckpointMode::Restart | CheckpointMode::Truncate
296 );
297
298 if wal_truncated {
299 let mut wal_writer = WalWriter::open(wal_path)?;
300 wal_writer.truncate()?;
301
302 let checkpoint_record = WalRecord::Checkpoint { lsn: last_lsn };
304 wal_writer.append(&checkpoint_record)?;
305 wal_writer.sync()?;
306 }
307
308 Ok(CheckpointResult {
309 transactions_processed: committed_txs.len() as u64,
310 max_transaction_id,
311 pages_checkpointed,
312 records_processed,
313 checkpoint_lsn: last_lsn,
314 wal_truncated,
315 })
316 }
317
318 pub fn recover(pager: &Pager, wal_path: &Path) -> Result<CheckpointResult, CheckpointError> {
333 if let Ok(header) = pager.header() {
335 if header.checkpoint_in_progress {
336 let _ = pager.set_checkpoint_in_progress(false, 0);
339 }
340 }
341 let checkpointer = Self::new(CheckpointMode::Truncate);
342 checkpointer.checkpoint(pager, wal_path)
343 }
344}
345
346#[cfg(test)]
347mod tests {
348 use super::*;
349 use crate::storage::engine::PageType;
350 use std::fs;
351 use std::time::{SystemTime, UNIX_EPOCH};
352
353 fn temp_dir() -> std::path::PathBuf {
354 let timestamp = SystemTime::now()
355 .duration_since(UNIX_EPOCH)
356 .unwrap()
357 .as_nanos();
358 std::env::temp_dir().join(format!("reddb_checkpoint_test_{}", timestamp))
359 }
360
361 fn cleanup(dir: &Path) {
362 let _ = fs::remove_dir_all(dir);
363 }
364
365 fn temp_wal_path(dir: &Path, name: &str) -> std::path::PathBuf {
366 reddb_file::layout::wal_component_temp_path(dir, "checkpoint", name, std::process::id())
367 }
368
369 #[test]
370 fn test_checkpoint_empty_wal() {
371 let dir = temp_dir();
372 let _ = fs::create_dir_all(&dir);
373 let db_path = dir.join("test.db");
374 let wal_path = temp_wal_path(&dir, "empty");
375
376 let pager = Pager::open_default(&db_path).unwrap();
378
379 let checkpointer = Checkpointer::default_mode();
381 let result = checkpointer.checkpoint(&pager, &wal_path).unwrap();
382
383 assert_eq!(result.transactions_processed, 0);
384 assert_eq!(result.pages_checkpointed, 0);
385
386 cleanup(&dir);
387 }
388
389 #[test]
390 fn test_checkpoint_committed_transaction() {
391 let dir = temp_dir();
392 let _ = fs::create_dir_all(&dir);
393 let db_path = dir.join("test.db");
394 let wal_path = temp_wal_path(&dir, "committed");
395
396 let pager = Pager::open_default(&db_path).unwrap();
398
399 let page = pager.allocate_page(PageType::BTreeLeaf).unwrap();
401 let page_id = page.page_id();
402
403 {
405 let mut wal_writer = WalWriter::open(&wal_path).unwrap();
406
407 wal_writer.append(&WalRecord::Begin { tx_id: 1 }).unwrap();
409
410 let mut page_data = [0u8; PAGE_SIZE];
412 page_data[0] = 0x42; wal_writer
414 .append(&WalRecord::PageWrite {
415 tx_id: 1,
416 page_id,
417 data: page_data.to_vec(),
418 })
419 .unwrap();
420
421 wal_writer.append(&WalRecord::Commit { tx_id: 1 }).unwrap();
423
424 wal_writer.sync().unwrap();
425 }
426
427 let checkpointer = Checkpointer::new(CheckpointMode::Full);
429 let result = checkpointer.checkpoint(&pager, &wal_path).unwrap();
430
431 assert_eq!(result.transactions_processed, 1);
432 assert_eq!(result.pages_checkpointed, 1);
433 assert_eq!(result.records_processed, 3);
434
435 let read_page = pager.read_page(page_id).unwrap();
437 assert_eq!(read_page.as_bytes()[0], 0x42);
438
439 cleanup(&dir);
440 }
441
442 #[test]
443 fn test_checkpoint_aborted_transaction() {
444 let dir = temp_dir();
445 let _ = fs::create_dir_all(&dir);
446 let db_path = dir.join("test.db");
447 let wal_path = temp_wal_path(&dir, "aborted");
448
449 let pager = Pager::open_default(&db_path).unwrap();
451
452 let page = pager.allocate_page(PageType::BTreeLeaf).unwrap();
454 let page_id = page.page_id();
455
456 {
458 let mut wal_writer = WalWriter::open(&wal_path).unwrap();
459
460 wal_writer.append(&WalRecord::Begin { tx_id: 1 }).unwrap();
462
463 let mut page_data = [0u8; PAGE_SIZE];
465 page_data[0] = 0x42;
466 wal_writer
467 .append(&WalRecord::PageWrite {
468 tx_id: 1,
469 page_id,
470 data: page_data.to_vec(),
471 })
472 .unwrap();
473
474 wal_writer
476 .append(&WalRecord::Rollback { tx_id: 1 })
477 .unwrap();
478
479 wal_writer.sync().unwrap();
480 }
481
482 let checkpointer = Checkpointer::new(CheckpointMode::Full);
484 let result = checkpointer.checkpoint(&pager, &wal_path).unwrap();
485
486 assert_eq!(result.transactions_processed, 0);
488 assert_eq!(result.pages_checkpointed, 0);
489
490 let read_page = pager.read_page(page_id).unwrap();
492 assert_ne!(read_page.as_bytes()[0], 0x42);
493
494 cleanup(&dir);
495 }
496
497 #[test]
498 fn test_checkpoint_mixed_transactions() {
499 let dir = temp_dir();
500 let _ = fs::create_dir_all(&dir);
501 let db_path = dir.join("test.db");
502 let wal_path = temp_wal_path(&dir, "truncate");
503
504 let pager = Pager::open_default(&db_path).unwrap();
506
507 let page1 = pager.allocate_page(PageType::BTreeLeaf).unwrap();
509 let page2 = pager.allocate_page(PageType::BTreeLeaf).unwrap();
510 let page1_id = page1.page_id();
511 let page2_id = page2.page_id();
512
513 {
515 let mut wal_writer = WalWriter::open(&wal_path).unwrap();
516
517 wal_writer.append(&WalRecord::Begin { tx_id: 1 }).unwrap();
519 let mut page_data1 = [0u8; PAGE_SIZE];
520 page_data1[0] = 0x11;
521 wal_writer
522 .append(&WalRecord::PageWrite {
523 tx_id: 1,
524 page_id: page1_id,
525 data: page_data1.to_vec(),
526 })
527 .unwrap();
528 wal_writer.append(&WalRecord::Commit { tx_id: 1 }).unwrap();
529
530 wal_writer.append(&WalRecord::Begin { tx_id: 2 }).unwrap();
532 let mut page_data2 = [0u8; PAGE_SIZE];
533 page_data2[0] = 0x22;
534 wal_writer
535 .append(&WalRecord::PageWrite {
536 tx_id: 2,
537 page_id: page2_id,
538 data: page_data2.to_vec(),
539 })
540 .unwrap();
541 wal_writer
542 .append(&WalRecord::Rollback { tx_id: 2 })
543 .unwrap();
544
545 wal_writer.append(&WalRecord::Begin { tx_id: 3 }).unwrap();
547 let mut page_data3 = [0u8; PAGE_SIZE];
548 page_data3[0] = 0x33;
549 wal_writer
550 .append(&WalRecord::PageWrite {
551 tx_id: 3,
552 page_id: page2_id,
553 data: page_data3.to_vec(),
554 })
555 .unwrap();
556 wal_writer.append(&WalRecord::Commit { tx_id: 3 }).unwrap();
557
558 wal_writer.sync().unwrap();
559 }
560
561 let checkpointer = Checkpointer::new(CheckpointMode::Full);
563 let result = checkpointer.checkpoint(&pager, &wal_path).unwrap();
564
565 assert_eq!(result.transactions_processed, 2);
567 assert_eq!(result.max_transaction_id, 3);
568 assert_eq!(result.pages_checkpointed, 2);
569
570 let read_page1 = pager.read_page(page1_id).unwrap();
572 assert_eq!(read_page1.as_bytes()[0], 0x11);
573
574 let read_page2 = pager.read_page(page2_id).unwrap();
575 assert_eq!(read_page2.as_bytes()[0], 0x33); cleanup(&dir);
578 }
579
580 #[test]
581 fn test_checkpoint_truncate() {
582 let dir = temp_dir();
583 let _ = fs::create_dir_all(&dir);
584 let db_path = dir.join("test.db");
585 let wal_path = temp_wal_path(&dir, "full-page-images");
586
587 let pager = Pager::open_default(&db_path).unwrap();
589
590 {
592 let mut wal_writer = WalWriter::open(&wal_path).unwrap();
593 wal_writer.append(&WalRecord::Begin { tx_id: 1 }).unwrap();
594 wal_writer.append(&WalRecord::Commit { tx_id: 1 }).unwrap();
595 wal_writer.sync().unwrap();
596 }
597
598 let checkpointer = Checkpointer::new(CheckpointMode::Truncate);
600 let result = checkpointer.checkpoint(&pager, &wal_path).unwrap();
601
602 assert!(result.wal_truncated);
603
604 let records: Vec<_> = WalReader::open(&wal_path)
605 .unwrap()
606 .iter()
607 .collect::<Result<Vec<_>, _>>()
608 .unwrap();
609 assert_eq!(records.len(), 1);
610 assert_eq!(
611 records[0].1,
612 WalRecord::Checkpoint {
613 lsn: result.checkpoint_lsn
614 }
615 );
616
617 cleanup(&dir);
618 }
619}