#![allow(clippy::trivial_regex, clippy::collapsible_if)]
use crate::page::{DEFAULT_PAGE_SIZE, Frame, PageNum};
use akar_common::memory::MemoryManager;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
#[derive(Debug, Clone)]
pub struct NumaInfo {
pub num_nodes: u32,
}
impl NumaInfo {
pub fn detect() -> Self {
let num_nodes = Self::detect_num_nodes();
Self { num_nodes }
}
#[cfg(target_os = "linux")]
fn detect_num_nodes() -> u32 {
use std::fs;
let node_dir = "/sys/devices/system/node";
if let Ok(entries) = fs::read_dir(node_dir) {
let count = entries
.filter_map(|e| e.ok())
.filter(|e| {
e.file_name()
.to_str()
.map(|s| s.starts_with("node") && s != "node")
.unwrap_or(false)
})
.count();
if count > 0 {
return count as u32;
}
}
1
}
#[cfg(not(target_os = "linux"))]
fn detect_num_nodes() -> u32 {
1
}
}
#[derive(Debug, Clone)]
pub struct ReadaheadPolicy {
pub enabled: bool,
pub window: usize,
}
impl Default for ReadaheadPolicy {
fn default() -> Self {
Self {
enabled: true,
window: 4,
}
}
}
#[derive(Debug, Clone)]
pub struct BufferManagerConfig {
pub max_memory: u64,
pub page_size: usize,
pub use_mmap: bool,
pub numa_aware: bool,
pub readahead: ReadaheadPolicy,
}
impl Default for BufferManagerConfig {
fn default() -> Self {
Self {
max_memory: 64 * 1024 * 1024, page_size: DEFAULT_PAGE_SIZE,
use_mmap: false,
numa_aware: false,
readahead: ReadaheadPolicy::default(),
}
}
}
#[derive(Debug, Default, Clone, Copy)]
pub struct BufferManagerStats {
pub page_faults: u64,
pub page_writes: u64,
pub num_frames: usize,
pub dirty_frames: usize,
pub pinned_frames: usize,
}
#[derive(Debug)]
#[allow(dead_code)]
pub struct BufferManager {
db_path: PathBuf,
page_size: usize,
max_frames: usize,
frames: HashMap<(String, PageNum), Frame>,
files: HashMap<String, FileHandleInfo>,
clock_hand: usize,
clock_order: Vec<(String, PageNum)>,
memory_manager: Arc<MemoryManager>,
stats: BufferManagerStats,
use_mmap: bool,
#[cfg(not(target_arch = "wasm32"))]
mmap_regions: HashMap<String, memmap2::Mmap>,
numa_info: NumaInfo,
readahead: ReadaheadPolicy,
last_accessed: HashMap<String, PageNum>,
prev_last_accessed: HashMap<String, PageNum>,
}
#[allow(dead_code)]
#[derive(Debug, Clone)]
struct FileHandleInfo {
path: PathBuf,
num_pages: u64,
}
impl BufferManager {
pub fn new(db_path: PathBuf, memory_manager: Arc<MemoryManager>, config: BufferManagerConfig) -> Self {
let max_frames = if config.max_memory > 0 {
(config.max_memory / config.page_size as u64) as usize
} else {
1000
};
let numa_info = if config.numa_aware {
NumaInfo::detect()
} else {
NumaInfo { num_nodes: 1 }
};
Self {
db_path,
page_size: config.page_size,
max_frames,
frames: HashMap::new(),
files: HashMap::new(),
clock_hand: 0,
clock_order: Vec::new(),
memory_manager,
stats: BufferManagerStats::default(),
use_mmap: config.use_mmap,
#[cfg(not(target_arch = "wasm32"))]
mmap_regions: HashMap::new(),
numa_info,
readahead: config.readahead,
last_accessed: HashMap::new(),
prev_last_accessed: HashMap::new(),
}
}
fn key(file_name: &str, page_num: PageNum) -> (String, PageNum) {
(file_name.to_string(), page_num)
}
pub fn stats(&self) -> &BufferManagerStats {
&self.stats
}
pub fn page_size(&self) -> usize {
self.page_size
}
pub fn num_frames(&self) -> usize {
self.frames.len()
}
pub fn is_file_registered(&self, name: &str) -> bool {
self.files.contains_key(name)
}
pub fn register_file(&mut self, name: &str, path: PathBuf) {
let num_pages = if path.exists() {
let len = std::fs::metadata(&path).map(|m| m.len()).unwrap_or(0);
len / self.page_size as u64
} else {
0
};
self.files.insert(name.to_string(), FileHandleInfo { path, num_pages });
}
pub fn pin(&mut self, file_name: &str, page_num: PageNum) -> std::io::Result<Frame> {
let k = Self::key(file_name, page_num);
if let Some(frame) = self.frames.get_mut(&k) {
frame.pin();
frame.clock_ref = true;
let prev = self.last_accessed.remove(file_name);
if let Some(p) = prev {
self.prev_last_accessed.insert(file_name.to_string(), p);
}
self.last_accessed.insert(file_name.to_string(), page_num);
return Ok(frame.clone());
}
self.stats.page_faults += 1;
if self.frames.len() >= self.max_frames {
self.evict()?;
}
let data = self.read_from_disk(file_name, page_num)?;
let mut frame = Frame::new(page_num, data);
frame.pin();
if self.numa_info.num_nodes > 1 {
frame.numa_node = self.current_numa_node();
}
let k = Self::key(file_name, page_num);
self.clock_order.push(k.clone());
self.frames.insert(k.clone(), frame.clone());
self.memory_manager.allocate(self.page_size as u64);
self.update_stats();
let prev = self.last_accessed.remove(file_name);
if let Some(p) = prev {
self.prev_last_accessed.insert(file_name.to_string(), p);
}
self.last_accessed.insert(file_name.to_string(), page_num);
self.maybe_readahead(file_name, page_num);
Ok(frame)
}
pub fn pin_mut(&mut self, file_name: &str, page_num: PageNum) -> std::io::Result<&mut Frame> {
let k = Self::key(file_name, page_num);
if !self.frames.contains_key(&k) {
let _ = self.pin(file_name, page_num)?;
}
if let Some(frame) = self.frames.get_mut(&k) {
frame.pin();
frame.clock_ref = true;
return Ok(frame);
}
unreachable!()
}
pub fn unpin(&mut self, file_name: &str, page_num: PageNum) {
let k = Self::key(file_name, page_num);
if let Some(frame) = self.frames.get_mut(&k) {
frame.unpin();
}
self.update_stats();
}
pub fn flush(&mut self, file_name: &str, page_num: PageNum) -> std::io::Result<()> {
let k = Self::key(file_name, page_num);
if let Some(frame) = self.frames.get(&k) {
if frame.is_dirty {
self.write_to_disk(file_name, page_num, &frame.data)?;
if let Some(f) = self.frames.get_mut(&k) {
f.is_dirty = false;
}
self.stats.page_writes += 1;
}
}
self.update_stats();
Ok(())
}
pub fn flush_all(&mut self) -> std::io::Result<()> {
let dirty: Vec<(String, PageNum)> = self
.frames
.iter()
.filter(|(_, f)| f.is_dirty)
.map(|(k, _)| k.clone())
.collect();
for (file_name, page_num) in dirty {
self.flush(&file_name, page_num)?;
}
Ok(())
}
pub fn dirty_page_nums_for_file(&self, file_name: &str) -> Vec<PageNum> {
self.frames
.iter()
.filter(|(k, f)| k.0 == file_name && f.is_dirty)
.map(|(k, _)| k.1)
.collect()
}
pub fn drop_file(&mut self, file_name: &str) {
let path = self.files.get(file_name).map(|f| f.path.clone());
self.frames.retain(|(f, _), _| f != file_name);
self.clock_order.retain(|(f, _)| f != file_name);
self.files.remove(file_name);
self.last_accessed.remove(file_name);
self.prev_last_accessed.remove(file_name);
#[cfg(not(target_arch = "wasm32"))]
if let Some(p) = path {
let pstr = p.to_string_lossy().to_string();
self.mmap_regions.remove(&pstr);
}
self.update_stats();
}
pub fn numa_info(&self) -> &NumaInfo {
&self.numa_info
}
fn maybe_readahead(&mut self, file_name: &str, page_num: PageNum) {
if !self.readahead.enabled || self.readahead.window == 0 {
return;
}
if let Some(&prev) = self.prev_last_accessed.get(file_name) {
if page_num > 0 && prev == page_num - 1 {
for offset in 1..=self.readahead.window as u64 {
let prefetch_page = page_num + offset;
let pk = Self::key(file_name, prefetch_page);
if self.frames.contains_key(&pk) {
continue; }
if self.frames.len() >= self.max_frames {
break; }
if let Ok(data) = self.read_from_disk(file_name, prefetch_page) {
self.stats.page_faults += 1;
let mut frame = Frame::new(prefetch_page, data);
if self.numa_info.num_nodes > 1 {
frame.numa_node = self.current_numa_node();
}
self.clock_order.push(pk.clone());
self.frames.insert(pk, frame);
self.memory_manager.allocate(self.page_size as u64);
}
}
}
}
}
fn evict(&mut self) -> std::io::Result<()> {
let n = self.clock_order.len();
if n == 0 {
return Ok(());
}
for _ in 0..=n {
if self.clock_hand >= n {
self.clock_hand = 0;
}
let (ref file_name, ref page_num) = self.clock_order[self.clock_hand];
let k = Self::key(file_name, *page_num);
if let Some(frame) = self.frames.get(&k) {
if frame.is_pinned() {
self.clock_hand += 1;
continue;
}
if frame.clock_ref {
self.frames.get_mut(&k).unwrap().clock_ref = false;
self.clock_hand += 1;
continue;
}
if frame.is_dirty {
self.write_to_disk(file_name, *page_num, &frame.data)?;
self.stats.page_writes += 1;
}
self.frames.remove(&k);
self.clock_order.remove(self.clock_hand);
self.memory_manager.deallocate(self.page_size as u64);
return Ok(());
}
}
Ok(())
}
fn read_from_disk(&mut self, file_name: &str, page_num: PageNum) -> std::io::Result<Vec<u8>> {
if self.use_mmap {
#[cfg(not(target_arch = "wasm32"))]
{
return self.read_mmap(file_name, page_num);
}
#[cfg(target_arch = "wasm32")]
{
return self.read_syscall(file_name, page_num);
}
}
self.read_syscall(file_name, page_num)
}
fn read_syscall(&self, file_name: &str, page_num: PageNum) -> std::io::Result<Vec<u8>> {
if let Some(fh) = self.files.get(file_name) {
use std::io::{Read, Seek, SeekFrom};
let mut buf = vec![0u8; self.page_size];
match std::fs::File::open(&fh.path) {
Ok(mut file) => {
file.seek(SeekFrom::Start(page_num * self.page_size as u64))?;
let _ = file.read(&mut buf);
Ok(buf)
}
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(buf),
Err(e) => Err(e),
}
} else {
Ok(vec![0u8; self.page_size])
}
}
#[cfg(not(target_arch = "wasm32"))]
fn read_mmap(&mut self, file_name: &str, page_num: PageNum) -> std::io::Result<Vec<u8>> {
if let Some(fh) = self.files.get(file_name) {
let path_str = fh.path.to_string_lossy().to_string();
let page_size = self.page_size;
if !self.mmap_regions.contains_key(&path_str) {
match std::fs::File::open(&fh.path) {
Ok(file) => {
let mmap = unsafe { memmap2::Mmap::map(&file)? };
self.mmap_regions.insert(path_str.clone(), mmap);
}
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
return Ok(vec![0u8; page_size]);
}
Err(e) => return Err(e),
}
}
let mmap = &self.mmap_regions[&path_str];
let offset = page_num as usize * page_size;
let end = offset + page_size;
if end > mmap.len() {
return Ok(vec![0u8; page_size]);
}
let slice = unsafe { std::slice::from_raw_parts(mmap.as_ptr().add(offset), page_size) };
Ok(slice.to_vec())
} else {
Ok(vec![0u8; self.page_size])
}
}
fn write_to_disk(&self, file_name: &str, page_num: PageNum, data: &[u8]) -> std::io::Result<()> {
if let Some(fh) = self.files.get(file_name) {
use std::fs::OpenOptions;
use std::io::{Seek, SeekFrom, Write};
let mut file = OpenOptions::new()
.create(true)
.write(true)
.truncate(false)
.open(&fh.path)?;
file.seek(SeekFrom::Start(page_num * self.page_size as u64))?;
file.write_all(data)?;
}
Ok(())
}
fn current_numa_node(&self) -> u32 {
0
}
fn update_stats(&mut self) {
self.stats.num_frames = self.frames.len();
self.stats.dirty_frames = self.frames.values().filter(|f| f.is_dirty).count();
self.stats.pinned_frames = self.frames.values().filter(|f| f.is_pinned()).count();
}
}
#[cfg(test)]
mod tests {
use super::*;
const TEST_FILE: &str = "main";
fn create_test_bm() -> (BufferManager, tempfile::TempDir) {
let dir = tempfile::tempdir().unwrap();
let mm = Arc::new(MemoryManager::new(1024 * 1024));
let config = BufferManagerConfig {
max_memory: 256 * 1024,
page_size: DEFAULT_PAGE_SIZE,
readahead: ReadaheadPolicy {
enabled: false,
..Default::default()
},
..Default::default()
};
let mut bm = BufferManager::new(dir.path().to_path_buf(), mm, config);
let db_path = dir.path().join("test.db");
std::fs::write(&db_path, vec![0u8; 8192 * 10]).unwrap();
bm.register_file(TEST_FILE, db_path);
(bm, dir)
}
#[test]
fn test_pin_unpin() {
let (mut bm, _dir) = create_test_bm();
let frame = bm.pin(TEST_FILE, 0).unwrap();
assert_eq!(frame.page_num, 0);
assert!(frame.is_pinned());
bm.unpin(TEST_FILE, 0);
assert_eq!(bm.stats().page_faults, 1);
}
#[test]
fn test_cached_page_no_fault() {
let (mut bm, _dir) = create_test_bm();
bm.pin(TEST_FILE, 0).unwrap();
bm.unpin(TEST_FILE, 0);
bm.pin(TEST_FILE, 0).unwrap();
bm.unpin(TEST_FILE, 0);
assert_eq!(bm.stats().page_faults, 1);
}
#[test]
fn test_multiple_pages() {
let (mut bm, _dir) = create_test_bm();
for i in 0..5 {
bm.pin(TEST_FILE, i).unwrap();
bm.unpin(TEST_FILE, i);
}
assert_eq!(bm.stats().page_faults, 5);
assert_eq!(bm.num_frames(), 5);
}
#[test]
fn test_dirty_and_flush() {
let (mut bm, _dir) = create_test_bm();
let frame = bm.pin_mut(TEST_FILE, 1).unwrap();
frame.mark_dirty();
bm.unpin(TEST_FILE, 1);
assert_eq!(bm.stats().dirty_frames, 1);
bm.flush(TEST_FILE, 1).unwrap();
assert_eq!(bm.stats().dirty_frames, 0);
assert_eq!(bm.stats().page_writes, 1);
}
#[test]
fn test_clock_eviction() {
let dir = tempfile::tempdir().unwrap();
let mm = Arc::new(MemoryManager::new(3 * DEFAULT_PAGE_SIZE as u64));
let config = BufferManagerConfig {
max_memory: 3 * DEFAULT_PAGE_SIZE as u64,
page_size: DEFAULT_PAGE_SIZE,
..Default::default()
};
let mut bm = BufferManager::new(dir.path().to_path_buf(), mm, config);
let db_path = dir.path().join("test.db");
std::fs::write(&db_path, vec![0u8; 8192 * 20]).unwrap();
bm.register_file(TEST_FILE, db_path);
for i in 0..3 {
bm.pin(TEST_FILE, i).unwrap();
bm.unpin(TEST_FILE, i);
}
assert_eq!(bm.num_frames(), 3);
bm.pin(TEST_FILE, 3).unwrap();
bm.unpin(TEST_FILE, 3);
assert_eq!(bm.num_frames(), 3);
}
#[test]
fn test_flush_all() {
let (mut bm, _dir) = create_test_bm();
for i in 0..3 {
let frame = bm.pin_mut(TEST_FILE, i).unwrap();
frame.mark_dirty();
bm.unpin(TEST_FILE, i);
}
bm.flush_all().unwrap();
assert_eq!(bm.stats().dirty_frames, 0);
}
#[test]
fn test_mmap_read() {
let dir = tempfile::tempdir().unwrap();
let mm = Arc::new(MemoryManager::new(1024 * 1024));
let config = BufferManagerConfig {
max_memory: 256 * 1024,
page_size: DEFAULT_PAGE_SIZE,
use_mmap: true,
..Default::default()
};
let mut bm = BufferManager::new(dir.path().to_path_buf(), mm, config);
let db_path = dir.path().join("mmap_test.db");
let mut data = vec![0u8; DEFAULT_PAGE_SIZE * 4];
for i in 0..DEFAULT_PAGE_SIZE {
data[i] = (i % 256) as u8;
}
std::fs::write(&db_path, &data).unwrap();
bm.register_file(TEST_FILE, db_path);
let frame = bm.pin(TEST_FILE, 0).unwrap();
assert_eq!(&frame.data[..], &data[..DEFAULT_PAGE_SIZE]);
bm.unpin(TEST_FILE, 0);
let frame = bm.pin(TEST_FILE, 1).unwrap();
assert_eq!(&frame.data[..], &data[DEFAULT_PAGE_SIZE..DEFAULT_PAGE_SIZE * 2]);
bm.unpin(TEST_FILE, 1);
}
#[test]
fn test_readahead_sequential() {
let dir = tempfile::tempdir().unwrap();
let mm = Arc::new(MemoryManager::new(1024 * 1024));
let config = BufferManagerConfig {
max_memory: 512 * 1024,
page_size: DEFAULT_PAGE_SIZE,
readahead: ReadaheadPolicy {
enabled: true,
window: 4,
},
..Default::default()
};
let mut bm = BufferManager::new(dir.path().to_path_buf(), mm, config);
let db_path = dir.path().join("seq_test.db");
std::fs::write(&db_path, vec![42u8; DEFAULT_PAGE_SIZE * 20]).unwrap();
bm.register_file(TEST_FILE, db_path);
bm.pin(TEST_FILE, 0).unwrap();
bm.unpin(TEST_FILE, 0);
let faults_after_0 = bm.stats().page_faults;
bm.pin(TEST_FILE, 1).unwrap();
bm.unpin(TEST_FILE, 1);
let faults_after_1 = bm.stats().page_faults;
assert!(
faults_after_1 > faults_after_0 + 1,
"Expected readahead faults: after_page_0={}, after_page_1={}",
faults_after_0,
faults_after_1
);
for p in 2..=5 {
let k = BufferManager::key(TEST_FILE, p);
assert!(
bm.frames.contains_key(&k),
"Page {} should be in cache after readahead",
p
);
}
let faults_before = bm.stats().page_faults;
bm.pin(TEST_FILE, 2).unwrap();
bm.unpin(TEST_FILE, 2);
assert_eq!(bm.stats().page_faults, faults_before);
}
#[test]
fn test_readahead_random() {
let dir = tempfile::tempdir().unwrap();
let mm = Arc::new(MemoryManager::new(1024 * 1024));
let config = BufferManagerConfig {
max_memory: 512 * 1024,
page_size: DEFAULT_PAGE_SIZE,
readahead: ReadaheadPolicy {
enabled: true,
window: 4,
},
..Default::default()
};
let mut bm = BufferManager::new(dir.path().to_path_buf(), mm, config);
let db_path = dir.path().join("rand_test.db");
std::fs::write(&db_path, vec![0u8; DEFAULT_PAGE_SIZE * 20]).unwrap();
bm.register_file(TEST_FILE, db_path);
bm.pin(TEST_FILE, 0).unwrap();
bm.unpin(TEST_FILE, 0);
bm.pin(TEST_FILE, 5).unwrap();
bm.unpin(TEST_FILE, 5);
let faults = bm.stats().page_faults;
assert_eq!(faults, 2, "Random access should not trigger readahead");
for p in 6..=9 {
let k = BufferManager::key(TEST_FILE, p);
assert!(
!bm.frames.contains_key(&k),
"Page {} should NOT be in cache after random access",
p
);
}
}
#[test]
fn test_numa_detection() {
let numa = NumaInfo::detect();
assert!(
numa.num_nodes >= 1,
"NUMA detection should return at least 1 node, got {}",
numa.num_nodes
);
}
}