Skip to content

Commit

Permalink
Implement simple allocator using enif_{alloc,free}
Browse files Browse the repository at this point in the history
Since we can only expect a 64bit alignment from the Erlang allocator, we
have to ensure larger alignments by overallocating. The allocator for
this case behaves the same way as the zigler "large beam allocator", see
https://github.com/E-xyza/zigler/blob/main/priv/beam/allocator.zig.

If the alignment is greater than 8, we allocate enough memory to store
an additional pointer. The result of the initial allocation is then
written immediately before the aligned pointer, which is returned from
the allocator. When deallocating, we can retrieve the original pointer
and pass it on to `enif_free`.
  • Loading branch information
filmor committed May 2, 2024
1 parent 6bd2a80 commit 182f812
Show file tree
Hide file tree
Showing 9 changed files with 57 additions and 5 deletions.
5 changes: 5 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -13,3 +13,8 @@ members = [
"rustler_tests/native/rustler_compile_tests",
"rustler_benchmarks/native/benchmark",
]
default-members = [
"rustler",
"rustler_codegen",
"rustler_sys",
]
1 change: 1 addition & 0 deletions rustler/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ big_integer = ["dep:num-bigint"]
default = ["derive", "nif_version_2_15"]
derive = ["rustler_codegen"]
alternative_nif_init_name = []
allocator = []
nif_version_2_14 = ["rustler_sys/nif_version_2_14"]
nif_version_2_15 = ["nif_version_2_14", "rustler_sys/nif_version_2_15"]
nif_version_2_16 = ["nif_version_2_15", "rustler_sys/nif_version_2_16"]
Expand Down
44 changes: 44 additions & 0 deletions rustler/src/alloc.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
use std::alloc::{GlobalAlloc, Layout};

const SIZEOF_USIZE: usize = std::mem::size_of::<usize>();
const MAX_ALIGN: usize = 8;

#[cfg(feature = "allocator")]
#[global_allocator]
static ALLOCATOR: EnifAllocator = EnifAllocator;

/// Allocator implementation that forwards all allocation calls to Erlang's allocator. Allows the
/// memory usage to be tracked by the BEAM.
pub struct EnifAllocator;

unsafe impl GlobalAlloc for EnifAllocator {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
if layout.align() > MAX_ALIGN {
// overallocate
let padded = layout.pad_to_align();
let total_size = SIZEOF_USIZE + padded.size();
let ptr = rustler_sys::enif_alloc(total_size) as *mut u8;

let ptr1 = ptr.wrapping_add(SIZEOF_USIZE);
let aligned_ptr = ptr1.wrapping_add(ptr1.align_offset(layout.align()));

let header = aligned_ptr.wrapping_sub(SIZEOF_USIZE);
*(header as *mut usize) = ptr as usize;

aligned_ptr
} else {
rustler_sys::enif_alloc(layout.size()) as *mut u8
}
}

unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
let ptr = if layout.align() > MAX_ALIGN {
let header = ptr.wrapping_sub(SIZEOF_USIZE);
let ptr = *(header as *mut usize);
ptr as *mut rustler_sys::c_void
} else {
ptr as *mut rustler_sys::c_void
};
rustler_sys::enif_free(ptr);
}
}
2 changes: 2 additions & 0 deletions rustler/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,8 @@ pub mod wrapper;
#[doc(hidden)]
pub mod codegen_runtime;

mod alloc;

pub use lazy_static;

#[macro_use]
Expand Down
2 changes: 1 addition & 1 deletion rustler_tests/native/deprecated_macros/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -10,4 +10,4 @@ path = "src/lib.rs"
crate-type = ["cdylib"]

[dependencies]
rustler = { path = "../../../rustler" }
rustler = { path = "../../../rustler", features = ["allocator"]}
2 changes: 1 addition & 1 deletion rustler_tests/native/dynamic_load/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -9,4 +9,4 @@ path = "src/lib.rs"
crate-type = ["cdylib"]

[dependencies]
rustler = { path = "../../../rustler", features = ["big_integer"] }
rustler = { path = "../../../rustler", features = ["big_integer", "allocator"] }
2 changes: 1 addition & 1 deletion rustler_tests/native/rustler_bigint_test/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -9,4 +9,4 @@ path = "src/lib.rs"
crate-type = ["cdylib"]

[dependencies]
rustler = { path = "../../../rustler", features = ["big_integer"] }
rustler = { path = "../../../rustler", features = ["big_integer", "allocator"] }
2 changes: 1 addition & 1 deletion rustler_tests/native/rustler_compile_tests/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -11,4 +11,4 @@ crate-type = ["cdylib"]

[dependencies]
lazy_static = "1.4"
rustler = { path = "../../../rustler" }
rustler = { path = "../../../rustler", features = ["allocator"] }
2 changes: 1 addition & 1 deletion rustler_tests/native/rustler_test/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -21,4 +21,4 @@ nif_version_2_17 = ["nif_version_2_16", "rustler/nif_version_2_17"]

[dependencies]
lazy_static = "1.4"
rustler = { path = "../../../rustler" }
rustler = { path = "../../../rustler", features = ["allocator"] }

0 comments on commit 182f812

Please sign in to comment.