Skip to main content

akar_storage/
checkpoint.rs

1//! Checkpoint logic — flushes WAL to main database files.
2//!
3//! During a checkpoint:
4//! 1. WAL records are flushed to the WAL file on disk
5//! 2. All dirty BufferManager pages are flushed to their respective files
6//! 3. The WAL in-memory buffer is cleared
7//! 4. A Checkpoint marker is appended to the new WAL
8
9use crate::buffer_manager::BufferManager;
10use crate::wal::{WAL, WALRecord};
11use std::sync::{Arc, Mutex};
12
13/// Result of a checkpoint operation.
14#[derive(Debug)]
15pub struct CheckpointResult {
16    pub wal_entries_processed: usize,
17    pub pages_flushed: usize,
18    pub success: bool,
19}
20
21/// Flush all dirty pages for a given table's column file to disk.
22///
23/// Iterates all frames belonging to `file_name` in the BufferManager
24/// and flushes each dirty page. Returns the number of pages flushed.
25pub fn flush_table(buffer_manager: &mut BufferManager, file_name: &str) -> std::io::Result<usize> {
26    let dirty_pages: Vec<u64> = buffer_manager.dirty_page_nums_for_file(file_name).into_iter().collect();
27
28    let count = dirty_pages.len();
29    for page_num in dirty_pages {
30        buffer_manager.flush(file_name, page_num)?;
31    }
32    Ok(count)
33}
34
35/// Perform a full checkpoint:
36///
37/// 1. Flush the WAL to disk (durability)
38/// 2. Flush all dirty BufferManager pages to disk
39/// 3. Clear the in-memory WAL buffer
40/// 4. Append a `Checkpoint` record to the new WAL
41///
42/// Returns the number of WAL entries processed and pages flushed.
43pub fn checkpoint(wal: &mut WAL, buffer_manager: &Arc<Mutex<BufferManager>>) -> std::io::Result<CheckpointResult> {
44    let wal_count = wal.len();
45
46    // Step 1: Flush the WAL to disk first (write-ahead: log before data).
47    if !wal.is_empty() {
48        wal.flush_to_disk()?;
49    }
50
51    // Step 2: Flush all dirty BufferManager pages to disk.
52    {
53        let mut bm = buffer_manager
54            .lock()
55            .map_err(|e| std::io::Error::other(format!("Lock poisoned: {e}")))?;
56        let stats_before = *bm.stats();
57        bm.flush_all()?;
58        let stats_after = *bm.stats();
59        let pages_flushed = stats_after.page_writes - stats_before.page_writes;
60
61        // Step 3: Clear the WAL after pages are durable.
62        wal.clear()?;
63
64        // Step 4: Append a Checkpoint marker and flush again.
65        wal.append(WALRecord::Checkpoint);
66        wal.flush_to_disk()?;
67
68        Ok(CheckpointResult {
69            wal_entries_processed: wal_count,
70            pages_flushed: pages_flushed as usize,
71            success: true,
72        })
73    }
74}
75
76#[cfg(test)]
77mod tests {
78    use super::*;
79    use crate::buffer_manager::BufferManagerConfig;
80    use akar_common::memory::MemoryManager;
81
82    #[test]
83    fn test_checkpoint_clears_wal() {
84        let dir = tempfile::tempdir().unwrap();
85        let wal_path = dir.path().join("wal.log");
86        let mut wal = WAL::new(wal_path);
87
88        let mm = Arc::new(MemoryManager::new(64 * 1024 * 1024));
89        let config = BufferManagerConfig::default();
90        let bm = Arc::new(Mutex::new(crate::buffer_manager::BufferManager::new(
91            dir.path().to_path_buf(),
92            mm,
93            config,
94        )));
95
96        wal.append(WALRecord::Insert {
97            table_id: 1,
98            data: vec![1, 2, 3],
99        });
100        wal.append(WALRecord::Commit { transaction_id: 42 });
101        assert!(!wal.is_empty());
102
103        let result = checkpoint(&mut wal, &bm).unwrap();
104        assert!(result.success);
105        assert_eq!(result.wal_entries_processed, 2);
106        assert!(wal.is_empty() || wal.len() == 1); // Checkpoint marker may remain
107    }
108
109    #[test]
110    fn test_checkpoint_flush_dirty_pages() {
111        let dir = tempfile::tempdir().unwrap();
112        let wal_path = dir.path().join("wal.log");
113        let mut wal = WAL::new(wal_path);
114
115        let mm = Arc::new(MemoryManager::new(64 * 1024 * 1024));
116        let config = BufferManagerConfig::default();
117        let bm = Arc::new(Mutex::new(crate::buffer_manager::BufferManager::new(
118            dir.path().to_path_buf(),
119            mm,
120            config,
121        )));
122
123        // Register a file and create a dirty page
124        {
125            let mut bm_lock = bm.lock().unwrap();
126            let db_file = dir.path().join("test.db");
127            std::fs::write(&db_file, vec![0u8; 8192 * 10]).unwrap();
128            bm_lock.register_file("test", db_file);
129            let frame = bm_lock.pin_mut("test", 0).unwrap();
130            frame.data[0..4].copy_from_slice(&[1, 2, 3, 4]);
131            frame.mark_dirty();
132            bm_lock.unpin("test", 0);
133        }
134
135        // Add WAL entries
136        wal.append(WALRecord::ColumnWrite {
137            table_id: 0,
138            col_id: 0,
139            page_id: 0,
140            data: vec![1, 2, 3, 4],
141        });
142
143        let result = checkpoint(&mut wal, &bm).unwrap();
144        assert!(result.success);
145
146        // Verify WAL file exists
147        assert!(dir.path().join("wal.log").exists());
148    }
149
150    #[test]
151    fn test_checkpoint_with_column_roundtrip() {
152        use crate::column::Column;
153        use crate::page::DEFAULT_PAGE_SIZE;
154        use akar_common::types::{LogicalTypeID, Value};
155
156        let dir = tempfile::tempdir().unwrap();
157        let wal_path = dir.path().join("wal.log");
158        let mut wal = WAL::new(wal_path);
159
160        let _mm = Arc::new(MemoryManager::new(64 * 1024 * 1024));
161        let config = BufferManagerConfig::default();
162        let bm = Arc::new(Mutex::new(crate::buffer_manager::BufferManager::new(
163            dir.path().to_path_buf(),
164            _mm,
165            config,
166        )));
167
168        // Create a column and write values via BufferManager
169        let mut col = Column::new(LogicalTypeID::Int64, 0, 0, dir.path(), bm.clone(), DEFAULT_PAGE_SIZE);
170
171        // Write data through the column (this goes through BufferManager pages)
172        for i in 0i64..10 {
173            col.append_value(&Value::Int64(i)).unwrap();
174        }
175        assert_eq!(col.num_values, 10);
176
177        // Log the column writes to WAL (simulating what the storage manager does)
178        for page_idx in 0..col.num_pages {
179            // Read the page data and log it
180            if let Ok(page_data) = col.read_value_bytes(0) {
181                wal.log_column_write(0, 0, page_idx, &page_data);
182            }
183        }
184        wal.append(WALRecord::Commit { transaction_id: 1 });
185
186        // Perform checkpoint: flush WAL + dirty pages to disk
187        let result = checkpoint(&mut wal, &bm).unwrap();
188        assert!(result.success);
189        assert!(result.wal_entries_processed > 0);
190
191        // Verify data is still readable after checkpoint
192        for i in 0i64..10 {
193            let v = col.get_value(i as u64).unwrap();
194            assert_eq!(v, Value::Int64(i));
195        }
196
197        // Verify WAL file exists and has content
198        assert!(dir.path().join("wal.log").exists());
199        let wal_meta = std::fs::metadata(dir.path().join("wal.log")).unwrap();
200        assert!(wal_meta.len() > 0, "WAL file should have content after flush");
201    }
202
203    #[test]
204    fn test_wal_column_write_replay() {
205        let dir = tempfile::tempdir().unwrap();
206        let wal_path = dir.path().join("wal.log");
207        let mut wal = WAL::new(wal_path);
208
209        // Write ColumnWrite records to WAL
210        wal.log_column_write(1, 0, 0, &[0x02, 0x2A, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]);
211        wal.log_column_write(1, 0, 1, &[0x02, 0x2B, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]);
212        wal.append(WALRecord::Commit { transaction_id: 42 });
213
214        // Replay and count ColumnWrite records
215        let mut write_count = 0;
216        wal.replay(|record| {
217            match record {
218                WALRecord::ColumnWrite {
219                    table_id,
220                    col_id,
221                    page_id,
222                    data,
223                } => {
224                    assert_eq!(*table_id, 1);
225                    assert_eq!(*col_id, 0);
226                    write_count += 1;
227                    // Verify the first record has Int64 value 42 (tag 0x02 + 0x2A)
228                    if *page_id == 0 {
229                        assert_eq!(data[0], 0x02); // tag
230                        assert_eq!(data[1], 0x2A); // 42
231                    }
232                }
233                WALRecord::Commit { transaction_id } => {
234                    assert_eq!(*transaction_id, 42);
235                }
236                _ => {}
237            }
238            Ok(())
239        })
240        .unwrap();
241        assert_eq!(write_count, 2);
242    }
243
244    // =================================================================
245    // P36.7 — Checkpoint persistence tests
246    // =================================================================
247
248    #[test]
249    fn test_flush_table_per_file() {
250        let dir = tempfile::tempdir().unwrap();
251
252        let mm = Arc::new(MemoryManager::new(64 * 1024 * 1024));
253        let config = BufferManagerConfig::default();
254        let bm = Arc::new(Mutex::new(crate::buffer_manager::BufferManager::new(
255            dir.path().to_path_buf(),
256            mm,
257            config,
258        )));
259
260        // Register a file and create dirty pages
261        {
262            let mut bm_lock = bm.lock().unwrap();
263            let db_file = dir.path().join("col_0_0");
264            std::fs::write(&db_file, vec![0u8; 8192 * 3]).unwrap();
265            bm_lock.register_file("col_0_0", db_file);
266
267            // Create 2 dirty pages
268            for page in 0..2u64 {
269                let frame = bm_lock.pin_mut("col_0_0", page).unwrap();
270                frame.data[0..4].copy_from_slice(&[1, 2, 3, 4]);
271                frame.mark_dirty();
272                bm_lock.unpin("col_0_0", page);
273            }
274        }
275
276        // flush_table should flush exactly 2 dirty pages
277        {
278            let mut bm_lock = bm.lock().unwrap();
279            let flushed = super::flush_table(&mut bm_lock, "col_0_0").unwrap();
280            assert_eq!(flushed, 2);
281        }
282
283        // After flush, no more dirty pages for this file
284        {
285            let bm_lock = bm.lock().unwrap();
286            assert!(bm_lock.dirty_page_nums_for_file("col_0_0").is_empty());
287        }
288    }
289
290    #[test]
291    fn test_column_metadata_save_load_roundtrip() {
292        use crate::column::Column;
293        use crate::page::DEFAULT_PAGE_SIZE;
294        use akar_common::types::{LogicalTypeID, Value};
295
296        let dir = tempfile::tempdir().unwrap();
297        let db_path = dir.path().to_path_buf();
298        let mm = Arc::new(MemoryManager::new(64 * 1024 * 1024));
299        let config = BufferManagerConfig::default();
300        let bm = Arc::new(Mutex::new(crate::buffer_manager::BufferManager::new(
301            db_path.clone(),
302            mm,
303            config,
304        )));
305
306        // Create a column and write data
307        let mut col = Column::new(LogicalTypeID::Int64, 0, 0, &db_path, bm.clone(), DEFAULT_PAGE_SIZE);
308        for i in 0i64..100 {
309            col.append_value(&Value::Int64(i)).unwrap();
310        }
311        assert_eq!(col.num_values, 100);
312        let orig_pages = col.num_pages;
313        let orig_offsets = col.page_row_offsets.clone();
314
315        // Save metadata
316        col.save_metadata().unwrap();
317
318        // Verify the .meta file exists
319        let meta_path = dir.path().join("col_0_0.meta");
320        assert!(meta_path.exists(), ".meta file should be created");
321
322        // Create a fresh column (no data) and load metadata
323        let mut col2 = Column::new(LogicalTypeID::Int64, 0, 0, &db_path, bm.clone(), DEFAULT_PAGE_SIZE);
324        assert_eq!(col2.num_values, 0);
325        assert_eq!(col2.num_pages, 0);
326
327        let loaded = col2.load_metadata().unwrap();
328        assert!(loaded, "metadata should be loaded");
329        assert_eq!(col2.num_values, 100);
330        assert_eq!(col2.num_pages, orig_pages);
331        assert_eq!(col2.page_row_offsets, orig_offsets);
332    }
333
334    #[test]
335    fn test_column_persistence_full_roundtrip() {
336        use crate::column::Column;
337        use crate::page::DEFAULT_PAGE_SIZE;
338        use akar_common::types::{LogicalTypeID, Value};
339
340        let dir = tempfile::tempdir().unwrap();
341        let db_path = dir.path().to_path_buf();
342        let mm = Arc::new(MemoryManager::new(64 * 1024 * 1024));
343        let config = BufferManagerConfig::default();
344        let bm = Arc::new(Mutex::new(crate::buffer_manager::BufferManager::new(
345            db_path.clone(),
346            mm,
347            config,
348        )));
349
350        // Phase 1: Create column, write data, flush to disk, save metadata
351        {
352            let mut col = Column::new(LogicalTypeID::Int64, 0, 0, &db_path, bm.clone(), DEFAULT_PAGE_SIZE);
353            for i in 0i64..256 {
354                col.append_value(&Value::Int64(i)).unwrap();
355            }
356            assert_eq!(col.num_values, 256);
357            col.flush().unwrap();
358            col.save_metadata().unwrap();
359        }
360
361        // Phase 2: Drop everything, create a fresh column and BufferManager
362        drop(bm);
363        let mm = Arc::new(MemoryManager::new(64 * 1024 * 1024));
364        let config = BufferManagerConfig::default();
365        let bm2 = Arc::new(Mutex::new(crate::buffer_manager::BufferManager::new(
366            db_path.clone(),
367            mm,
368            config,
369        )));
370
371        // Phase 3: Recreate column, load metadata, read back data
372        {
373            let mut col = Column::new(LogicalTypeID::Int64, 0, 0, &db_path, bm2.clone(), DEFAULT_PAGE_SIZE);
374            let loaded = col.load_metadata().unwrap();
375            assert!(loaded, "metadata should exist from Phase 1");
376            assert_eq!(col.num_values, 256);
377            assert!(col.num_pages > 0, "should have pages on disk");
378
379            // Read back all values from disk
380            for i in 0i64..256 {
381                let v = col.get_value(i as u64).unwrap();
382                assert_eq!(v, Value::Int64(i), "data mismatch at row {} after restart", i);
383            }
384        }
385    }
386
387    #[test]
388    fn test_checkpoint_with_column_write_to_disk() {
389        use crate::column::Column;
390        use crate::page::DEFAULT_PAGE_SIZE;
391        use akar_common::types::{LogicalTypeID, Value};
392
393        let dir = tempfile::tempdir().unwrap();
394        let wal_path = dir.path().join("wal.log");
395        let mut wal = WAL::new(wal_path);
396
397        let mm = Arc::new(MemoryManager::new(64 * 1024 * 1024));
398        let config = BufferManagerConfig::default();
399        let bm = Arc::new(Mutex::new(crate::buffer_manager::BufferManager::new(
400            dir.path().to_path_buf(),
401            mm,
402            config,
403        )));
404
405        // Write data through Column
406        let mut col = Column::new(LogicalTypeID::Int64, 0, 0, dir.path(), bm.clone(), DEFAULT_PAGE_SIZE);
407        for i in 0i64..50 {
408            col.append_value(&Value::Int64(i * 10)).unwrap();
409        }
410
411        // Log ColumnWrite to WAL
412        for page_idx in 0..col.num_pages {
413            if let Ok(page_data) = col.read_value_bytes(page_idx * 256) {
414                wal.log_column_write(0, 0, page_idx, &page_data);
415            }
416        }
417        wal.append(WALRecord::Commit { transaction_id: 1 });
418
419        // Checkpoint: flush WAL + dirty pages
420        let result = checkpoint(&mut wal, &bm).unwrap();
421        assert!(result.success);
422        assert!(result.pages_flushed > 0);
423
424        // Save column metadata so it can be reconstructed
425        col.save_metadata().unwrap();
426
427        // Verify column files exist on disk
428        let col_file = dir.path().join("col_0_0");
429        assert!(col_file.exists(), "column data file should exist after checkpoint");
430        let col_meta = dir.path().join("col_0_0.meta");
431        assert!(col_meta.exists(), "column metadata file should exist");
432
433        // Verify data is still readable
434        for i in 0i64..50 {
435            let v = col.get_value(i as u64).unwrap();
436            assert_eq!(v, Value::Int64(i * 10));
437        }
438    }
439
440    #[test]
441    fn test_wal_replay_with_column_write_records() {
442        use crate::column::Column;
443        use crate::wal_replayer::WALReplayer;
444        use akar_common::types::Value;
445
446        let dir = tempfile::tempdir().unwrap();
447        let wal_path = dir.path().join("wal.log");
448
449        // Phase 1: Write ColumnWrite records to WAL
450        {
451            let mut wal = WAL::new(wal_path.clone());
452            for i in 0i64..5 {
453                let val = Value::Int64(i * 100);
454                let raw = Column::serialize_value(&val);
455                wal.log_column_write(0, 0, 0, &raw);
456            }
457            wal.append(WALRecord::Commit { transaction_id: 10 });
458            wal.flush_to_disk().unwrap();
459        }
460
461        // Phase 2: Replay the WAL
462        let mut replayed_writes = Vec::new();
463        let result = WALReplayer::replay(&wal_path, |record| {
464            if let WALRecord::ColumnWrite {
465                table_id,
466                col_id,
467                page_id,
468                data,
469            } = record
470            {
471                replayed_writes.push((*table_id, *col_id, *page_id, data.clone()));
472            }
473            Ok(())
474        })
475        .unwrap();
476
477        assert_eq!(result.records_replayed, 5, "should replay 5 ColumnWrite records");
478        assert_eq!(replayed_writes.len(), 5);
479
480        // Verify the data in replayed records
481        for (idx, (tid, cid, pid, data)) in replayed_writes.iter().enumerate() {
482            assert_eq!(*tid, 0);
483            assert_eq!(*cid, 0);
484            assert_eq!(*pid, 0);
485            // Each record should start with TAG_INT64 = 2
486            assert_eq!(data[0], 2, "should be Int64 tag");
487            let val = i64::from_le_bytes(data[1..9].try_into().unwrap());
488            assert_eq!(val, (idx as i64) * 100);
489        }
490    }
491
492    #[test]
493    fn test_multi_column_checkpoint_persistence() {
494        use crate::column::Column;
495        use crate::page::DEFAULT_PAGE_SIZE;
496        use akar_common::types::{LogicalTypeID, Value};
497
498        let dir = tempfile::tempdir().unwrap();
499        let db_path = dir.path().to_path_buf();
500        let mm = Arc::new(MemoryManager::new(64 * 1024 * 1024));
501        let config = BufferManagerConfig::default();
502        let bm = Arc::new(Mutex::new(crate::buffer_manager::BufferManager::new(
503            db_path.clone(),
504            mm,
505            config,
506        )));
507
508        // Create two columns (simulating a table with id + name)
509        let mut col_id = Column::new(LogicalTypeID::Int64, 1, 0, &db_path, bm.clone(), DEFAULT_PAGE_SIZE);
510        let mut col_name = Column::new(LogicalTypeID::Int64, 1, 1, &db_path, bm.clone(), DEFAULT_PAGE_SIZE);
511
512        // Write data
513        for i in 0i64..100 {
514            col_id.append_value(&Value::Int64(i)).unwrap();
515            col_name.append_value(&Value::Int64(i * 1000)).unwrap();
516        }
517
518        // Flush both columns
519        col_id.flush().unwrap();
520        col_name.flush().unwrap();
521
522        // Save metadata for both
523        col_id.save_metadata().unwrap();
524        col_name.save_metadata().unwrap();
525
526        // Drop everything
527        drop(bm);
528        drop(col_id);
529        drop(col_name);
530
531        // Create fresh BufferManager and Columns
532        let mm = Arc::new(MemoryManager::new(64 * 1024 * 1024));
533        let config = BufferManagerConfig::default();
534        let bm2 = Arc::new(Mutex::new(crate::buffer_manager::BufferManager::new(
535            db_path.clone(),
536            mm,
537            config,
538        )));
539
540        let mut col_id2 = Column::new(LogicalTypeID::Int64, 1, 0, &db_path, bm2.clone(), DEFAULT_PAGE_SIZE);
541        let mut col_name2 = Column::new(LogicalTypeID::Int64, 1, 1, &db_path, bm2.clone(), DEFAULT_PAGE_SIZE);
542
543        col_id2.load_metadata().unwrap();
544        col_name2.load_metadata().unwrap();
545
546        assert_eq!(col_id2.num_values, 100);
547        assert_eq!(col_name2.num_values, 100);
548
549        // Read back all values from both columns
550        for i in 0i64..100 {
551            assert_eq!(col_id2.get_value(i as u64).unwrap(), Value::Int64(i));
552            assert_eq!(col_name2.get_value(i as u64).unwrap(), Value::Int64(i * 1000));
553        }
554    }
555}