130 lines
3.2 KiB
Rust
130 lines
3.2 KiB
Rust
use std::{
|
|
sync::{LazyLock, Mutex},
|
|
usize,
|
|
};
|
|
|
|
use libc::{c_char, c_void, ftruncate, memfd_create, mmap, MAP_SHARED, PROT_READ, PROT_WRITE};
|
|
|
|
pub(crate) const MEMFD_INITIAL_SIZE: usize = 1024 * 1024 * 4;
|
|
const MMAP_SIZE: usize = 1024 * 1024 * 1024;
|
|
|
|
pub(crate) static BUMP_ALLOCATOR: LazyLock<Mutex<BumpAllocator>> =
|
|
LazyLock::new(|| unsafe { Mutex::new(BumpAllocator::new()) });
|
|
|
|
pub struct BumpAllocator {
|
|
start_of_mem: *mut u8,
|
|
head: *mut u8,
|
|
end_of_mem: *mut u8,
|
|
number_of_allocated_chunks: usize,
|
|
backing_fd: i32,
|
|
fd_size: usize,
|
|
}
|
|
|
|
unsafe impl Send for BumpAllocator {}
|
|
|
|
impl BumpAllocator {
|
|
unsafe fn new() -> Self {
|
|
assert!(MMAP_SIZE >= MEMFD_INITIAL_SIZE);
|
|
|
|
let data_fd = memfd_create("data\x00".as_ptr() as *const c_char, 0);
|
|
|
|
assert!(data_fd > 0);
|
|
|
|
assert_eq!(ftruncate(data_fd, MEMFD_INITIAL_SIZE as i64), 0);
|
|
|
|
let start_of_mem = mmap(
|
|
0 as *mut c_void,
|
|
MMAP_SIZE,
|
|
PROT_READ | PROT_WRITE,
|
|
MAP_SHARED,
|
|
data_fd,
|
|
0,
|
|
) as *mut u8;
|
|
|
|
assert_ne!(start_of_mem, 0 as *mut u8);
|
|
|
|
let end_of_mem = start_of_mem.byte_add(MEMFD_INITIAL_SIZE);
|
|
|
|
BumpAllocator {
|
|
start_of_mem,
|
|
head: start_of_mem,
|
|
end_of_mem,
|
|
number_of_allocated_chunks: 0,
|
|
backing_fd: data_fd,
|
|
fd_size: MEMFD_INITIAL_SIZE,
|
|
}
|
|
}
|
|
|
|
pub(crate) unsafe fn alloc(&mut self, size: usize) -> Option<*mut u8> {
|
|
let new_head = self.head.byte_add(size);
|
|
|
|
if new_head > self.end_of_mem {
|
|
if self.end_of_mem.byte_add(self.fd_size) < self.start_of_mem.byte_add(MMAP_SIZE) {
|
|
assert_eq!(ftruncate(self.backing_fd, (self.fd_size * 2) as i64), 0);
|
|
self.fd_size *= 2;
|
|
} else {
|
|
return None;
|
|
}
|
|
}
|
|
|
|
let ret = Some(self.head);
|
|
|
|
self.head = new_head;
|
|
self.number_of_allocated_chunks += 1;
|
|
|
|
ret
|
|
}
|
|
|
|
pub(crate) fn dealloc(&mut self) {
|
|
self.number_of_allocated_chunks -= 1;
|
|
|
|
if self.number_of_allocated_chunks == 0 {
|
|
self.head = self.start_of_mem;
|
|
}
|
|
}
|
|
|
|
pub(crate) unsafe fn get_offset(&self, ptr: *const u8) -> usize {
|
|
let offset = ptr.byte_offset_from(self.start_of_mem);
|
|
|
|
debug_assert!(offset >= 0);
|
|
debug_assert!(offset < self.end_of_mem.byte_offset_from(self.start_of_mem));
|
|
|
|
offset as usize
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
extern crate test;
|
|
|
|
use core::slice;
|
|
use test::Bencher;
|
|
|
|
use super::BUMP_ALLOCATOR;
|
|
|
|
#[test]
|
|
fn test() {
|
|
let mut allocator = BUMP_ALLOCATOR.lock().unwrap();
|
|
|
|
unsafe {
|
|
let x = allocator.alloc(10).unwrap();
|
|
let x = slice::from_raw_parts_mut(x, 10);
|
|
x[0] = 1;
|
|
assert_eq!(x[0], 1);
|
|
}
|
|
}
|
|
|
|
#[bench]
|
|
fn bench(b: &mut Bencher) {
|
|
let mut allocator = BUMP_ALLOCATOR.lock().unwrap();
|
|
|
|
b.iter(|| unsafe {
|
|
let x = allocator.alloc(10).unwrap();
|
|
let x = slice::from_raw_parts_mut(x, 10);
|
|
x[0] = 1;
|
|
assert_eq!(x[0], 1);
|
|
allocator.dealloc();
|
|
});
|
|
}
|
|
}
|