diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index b0ca0b68f..70e829703 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -114,7 +114,7 @@ jobs: path: "~/.cache/bazel" key: bazel - run: bazel build -c opt //:snmalloc - - run: bazel build -c opt //:snmalloc-rs + - run: bazel build -c opt //:snmalloc-rust-support # If this looks remarkably familiar, that's because it is. Sigh. macos: @@ -166,7 +166,7 @@ jobs: path: "~/.cache/bazel" key: bazel - run: bazel build -c opt //:snmalloc - - run: bazel build -c opt //:snmalloc-rs + - run: bazel build -c opt //:snmalloc-rust-support # GitHub doesn't natively support *BSD, but we can run them in VMs on Mac / diff --git a/BUILD.bazel b/BUILD.bazel index 0d734ad4f..6ff64e901 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -1,4 +1,7 @@ +load("@bazel_skylib//rules:native_binary.bzl", "native_binary") load("@rules_foreign_cc//foreign_cc:defs.bzl", "cmake") +load("@rules_cc//cc:defs.bzl", "cc_import") +load("@rules_rust_bindgen//:defs.bzl", "rust_bindgen_toolchain") filegroup( name = "srcs", @@ -38,7 +41,6 @@ CMAKE_FLAGS = { "SNMALLOC_OPTIMISE_FOR_CURRENT_MACHINE": "ON", "SNMALLOC_USE_SELF_VENDORED_STL": "OFF", "SNMALLOC_IPO": "ON", - "USE_SNMALLOC_STATS": "ON", } | select({ ":release_with_debug": {"CMAKE_BUILD_TYPE": "RelWithDebInfo"}, ":release": {"CMAKE_BUILD_TYPE": "Release"}, @@ -65,10 +67,11 @@ cmake( ], postfix_script = "ninja", visibility = ["//visibility:public"], + alwayslink = True, ) cmake( - name = "snmalloc-rs", + name = "snmalloc-rust-support", cache_entries = CMAKE_FLAGS | { "SNMALLOC_RUST_SUPPORT": "ON", }, @@ -87,5 +90,39 @@ cmake( "libsnmalloc-new-override.a", ], postfix_script = "ninja", + visibility = ["//src/snmalloc_rs:__pkg__"], + alwayslink = True, +) + +alias( + name = "snmalloc_rs", + actual = "//src/snmalloc_rs", visibility = ["//visibility:public"], ) + +native_binary( + name = "clang", + src = "@llvm_toolchain_llvm//:bin/clang", + visibility = ["//snmalloc_rs:__subpackages__"], +) + +cc_import( + name = "libclang", + shared_library = "@llvm_toolchain_llvm//:libclang", + visibility = ["//snmalloc_rs:__subpackages__"], +) + +rust_bindgen_toolchain( + name = "rust_bindgen_toolchain", + bindgen = "@rules_rust_bindgen//3rdparty:bindgen", + clang = ":clang", + libclang = ":libclang", + visibility = ["//snmalloc_rs:__subpackages__"], +) + +toolchain( + name = "default_bindgen_toolchain", + toolchain = ":rust_bindgen_toolchain", + toolchain_type = "@rules_rust_bindgen//:toolchain_type", + visibility = ["//snmalloc_rs:__subpackages__"], +) diff --git a/MODULE.bazel b/MODULE.bazel index 4afb06b69..1249a5fdb 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -1,6 +1,43 @@ module(name = "snmalloc") +bazel_dep(name = "bazel_skylib", version = "1.7.1") bazel_dep(name = "rules_cc", version = "0.1.1") bazel_dep(name = "rules_foreign_cc", version = "0.14.0") bazel_dep(name = "fuzztest", version = "20250214.0") bazel_dep(name = "googletest", version = "1.16.0") +bazel_dep(name = "toolchains_llvm", version = "1.4.0") +bazel_dep(name = "rules_rust", version = "0.61.0") +bazel_dep(name = "rules_rust_bindgen", version = "0.61.0") + +# Configure and register the toolchain. +# https://github.com/bazel-contrib/toolchains_llvm/blob/master/toolchain/internal/llvm_distributions.bzl +llvm = use_extension("@toolchains_llvm//toolchain/extensions:llvm.bzl", "llvm", dev_dependency = True) + +# LLVM toolchain. +llvm.toolchain( + name = "llvm_toolchain", + cxx_standard = {"": "c++20"}, + llvm_version = "17.0.6", +) +use_repo(llvm, "llvm_toolchain", "llvm_toolchain_llvm") + +#register_toolchains("@llvm_toolchain//:all") + +rust = use_extension("@rules_rust//rust:extensions.bzl", "rust", dev_dependency = True) +rust.toolchain( + edition = "2021", + versions = [ + "1.86.0", + ], +) +use_repo(rust, "rust_toolchains") + +register_toolchains( + "@rust_toolchains//:all", + dev_dependency = True, +) + +register_toolchains( + "//:default_bindgen_toolchain", + dev_dependency = True, +) diff --git a/README.md b/README.md index 52dc46ccf..2a80b8732 100644 --- a/README.md +++ b/README.md @@ -3,7 +3,7 @@ snmalloc is a high-performance allocator. snmalloc can be used directly in a project as a header-only C++ library, it can be `LD_PRELOAD`ed on Elf platforms (e.g. Linux, BSD), -and there is a [crate](https://crates.io/crates/snmalloc-rs) to use it from Rust. +and there is a [crate](https://crates.io/crates/snmalloc_rs) to use it from Rust. Its key design features are: diff --git a/fuzzing/BUILD.bazel b/fuzzing/BUILD.bazel index 0ffd8e878..841212eac 100644 --- a/fuzzing/BUILD.bazel +++ b/fuzzing/BUILD.bazel @@ -9,7 +9,7 @@ cc_test( copts = [ "-fsanitize=address", ] + select({ - "@bazel_tools//tools/cpp:clang-cl": ["-fexperimental-library"], # needed for std::execution::unseq, + "@bazel_tools//src/conditions:darwin": ["-fexperimental-library"], "//conditions:default": ["-mcx16"], }), defines = [ diff --git a/src/snmalloc_rs/BUILD.bazel b/src/snmalloc_rs/BUILD.bazel new file mode 100644 index 000000000..0fd0132dd --- /dev/null +++ b/src/snmalloc_rs/BUILD.bazel @@ -0,0 +1,65 @@ +load("@bazel_skylib//rules:common_settings.bzl", "string_flag") +load("@rules_rust//rust:defs.bzl", "rust_binary", "rust_library", "rust_test") +load("@rules_rust_bindgen//:defs.bzl", "rust_bindgen_library") +load("@rules_cc//cc:defs.bzl", "cc_library") + +cc_library( + name = "wrapper", + hdrs = ["wrapper.h"], + defines = ["SNMALLOC_USE_WAIT_ON_ADDRESS"], + linkstatic = True, + visibility = ["//visibility:private"], + deps = ["//:snmalloc-rust-support"], +) + +rust_bindgen_library( + name = "snmalloc_sys", + bindgen_flags = [ + "--allowlist-function=aligned_alloc", + "--allowlist-function=calloc", + "--allowlist-function=get_malloc_info_v1", # getting mangles + "--allowlist-function=malloc", + "--allowlist-function=malloc_usable_size", # getting mangles + "--allowlist-function=realloc", + "--allowlist-function=free", + # comment out the above and uncomment the two lines below to generate bindings for everything + # "--opaque-type=std::.*", + # "--blocklist-item=std::value", + "--use-core", # don't compile with std, allows baremetal envs + ], + cc_lib = ":wrapper", + clang_flags = [ + # "-v", # enable for debugging + "-x", + "c++", + "-std=c++20", + "-I", + ], + header = ":wrapper.h", +) + +rust_library( + name = "snmalloc_rs", + srcs = ["lib.rs"], + visibility = ["//visibility:public"], + deps = [":snmalloc_sys"], +) + +# bazel test //src/snmalloc_rs/... --extra_toolchains=@llvm_toolchain//:cc-toolchain-x86_64-linux +rust_test( + name = "test", + size = "small", + crate = ":snmalloc_rs", +) + +# bazel run //snmalloc_rs:main --extra_toolchains=@llvm_toolchain//:cc-toolchain-x86_64-linux +rust_binary( + name = "main", + srcs = ["main.rs"], + rustc_flags = [ + "-Ccodegen-units=1", + "-Copt-level=3", + "-Cstrip=symbols", + ], + deps = [":snmalloc_rs"], +) diff --git a/src/snmalloc_rs/Cargo.lock b/src/snmalloc_rs/Cargo.lock new file mode 100644 index 000000000..a0510b87c --- /dev/null +++ b/src/snmalloc_rs/Cargo.lock @@ -0,0 +1,16 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "libc" +version = "0.2.172" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d750af042f7ef4f724306de029d18836c26c1765a54a6a3f094cbd23a7267ffa" + +[[package]] +name = "snmalloc_rs" +version = "0.0.0" +dependencies = [ + "libc", +] diff --git a/src/snmalloc_rs/Cargo.toml b/src/snmalloc_rs/Cargo.toml new file mode 100644 index 000000000..9d3dc5f8c --- /dev/null +++ b/src/snmalloc_rs/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "snmalloc_rs" +version = "0.0.0" +edition = "2024" +publish = false +homepage = "https://github.com/microsoft/snmalloc" +repository = "https://github.com/microsoft/snmalloc" + + +[lib] +path = "lib.rs" + +[dependencies] +libc = "0.2.172" \ No newline at end of file diff --git a/src/snmalloc_rs/lib.rs b/src/snmalloc_rs/lib.rs new file mode 100644 index 000000000..999de365e --- /dev/null +++ b/src/snmalloc_rs/lib.rs @@ -0,0 +1,245 @@ +#![cfg_attr(not(test), no_std)] +//! `snmalloc_rs` provides a wrapper for [`microsoft/snmalloc`](https://github.com/microsoft/snmalloc) to make it usable as a global allocator for rust. +//! snmalloc is a research allocator. Its key design features are: +//! - Memory that is freed by the same thread that allocated it does not require any synchronising operations. +//! - Freeing memory in a different thread to initially allocated it, does not take any locks and instead uses a novel message passing scheme to return the memory to the original allocator, where it is recycled. +//! - The allocator uses large ranges of pages to reduce the amount of meta-data required. +//! +//! The benchmark is available at the [paper](https://github.com/microsoft/snmalloc/blob/master/snmalloc.pdf) of `snmalloc` +//! There are three features defined in this crate: +//! - `debug`: Enable the `Debug` mode in `snmalloc`. +//! - `1mib`: Use the `1mib` chunk configuration. +//! - `cache-friendly`: Make the allocator more cache friendly (setting `CACHE_FRIENDLY_OFFSET` to `64` in building the library). +//! +//! The whole library supports `no_std`. +//! +//! To use `snmalloc_rs` add it as a dependency: +//! ```toml +//! # Cargo.toml +//! [dependencies] +//! snmalloc_rs = "0.1.0" +//! ``` +//! +//! To set `SnMalloc` as the global allocator add this to your project: +//! ```rust +//! #[global_allocator] +//! static ALLOC: snmalloc_rs::SnMalloc = snmalloc_rs::SnMalloc; +//! ``` + +use core::{ + alloc::{GlobalAlloc, Layout}, + ptr::NonNull, +}; +use snmalloc_sys; + +#[derive(Debug, Copy, Clone)] +#[repr(C)] +pub struct SnMalloc; + +unsafe impl Send for SnMalloc {} +unsafe impl Sync for SnMalloc {} + +impl SnMalloc { + #[inline(always)] + pub const fn new() -> Self { + Self + } + + /// Returns the available bytes in a memory block. + // #[inline(always)] + // pub fn usable_size(&self, ptr: *const u8) -> Option { + // match ptr.is_null() { + // true => None, + // false => { + // Some(unsafe { snmalloc_sys::malloc_usable_size(ptr.cast::() as *mut _) }) + // } + // } + // } + + /// Allocates memory with the given layout, returning a non-null pointer on success + #[inline(always)] + pub fn alloc_aligned(&self, layout: Layout) -> Option> { + match layout.size() { + 0 => NonNull::new(layout.align() as *mut u8), + size => NonNull::new(unsafe { snmalloc_sys::aligned_alloc(layout.align(), size) } as *mut u8) + } + } +} + +unsafe impl GlobalAlloc for SnMalloc { + /// Allocate the memory with the given alignment and size. + /// On success, it returns a pointer pointing to the required memory address. + /// On failure, it returns a null pointer. + /// The client must assure the following things: + /// - `alignment` is greater than zero + /// - Other constrains are the same as the rust standard library. + /// + /// The program may be forced to abort if the constraints are not full-filled. + #[inline(always)] + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + match layout.size() { + 0 => layout.align() as *mut u8, + _ => snmalloc_sys::malloc(layout.size()) as *mut u8, + } + } + + /// De-allocate the memory at the given address with the given alignment and size. + /// The client must assure the following things: + /// - the memory is acquired using the same allocator and the pointer points to the start position. + /// - Other constrains are the same as the rust standard library. + /// + /// The program may be forced to abort if the constraints are not full-filled. + #[inline(always)] + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + if layout.size() != 0 { + snmalloc_sys::free(ptr as *mut _) + } + } + + /// Behaves like alloc, but also ensures that the contents are set to zero before being returned. + #[inline(always)] + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + match layout.size() { + 0 => layout.align() as *mut u8, + size => snmalloc_sys::calloc(layout.align(), size) as *mut u8, + } + } + + /// Re-allocate the memory at the given address with the given alignment and size. + /// On success, it returns a pointer pointing to the required memory address. + /// The memory content within the `new_size` will remain the same as previous. + /// On failure, it returns a null pointer. In this situation, the previous memory is not returned to the allocator. + /// The client must assure the following things: + /// - the memory is acquired using the same allocator and the pointer points to the start position + /// - `alignment` fulfills all the requirements as `rust_alloc` + /// - Other constrains are the same as the rust standard library. + /// + /// The program may be forced to abort if the constraints are not full-filled. + #[inline(always)] + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + match new_size { + 0 => { + self.dealloc(ptr, layout); + layout.align() as *mut u8 + } + new_size if layout.size() == 0 => { + self.alloc(Layout::from_size_align_unchecked(new_size, layout.align())) + } + _ => snmalloc_sys::realloc(ptr.cast(), new_size) as *mut u8, + } + } +} + +pub type SnMallocInfo = snmalloc_sys::malloc_info_v1; + +pub fn load_stats(stats: &mut SnMallocInfo) { + unsafe { snmalloc_sys::get_malloc_info_v1(stats as *mut _) } +} + + +#[cfg(test)] +mod tests { + use super::{SnMalloc, SnMallocInfo, load_stats}; + use core::alloc::{GlobalAlloc, Layout}; + #[test] + fn allocation_lifecycle() { + let alloc = SnMalloc::new(); + unsafe { + let layout = Layout::from_size_align(8, 8).unwrap(); + + // Test regular allocation + let ptr = alloc.alloc(layout); + alloc.dealloc(ptr, layout); + + // Test zeroed allocation + let ptr = alloc.alloc_zeroed(layout); + alloc.dealloc(ptr, layout); + + // Test reallocation + let ptr = alloc.alloc(layout); + let ptr = alloc.realloc(ptr, layout, 16); + alloc.dealloc(ptr, layout); + + // Test large allocation + let large_layout = Layout::from_size_align(1 << 20, 32).unwrap(); + let ptr = alloc.alloc(large_layout); + alloc.dealloc(ptr, large_layout); + } + } + #[test] + fn it_frees_allocated_memory() { + unsafe { + let layout = Layout::from_size_align(8, 8).unwrap(); + let alloc = SnMalloc; + + let ptr = alloc.alloc(layout); + alloc.dealloc(ptr, layout); + } + } + + #[test] + fn it_frees_zero_allocated_memory() { + unsafe { + let layout = Layout::from_size_align(8, 8).unwrap(); + let alloc = SnMalloc; + + let ptr = alloc.alloc_zeroed(layout); + alloc.dealloc(ptr, layout); + } + } + + #[test] + fn it_frees_reallocated_memory() { + unsafe { + let layout = Layout::from_size_align(8, 8).unwrap(); + let alloc = SnMalloc; + + let ptr = alloc.alloc(layout); + let ptr = alloc.realloc(ptr, layout, 16); + alloc.dealloc(ptr, layout); + } + } + + #[test] + fn it_frees_large_alloc() { + unsafe { + let layout = Layout::from_size_align(1 << 20, 32).unwrap(); + let alloc = SnMalloc; + + let ptr = alloc.alloc(layout); + alloc.dealloc(ptr, layout); + } + } + + // #[test] + // fn test_usable_size() { + // let alloc = SnMalloc::new(); + // unsafe { + // let layout = Layout::from_size_align(8, 8).unwrap(); + // let ptr = alloc.alloc(layout); + // let usz = alloc.usable_size::(ptr).expect("usable_size returned None"); + // alloc.dealloc(ptr, layout); + // assert!(usz >= 8); + // } + // } + + // #[test] + // fn test_stats() { + // let alloc = SnMalloc::new(); + // let mut info = SnMallocInfo { + // current_memory_usage: 0, + // peak_memory_usage: 0, + // }; + // unsafe { + // let layout = Layout::from_size_align(8, 8).unwrap(); + // let ptr = alloc.alloc(layout); + // load_stats(&mut info); + // assert_ne!(0, info.current_memory_usage); + // assert_ne!(0, info.peak_memory_usage); + // alloc.dealloc(ptr, layout); + // load_stats(&mut info); + // assert_ne!(0, info.current_memory_usage); + // assert_ne!(0, info.peak_memory_usage); + // } + // } +} \ No newline at end of file diff --git a/src/snmalloc_rs/main.rs b/src/snmalloc_rs/main.rs new file mode 100644 index 000000000..6323e4fb5 --- /dev/null +++ b/src/snmalloc_rs/main.rs @@ -0,0 +1,21 @@ +use snmalloc_rs; + +#[global_allocator] +static GLOBAL: snmalloc_rs::SnMalloc = snmalloc_rs::SnMalloc; + +fn main() { + // This `Vec` will allocate memory through `GLOBAL` above + println!("allocation a new vec"); + let mut v = Vec::new(); + println!("push an element"); + v.push(1); + println!("done"); + + let mut stats = snmalloc_rs::SnMallocInfo { + current_memory_usage: 0, + peak_memory_usage: 0, + }; + // snmalloc_rs::load_stats(&mut stats); # gets mangled + println!("current_memory_usage: {}", stats.current_memory_usage); + println!("peak_memory_usage: {}", stats.peak_memory_usage); +} \ No newline at end of file diff --git a/src/snmalloc_rs/wrapper.h b/src/snmalloc_rs/wrapper.h new file mode 100644 index 000000000..f9350a0ce --- /dev/null +++ b/src/snmalloc_rs/wrapper.h @@ -0,0 +1,3 @@ +#include +#include "snmalloc/override/rust.cc" +#include "snmalloc/override/malloc-extensions.cc"