From 8cd55cd2ee4cdfceff27b9ea4cdf19d04229f5ca Mon Sep 17 00:00:00 2001 From: Benedikt Reinartz Date: Thu, 16 May 2024 14:05:29 +0200 Subject: [PATCH] Implement simple allocator using enif_{alloc,free} (#580) 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`. --- CHANGELOG.md | 4 ++ Cargo.toml | 5 ++ rustler/Cargo.toml | 1 + rustler/src/alloc.rs | 55 +++++++++++++++++++ rustler/src/lib.rs | 2 + rustler_tests/native/dynamic_load/Cargo.toml | 2 +- .../native/rustler_bigint_test/Cargo.toml | 2 +- .../native/rustler_compile_tests/Cargo.toml | 2 +- rustler_tests/native/rustler_test/Cargo.toml | 2 +- 9 files changed, 71 insertions(+), 4 deletions(-) create mode 100644 rustler/src/alloc.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index f33339d7..d0af23c7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,10 @@ versions. ## [unreleased] ### Added + +- Add optional support for using Erlang's allocator as Rust's global allocator + (#580). + ### Fixed ### Changed diff --git a/Cargo.toml b/Cargo.toml index 7f43e19e..6b494faa 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,3 +12,8 @@ members = [ "rustler_tests/native/rustler_compile_tests", "rustler_benchmarks/native/benchmark", ] +default-members = [ + "rustler", + "rustler_codegen", + "rustler_sys", +] diff --git a/rustler/Cargo.toml b/rustler/Cargo.toml index 15567975..f25bcb6f 100644 --- a/rustler/Cargo.toml +++ b/rustler/Cargo.toml @@ -14,6 +14,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"] diff --git a/rustler/src/alloc.rs b/rustler/src/alloc.rs new file mode 100644 index 00000000..61f3428c --- /dev/null +++ b/rustler/src/alloc.rs @@ -0,0 +1,55 @@ +use std::alloc::{GlobalAlloc, Layout}; + +const SIZEOF_USIZE: usize = std::mem::size_of::(); +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 and store the original pointer in memory immediately before the aligned + // section. + // + // The requested size is chosen such that we can always get an aligned buffer of size + // `layout.size()`: Ignoring `SIZEOF_USIZE`, there must always be an aligned pointer in + // the interval `[ptr, layout.align())`, so in the worst case, we have to pad with + // `layout.align() - 1`. The requirement for an additional `usize` just shifts the + // problem without changing the padding requirement. + let total_size = SIZEOF_USIZE + layout.size() + layout.align() - 1; + let ptr = rustler_sys::enif_alloc(total_size) as *mut u8; + + // Shift the returned pointer to make space for the original pointer + let ptr1 = ptr.wrapping_add(SIZEOF_USIZE); + + // Align the result to the requested alignment + let aligned_ptr = ptr1.wrapping_add(ptr1.align_offset(layout.align())); + + // Write the original pointer immediately in front of the aligned pointer + 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 { + // Retrieve the original pointer + 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); + } +} diff --git a/rustler/src/lib.rs b/rustler/src/lib.rs index fb0e1c91..46a3ffc0 100644 --- a/rustler/src/lib.rs +++ b/rustler/src/lib.rs @@ -29,6 +29,8 @@ pub mod wrapper; #[doc(hidden)] pub mod codegen_runtime; +mod alloc; + #[macro_use] pub mod types; diff --git a/rustler_tests/native/dynamic_load/Cargo.toml b/rustler_tests/native/dynamic_load/Cargo.toml index 4cc7c5d1..126263b3 100644 --- a/rustler_tests/native/dynamic_load/Cargo.toml +++ b/rustler_tests/native/dynamic_load/Cargo.toml @@ -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"] } diff --git a/rustler_tests/native/rustler_bigint_test/Cargo.toml b/rustler_tests/native/rustler_bigint_test/Cargo.toml index 9f64a529..47374e85 100644 --- a/rustler_tests/native/rustler_bigint_test/Cargo.toml +++ b/rustler_tests/native/rustler_bigint_test/Cargo.toml @@ -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"] } diff --git a/rustler_tests/native/rustler_compile_tests/Cargo.toml b/rustler_tests/native/rustler_compile_tests/Cargo.toml index f4b70a6a..9c3ca9fe 100644 --- a/rustler_tests/native/rustler_compile_tests/Cargo.toml +++ b/rustler_tests/native/rustler_compile_tests/Cargo.toml @@ -10,4 +10,4 @@ path = "src/lib.rs" crate-type = ["cdylib"] [dependencies] -rustler = { path = "../../../rustler" } +rustler = { path = "../../../rustler", features = ["allocator"] } diff --git a/rustler_tests/native/rustler_test/Cargo.toml b/rustler_tests/native/rustler_test/Cargo.toml index c2418189..acc27453 100644 --- a/rustler_tests/native/rustler_test/Cargo.toml +++ b/rustler_tests/native/rustler_test/Cargo.toml @@ -20,4 +20,4 @@ nif_version_2_16 = ["nif_version_2_15", "rustler/nif_version_2_16"] nif_version_2_17 = ["nif_version_2_16", "rustler/nif_version_2_17"] [dependencies] -rustler = { path = "../../../rustler" } +rustler = { path = "../../../rustler", features = ["allocator"] }