1#![allow(clippy::trivial_regex, clippy::collapsible_if)]
4
5use crate::page::{DEFAULT_PAGE_SIZE, Frame, PageNum};
6use akar_common::memory::MemoryManager;
7use std::collections::HashMap;
8use std::path::PathBuf;
9use std::sync::Arc;
10
11#[derive(Debug, Clone)]
13pub struct NumaInfo {
14 pub num_nodes: u32,
16}
17
18impl NumaInfo {
19 pub fn detect() -> Self {
25 let num_nodes = Self::detect_num_nodes();
26 Self { num_nodes }
27 }
28
29 #[cfg(target_os = "linux")]
30 fn detect_num_nodes() -> u32 {
31 use std::fs;
32 let node_dir = "/sys/devices/system/node";
33 if let Ok(entries) = fs::read_dir(node_dir) {
34 let count = entries
35 .filter_map(|e| e.ok())
36 .filter(|e| {
37 e.file_name()
38 .to_str()
39 .map(|s| s.starts_with("node") && s != "node")
40 .unwrap_or(false)
41 })
42 .count();
43 if count > 0 {
44 return count as u32;
45 }
46 }
47 1
48 }
49
50 #[cfg(not(target_os = "linux"))]
51 fn detect_num_nodes() -> u32 {
52 1
53 }
54}
55
56#[derive(Debug, Clone)]
58pub struct ReadaheadPolicy {
59 pub enabled: bool,
61 pub window: usize,
63}
64
65impl Default for ReadaheadPolicy {
66 fn default() -> Self {
67 Self {
68 enabled: true,
69 window: 4,
70 }
71 }
72}
73
74#[derive(Debug, Clone)]
76pub struct BufferManagerConfig {
77 pub max_memory: u64,
79 pub page_size: usize,
81 pub use_mmap: bool,
83 pub numa_aware: bool,
85 pub readahead: ReadaheadPolicy,
87}
88
89impl Default for BufferManagerConfig {
90 fn default() -> Self {
91 Self {
92 max_memory: 64 * 1024 * 1024, page_size: DEFAULT_PAGE_SIZE,
94 use_mmap: false,
95 numa_aware: false,
96 readahead: ReadaheadPolicy::default(),
97 }
98 }
99}
100
101#[derive(Debug, Default, Clone, Copy)]
103pub struct BufferManagerStats {
104 pub page_faults: u64,
106 pub page_writes: u64,
108 pub num_frames: usize,
110 pub dirty_frames: usize,
112 pub pinned_frames: usize,
114}
115
116#[derive(Debug)]
118#[allow(dead_code)]
119pub struct BufferManager {
120 db_path: PathBuf,
122 page_size: usize,
124 max_frames: usize,
126 frames: HashMap<(String, PageNum), Frame>,
129 files: HashMap<String, FileHandleInfo>,
131 clock_hand: usize,
133 clock_order: Vec<(String, PageNum)>,
135 memory_manager: Arc<MemoryManager>,
137 stats: BufferManagerStats,
139 use_mmap: bool,
141 #[cfg(not(target_arch = "wasm32"))]
143 mmap_regions: HashMap<String, memmap2::Mmap>,
144 numa_info: NumaInfo,
146 readahead: ReadaheadPolicy,
148 last_accessed: HashMap<String, PageNum>,
150 prev_last_accessed: HashMap<String, PageNum>,
152}
153
154#[allow(dead_code)]
155#[derive(Debug, Clone)]
156struct FileHandleInfo {
157 path: PathBuf,
158 num_pages: u64,
159}
160
161impl BufferManager {
162 pub fn new(db_path: PathBuf, memory_manager: Arc<MemoryManager>, config: BufferManagerConfig) -> Self {
163 let max_frames = if config.max_memory > 0 {
164 (config.max_memory / config.page_size as u64) as usize
165 } else {
166 1000
167 };
168
169 let numa_info = if config.numa_aware {
170 NumaInfo::detect()
171 } else {
172 NumaInfo { num_nodes: 1 }
173 };
174
175 Self {
176 db_path,
177 page_size: config.page_size,
178 max_frames,
179 frames: HashMap::new(),
180 files: HashMap::new(),
181 clock_hand: 0,
182 clock_order: Vec::new(),
183 memory_manager,
184 stats: BufferManagerStats::default(),
185 use_mmap: config.use_mmap,
186 #[cfg(not(target_arch = "wasm32"))]
187 mmap_regions: HashMap::new(),
188 numa_info,
189 readahead: config.readahead,
190 last_accessed: HashMap::new(),
191 prev_last_accessed: HashMap::new(),
192 }
193 }
194
195 fn key(file_name: &str, page_num: PageNum) -> (String, PageNum) {
197 (file_name.to_string(), page_num)
198 }
199
200 pub fn stats(&self) -> &BufferManagerStats {
201 &self.stats
202 }
203
204 pub fn page_size(&self) -> usize {
205 self.page_size
206 }
207
208 pub fn num_frames(&self) -> usize {
209 self.frames.len()
210 }
211
212 pub fn is_file_registered(&self, name: &str) -> bool {
214 self.files.contains_key(name)
215 }
216
217 pub fn register_file(&mut self, name: &str, path: PathBuf) {
219 let num_pages = if path.exists() {
220 let len = std::fs::metadata(&path).map(|m| m.len()).unwrap_or(0);
221 len / self.page_size as u64
222 } else {
223 0
224 };
225 self.files.insert(name.to_string(), FileHandleInfo { path, num_pages });
226 }
227
228 pub fn pin(&mut self, file_name: &str, page_num: PageNum) -> std::io::Result<Frame> {
230 let k = Self::key(file_name, page_num);
231 if let Some(frame) = self.frames.get_mut(&k) {
232 frame.pin();
233 frame.clock_ref = true;
234
235 let prev = self.last_accessed.remove(file_name);
237 if let Some(p) = prev {
238 self.prev_last_accessed.insert(file_name.to_string(), p);
239 }
240 self.last_accessed.insert(file_name.to_string(), page_num);
241
242 return Ok(frame.clone());
243 }
244
245 self.stats.page_faults += 1;
247
248 if self.frames.len() >= self.max_frames {
249 self.evict()?;
250 }
251
252 let data = self.read_from_disk(file_name, page_num)?;
253 let mut frame = Frame::new(page_num, data);
254 frame.pin();
255
256 if self.numa_info.num_nodes > 1 {
257 frame.numa_node = self.current_numa_node();
258 }
259
260 let k = Self::key(file_name, page_num);
261 self.clock_order.push(k.clone());
262 self.frames.insert(k.clone(), frame.clone());
263 self.memory_manager.allocate(self.page_size as u64);
264
265 self.update_stats();
266
267 let prev = self.last_accessed.remove(file_name);
269 if let Some(p) = prev {
270 self.prev_last_accessed.insert(file_name.to_string(), p);
271 }
272 self.last_accessed.insert(file_name.to_string(), page_num);
273
274 self.maybe_readahead(file_name, page_num);
276
277 Ok(frame)
278 }
279
280 pub fn pin_mut(&mut self, file_name: &str, page_num: PageNum) -> std::io::Result<&mut Frame> {
282 let k = Self::key(file_name, page_num);
283 if !self.frames.contains_key(&k) {
284 let _ = self.pin(file_name, page_num)?;
286 }
287 if let Some(frame) = self.frames.get_mut(&k) {
288 frame.pin();
289 frame.clock_ref = true;
290 return Ok(frame);
291 }
292 unreachable!()
293 }
294
295 pub fn unpin(&mut self, file_name: &str, page_num: PageNum) {
297 let k = Self::key(file_name, page_num);
298 if let Some(frame) = self.frames.get_mut(&k) {
299 frame.unpin();
300 }
301 self.update_stats();
302 }
303
304 pub fn flush(&mut self, file_name: &str, page_num: PageNum) -> std::io::Result<()> {
306 let k = Self::key(file_name, page_num);
307 if let Some(frame) = self.frames.get(&k) {
308 if frame.is_dirty {
309 self.write_to_disk(file_name, page_num, &frame.data)?;
310 if let Some(f) = self.frames.get_mut(&k) {
311 f.is_dirty = false;
312 }
313 self.stats.page_writes += 1;
314 }
315 }
316 self.update_stats();
317 Ok(())
318 }
319
320 pub fn flush_all(&mut self) -> std::io::Result<()> {
322 let dirty: Vec<(String, PageNum)> = self
323 .frames
324 .iter()
325 .filter(|(_, f)| f.is_dirty)
326 .map(|(k, _)| k.clone())
327 .collect();
328
329 for (file_name, page_num) in dirty {
330 self.flush(&file_name, page_num)?;
331 }
332 Ok(())
333 }
334
335 pub fn dirty_page_nums_for_file(&self, file_name: &str) -> Vec<PageNum> {
337 self.frames
338 .iter()
339 .filter(|(k, f)| k.0 == file_name && f.is_dirty)
340 .map(|(k, _)| k.1)
341 .collect()
342 }
343
344 pub fn drop_file(&mut self, file_name: &str) {
349 let path = self.files.get(file_name).map(|f| f.path.clone());
350 self.frames.retain(|(f, _), _| f != file_name);
351 self.clock_order.retain(|(f, _)| f != file_name);
352 self.files.remove(file_name);
353 self.last_accessed.remove(file_name);
354 self.prev_last_accessed.remove(file_name);
355 #[cfg(not(target_arch = "wasm32"))]
356 if let Some(p) = path {
357 let pstr = p.to_string_lossy().to_string();
358 self.mmap_regions.remove(&pstr);
359 }
360 self.update_stats();
361 }
362
363 pub fn numa_info(&self) -> &NumaInfo {
365 &self.numa_info
366 }
367
368 fn maybe_readahead(&mut self, file_name: &str, page_num: PageNum) {
373 if !self.readahead.enabled || self.readahead.window == 0 {
374 return;
375 }
376
377 if let Some(&prev) = self.prev_last_accessed.get(file_name) {
378 if page_num > 0 && prev == page_num - 1 {
379 for offset in 1..=self.readahead.window as u64 {
381 let prefetch_page = page_num + offset;
382 let pk = Self::key(file_name, prefetch_page);
383 if self.frames.contains_key(&pk) {
384 continue; }
386 if self.frames.len() >= self.max_frames {
387 break; }
389 if let Ok(data) = self.read_from_disk(file_name, prefetch_page) {
390 self.stats.page_faults += 1;
391 let mut frame = Frame::new(prefetch_page, data);
392 if self.numa_info.num_nodes > 1 {
393 frame.numa_node = self.current_numa_node();
394 }
395 self.clock_order.push(pk.clone());
396 self.frames.insert(pk, frame);
397 self.memory_manager.allocate(self.page_size as u64);
398 }
399 }
400 }
401 }
402 }
403
404 fn evict(&mut self) -> std::io::Result<()> {
407 let n = self.clock_order.len();
408 if n == 0 {
409 return Ok(());
410 }
411
412 for _ in 0..=n {
413 if self.clock_hand >= n {
414 self.clock_hand = 0;
415 }
416 let (ref file_name, ref page_num) = self.clock_order[self.clock_hand];
417 let k = Self::key(file_name, *page_num);
418
419 if let Some(frame) = self.frames.get(&k) {
420 if frame.is_pinned() {
421 self.clock_hand += 1;
422 continue;
423 }
424 if frame.clock_ref {
425 self.frames.get_mut(&k).unwrap().clock_ref = false;
426 self.clock_hand += 1;
427 continue;
428 }
429 if frame.is_dirty {
431 self.write_to_disk(file_name, *page_num, &frame.data)?;
432 self.stats.page_writes += 1;
433 }
434 self.frames.remove(&k);
435 self.clock_order.remove(self.clock_hand);
436 self.memory_manager.deallocate(self.page_size as u64);
437 return Ok(());
438 }
439 }
440 Ok(())
441 }
442
443 fn read_from_disk(&mut self, file_name: &str, page_num: PageNum) -> std::io::Result<Vec<u8>> {
447 if self.use_mmap {
448 #[cfg(not(target_arch = "wasm32"))]
449 {
450 return self.read_mmap(file_name, page_num);
451 }
452 #[cfg(target_arch = "wasm32")]
453 {
454 return self.read_syscall(file_name, page_num);
455 }
456 }
457 self.read_syscall(file_name, page_num)
458 }
459
460 fn read_syscall(&self, file_name: &str, page_num: PageNum) -> std::io::Result<Vec<u8>> {
462 if let Some(fh) = self.files.get(file_name) {
463 use std::io::{Read, Seek, SeekFrom};
464 let mut buf = vec![0u8; self.page_size];
465 match std::fs::File::open(&fh.path) {
466 Ok(mut file) => {
467 file.seek(SeekFrom::Start(page_num * self.page_size as u64))?;
468 let _ = file.read(&mut buf);
472 Ok(buf)
473 }
474 Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(buf),
475 Err(e) => Err(e),
476 }
477 } else {
478 Ok(vec![0u8; self.page_size])
479 }
480 }
481
482 #[cfg(not(target_arch = "wasm32"))]
484 fn read_mmap(&mut self, file_name: &str, page_num: PageNum) -> std::io::Result<Vec<u8>> {
485 if let Some(fh) = self.files.get(file_name) {
486 let path_str = fh.path.to_string_lossy().to_string();
487 let page_size = self.page_size;
488
489 if !self.mmap_regions.contains_key(&path_str) {
490 match std::fs::File::open(&fh.path) {
491 Ok(file) => {
492 let mmap = unsafe { memmap2::Mmap::map(&file)? };
493 self.mmap_regions.insert(path_str.clone(), mmap);
494 }
495 Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
496 return Ok(vec![0u8; page_size]);
498 }
499 Err(e) => return Err(e),
500 }
501 }
502
503 let mmap = &self.mmap_regions[&path_str];
504 let offset = page_num as usize * page_size;
505 let end = offset + page_size;
506
507 if end > mmap.len() {
508 return Ok(vec![0u8; page_size]);
511 }
512
513 let slice = unsafe { std::slice::from_raw_parts(mmap.as_ptr().add(offset), page_size) };
517 Ok(slice.to_vec())
518 } else {
519 Ok(vec![0u8; self.page_size])
520 }
521 }
522
523 fn write_to_disk(&self, file_name: &str, page_num: PageNum, data: &[u8]) -> std::io::Result<()> {
524 if let Some(fh) = self.files.get(file_name) {
525 use std::fs::OpenOptions;
526 use std::io::{Seek, SeekFrom, Write};
527 let mut file = OpenOptions::new()
528 .create(true)
529 .write(true)
530 .truncate(false)
531 .open(&fh.path)?;
532 file.seek(SeekFrom::Start(page_num * self.page_size as u64))?;
533 file.write_all(data)?;
534 }
535 Ok(())
536 }
537
538 fn current_numa_node(&self) -> u32 {
540 0
543 }
544
545 fn update_stats(&mut self) {
546 self.stats.num_frames = self.frames.len();
547 self.stats.dirty_frames = self.frames.values().filter(|f| f.is_dirty).count();
548 self.stats.pinned_frames = self.frames.values().filter(|f| f.is_pinned()).count();
549 }
550}
551
552#[cfg(test)]
553mod tests {
554 use super::*;
555
556 const TEST_FILE: &str = "main";
557
558 fn create_test_bm() -> (BufferManager, tempfile::TempDir) {
559 let dir = tempfile::tempdir().unwrap();
560 let mm = Arc::new(MemoryManager::new(1024 * 1024));
561 let config = BufferManagerConfig {
562 max_memory: 256 * 1024,
563 page_size: DEFAULT_PAGE_SIZE,
564 readahead: ReadaheadPolicy {
565 enabled: false,
566 ..Default::default()
567 },
568 ..Default::default()
569 };
570 let mut bm = BufferManager::new(dir.path().to_path_buf(), mm, config);
571 let db_path = dir.path().join("test.db");
572 std::fs::write(&db_path, vec![0u8; 8192 * 10]).unwrap();
573 bm.register_file(TEST_FILE, db_path);
574 (bm, dir)
575 }
576
577 #[test]
578 fn test_pin_unpin() {
579 let (mut bm, _dir) = create_test_bm();
580 let frame = bm.pin(TEST_FILE, 0).unwrap();
581 assert_eq!(frame.page_num, 0);
582 assert!(frame.is_pinned());
583 bm.unpin(TEST_FILE, 0);
584 assert_eq!(bm.stats().page_faults, 1);
585 }
586
587 #[test]
588 fn test_cached_page_no_fault() {
589 let (mut bm, _dir) = create_test_bm();
590 bm.pin(TEST_FILE, 0).unwrap();
591 bm.unpin(TEST_FILE, 0);
592 bm.pin(TEST_FILE, 0).unwrap();
593 bm.unpin(TEST_FILE, 0);
594 assert_eq!(bm.stats().page_faults, 1);
595 }
596
597 #[test]
598 fn test_multiple_pages() {
599 let (mut bm, _dir) = create_test_bm();
600 for i in 0..5 {
601 bm.pin(TEST_FILE, i).unwrap();
602 bm.unpin(TEST_FILE, i);
603 }
604 assert_eq!(bm.stats().page_faults, 5);
605 assert_eq!(bm.num_frames(), 5);
606 }
607
608 #[test]
609 fn test_dirty_and_flush() {
610 let (mut bm, _dir) = create_test_bm();
611 let frame = bm.pin_mut(TEST_FILE, 1).unwrap();
612 frame.mark_dirty();
613 bm.unpin(TEST_FILE, 1);
614 assert_eq!(bm.stats().dirty_frames, 1);
615 bm.flush(TEST_FILE, 1).unwrap();
616 assert_eq!(bm.stats().dirty_frames, 0);
617 assert_eq!(bm.stats().page_writes, 1);
618 }
619
620 #[test]
621 fn test_clock_eviction() {
622 let dir = tempfile::tempdir().unwrap();
623 let mm = Arc::new(MemoryManager::new(3 * DEFAULT_PAGE_SIZE as u64));
624 let config = BufferManagerConfig {
625 max_memory: 3 * DEFAULT_PAGE_SIZE as u64,
626 page_size: DEFAULT_PAGE_SIZE,
627 ..Default::default()
628 };
629 let mut bm = BufferManager::new(dir.path().to_path_buf(), mm, config);
630 let db_path = dir.path().join("test.db");
631 std::fs::write(&db_path, vec![0u8; 8192 * 20]).unwrap();
632 bm.register_file(TEST_FILE, db_path);
633
634 for i in 0..3 {
635 bm.pin(TEST_FILE, i).unwrap();
636 bm.unpin(TEST_FILE, i);
637 }
638 assert_eq!(bm.num_frames(), 3);
639
640 bm.pin(TEST_FILE, 3).unwrap();
641 bm.unpin(TEST_FILE, 3);
642 assert_eq!(bm.num_frames(), 3);
643 }
644
645 #[test]
646 fn test_flush_all() {
647 let (mut bm, _dir) = create_test_bm();
648 for i in 0..3 {
649 let frame = bm.pin_mut(TEST_FILE, i).unwrap();
650 frame.mark_dirty();
651 bm.unpin(TEST_FILE, i);
652 }
653 bm.flush_all().unwrap();
654 assert_eq!(bm.stats().dirty_frames, 0);
655 }
656
657 #[test]
660 fn test_mmap_read() {
661 let dir = tempfile::tempdir().unwrap();
662 let mm = Arc::new(MemoryManager::new(1024 * 1024));
663 let config = BufferManagerConfig {
664 max_memory: 256 * 1024,
665 page_size: DEFAULT_PAGE_SIZE,
666 use_mmap: true,
667 ..Default::default()
668 };
669 let mut bm = BufferManager::new(dir.path().to_path_buf(), mm, config);
670
671 let db_path = dir.path().join("mmap_test.db");
673 let mut data = vec![0u8; DEFAULT_PAGE_SIZE * 4];
674 for i in 0..DEFAULT_PAGE_SIZE {
675 data[i] = (i % 256) as u8;
676 }
677 std::fs::write(&db_path, &data).unwrap();
678 bm.register_file(TEST_FILE, db_path);
679
680 let frame = bm.pin(TEST_FILE, 0).unwrap();
682 assert_eq!(&frame.data[..], &data[..DEFAULT_PAGE_SIZE]);
683 bm.unpin(TEST_FILE, 0);
684
685 let frame = bm.pin(TEST_FILE, 1).unwrap();
687 assert_eq!(&frame.data[..], &data[DEFAULT_PAGE_SIZE..DEFAULT_PAGE_SIZE * 2]);
688 bm.unpin(TEST_FILE, 1);
689 }
690
691 #[test]
692 fn test_readahead_sequential() {
693 let dir = tempfile::tempdir().unwrap();
694 let mm = Arc::new(MemoryManager::new(1024 * 1024));
695 let config = BufferManagerConfig {
696 max_memory: 512 * 1024,
697 page_size: DEFAULT_PAGE_SIZE,
698 readahead: ReadaheadPolicy {
699 enabled: true,
700 window: 4,
701 },
702 ..Default::default()
703 };
704 let mut bm = BufferManager::new(dir.path().to_path_buf(), mm, config);
705
706 let db_path = dir.path().join("seq_test.db");
707 std::fs::write(&db_path, vec![42u8; DEFAULT_PAGE_SIZE * 20]).unwrap();
708 bm.register_file(TEST_FILE, db_path);
709
710 bm.pin(TEST_FILE, 0).unwrap();
712 bm.unpin(TEST_FILE, 0);
713 let faults_after_0 = bm.stats().page_faults;
714
715 bm.pin(TEST_FILE, 1).unwrap();
717 bm.unpin(TEST_FILE, 1);
718 let faults_after_1 = bm.stats().page_faults;
719
720 assert!(
722 faults_after_1 > faults_after_0 + 1,
723 "Expected readahead faults: after_page_0={}, after_page_1={}",
724 faults_after_0,
725 faults_after_1
726 );
727
728 for p in 2..=5 {
730 let k = BufferManager::key(TEST_FILE, p);
731 assert!(
732 bm.frames.contains_key(&k),
733 "Page {} should be in cache after readahead",
734 p
735 );
736 }
737
738 let faults_before = bm.stats().page_faults;
740 bm.pin(TEST_FILE, 2).unwrap();
741 bm.unpin(TEST_FILE, 2);
742 assert_eq!(bm.stats().page_faults, faults_before);
743 }
744
745 #[test]
746 fn test_readahead_random() {
747 let dir = tempfile::tempdir().unwrap();
748 let mm = Arc::new(MemoryManager::new(1024 * 1024));
749 let config = BufferManagerConfig {
750 max_memory: 512 * 1024,
751 page_size: DEFAULT_PAGE_SIZE,
752 readahead: ReadaheadPolicy {
753 enabled: true,
754 window: 4,
755 },
756 ..Default::default()
757 };
758 let mut bm = BufferManager::new(dir.path().to_path_buf(), mm, config);
759
760 let db_path = dir.path().join("rand_test.db");
761 std::fs::write(&db_path, vec![0u8; DEFAULT_PAGE_SIZE * 20]).unwrap();
762 bm.register_file(TEST_FILE, db_path);
763
764 bm.pin(TEST_FILE, 0).unwrap();
766 bm.unpin(TEST_FILE, 0);
767
768 bm.pin(TEST_FILE, 5).unwrap();
770 bm.unpin(TEST_FILE, 5);
771 let faults = bm.stats().page_faults;
772 assert_eq!(faults, 2, "Random access should not trigger readahead");
773
774 for p in 6..=9 {
776 let k = BufferManager::key(TEST_FILE, p);
777 assert!(
778 !bm.frames.contains_key(&k),
779 "Page {} should NOT be in cache after random access",
780 p
781 );
782 }
783 }
784
785 #[test]
786 fn test_numa_detection() {
787 let numa = NumaInfo::detect();
788 assert!(
789 numa.num_nodes >= 1,
790 "NUMA detection should return at least 1 node, got {}",
791 numa.num_nodes
792 );
793 }
794}