diff --git a/.bazelrc b/.bazelrc index ffcd8bbce..aaa55a557 100644 --- a/.bazelrc +++ b/.bazelrc @@ -23,6 +23,9 @@ build --incompatible_disallow_empty_glob # for details. common --incompatible_config_setting_private_default_visibility +# cc static lib for rust +build --experimental_cc_static_library + # Link Google Test against Abseil and RE2. build --define='absl=1' diff --git a/BUILD b/BUILD index 9061795cc..9ba304175 100644 --- a/BUILD +++ b/BUILD @@ -11,3 +11,41 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. + +load("@bazel_skylib//rules:native_binary.bzl", "native_binary") +load("@rules_cc//cc:defs.bzl", "cc_import") +load("@rules_rust_bindgen//:defs.bzl", "rust_bindgen_toolchain") + +native_binary( + name = "clang", + src = "@llvm_toolchain_llvm//:bin/clang", + visibility = ["//tcmalloc_rs:__subpackages__"], +) + +cc_import( + name = "libclang", + shared_library = "@llvm_toolchain_llvm//:libclang", + visibility = ["//tcmalloc_rs:__subpackages__"], +) + +cc_import( + name = "libstdcxx", + static_library = "@llvm_toolchain_llvm//:lib/x86_64-unknown-linux-gnu/libc++.a", + visibility = ["//tcmalloc_rs:__subpackages__"], +) + +rust_bindgen_toolchain( + name = "rust_bindgen_toolchain", + bindgen = "@rules_rust_bindgen//3rdparty:bindgen", + clang = ":clang", + libclang = ":libclang", + libstdcxx = ":libstdcxx", + visibility = ["//tcmalloc_rs:__subpackages__"], +) + +toolchain( + name = "default_bindgen_toolchain", + toolchain = ":rust_bindgen_toolchain", + toolchain_type = "@rules_rust_bindgen//:toolchain_type", + visibility = ["//tcmalloc_rs:__subpackages__"], +) \ No newline at end of file diff --git a/MODULE.bazel b/MODULE.bazel index e6f79d47c..0dd79aa52 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -29,3 +29,38 @@ bazel_dep(name = "fuzztest", version = "20241028.0", dev_dependency = True, repo bazel_dep(name = "google_benchmark", version = "1.9.2", dev_dependency = True, repo_name = "com_github_google_benchmark") bazel_dep(name = "googletest", version = "1.13.0", dev_dependency = True, repo_name = "com_google_googletest") bazel_dep(name = "rules_fuzzing", version = "0.5.2", dev_dependency = True) + +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++17"}, + llvm_version = "17.0.6", +) +use_repo(llvm, "llvm_toolchain", "llvm_toolchain_llvm") + +rust = use_extension("@rules_rust//rust:extensions.bzl", "rust", dev_dependency = True) +rust.toolchain( + edition = "2021", + versions = [ + "1.85.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/docs/quickstart.md b/docs/quickstart.md index cdfc71842..b2d994436 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -254,6 +254,36 @@ build --cxxopt='-std=c++17' Congratulations! You've created your first binary using TCMalloc. +## TCMalloc-rs +This is the rust wrapper for using `TCMalloc` as a global allocator. + +#### Pagesize +By default, `tcmalloc` build with support for `8k` pagesize. + +To build `tcmalloc` with `32k` pagesizes, build `tcmalloc_rs` with `--//tcmalloc_rs:pagesize=32k` + +To build `tcmalloc` with `512k` pagesizes, build `tcmalloc_rs` with `--//tcmalloc_rs:pagesize=512k` +### Running Tests +`bazel run //tcmalloc_rs:test --extra_toolchains=@llvm_toolchain//:cc-toolchain-x86_64-linux` + +### Usage +Add the bazel dependency: +```Starlark +//tcmalloc_rs +``` +To use the allocator, add the following to your code: +```rust +// To set TCMalloc as the global allocator add this to your project: +use tcmalloc_rs; + +#[global_allocator] +static GLOBAL: tcmalloc_rs::TCMalloc = tcmalloc_rs::TCMalloc; +``` + +#### Running the example +Take a look at the code in `tcmalloc_rs/main.rs` and run +`bazel run //tcmalloc_rs:main --extra_toolchains=@llvm_toolchain//:cc-toolchain-x86_64-linux`. + ## What's Next * Read our [overview](overview.md), if you haven't already. The overview diff --git a/tcmalloc/BUILD b/tcmalloc/BUILD index c6b01520d..57e8ab973 100644 --- a/tcmalloc/BUILD +++ b/tcmalloc/BUILD @@ -25,7 +25,11 @@ package(default_visibility = ["//visibility:private"]) licenses(["notice"]) -exports_files(["LICENSE"]) +exports_files([ + "LICENSE", + "tcmalloc.h", + "malloc_extension.h", +]) config_setting( name = "llvm", diff --git a/tcmalloc/selsan/BUILD b/tcmalloc/selsan/BUILD index 4c94450c6..d16b8cb30 100644 --- a/tcmalloc/selsan/BUILD +++ b/tcmalloc/selsan/BUILD @@ -38,7 +38,10 @@ cc_library( "//tcmalloc/internal:logging", "@com_google_absl//absl/base:core_headers", ], - alwayslink = 1, + alwayslink = select({ + "//tcmalloc:unlinked": 0, + "//conditions:default": 1, + }), ) cc_test( diff --git a/tcmalloc_rs/BUILD b/tcmalloc_rs/BUILD new file mode 100644 index 000000000..809f9cd79 --- /dev/null +++ b/tcmalloc_rs/BUILD @@ -0,0 +1,129 @@ +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") + +string_flag( + name = "pagesize", + build_setting_default = "8k", +) + +config_setting( + name = "large", + flag_values = { + ":pagesize": "32k", + }, + visibility = ["//visibility:public"], +) + +config_setting( + name = "xlarge", + flag_values = { + ":pagesize": "256k", + }, + visibility = ["//visibility:public"], +) + +cc_static_library( + name = "tcmalloc", + visibility = ["//visibility:public"], + deps = select({ + ":xlarge": ["//tcmalloc:tcmalloc_256k_pages"], + ":large": ["//tcmalloc:tcmalloc_large_pages"], + "//conditions:default": ["//tcmalloc"], + }), +) + +rust_bindgen_library( + name = "tcmalloc_sys_extensions", + bindgen_flags = [ + "--allowlist-function=.*MallocExtension_.*", + "--allowlist-type=.*MallocExtension_Ownership", + "--allowlist-type=.*MallocExtension_Property", + "--blocklist-function=.*MallocExtension_GetProfileSamplingRate", + "--blocklist-function=.*MallocExtension_SetProfileSamplingRate", + "--blocklist-function=.*MallocExtension_GetGuardedSamplingRate", + "--blocklist-function=.*MallocExtension_SetGuardedSamplingRate", + "--blocklist-function=.*MallocExtension_GetSkipSubreleaseInterval", + "--blocklist-function=.*MallocExtension_SetSkipSubreleaseInterval", + "--opaque-type=std::optional", # prevents unsafe comparisons which could result in overflow + "--opaque-type=std::basic_string", # prevents unsafe comparisons which could result in overflow + "--opaque-type=std::basic_string_view", # allow bindgen to use __BindgenOpaqueArray to repr std::string + "--opaque-type=std::unique_ptr", # prevents unsafe comparisons which could result in overflow + "--use-core", # don't compile with std, allows baremetal envs + ], + cc_lib = "//tcmalloc:malloc_extension", + clang_flags = [ + # "-v", # enable for debugging + "-x", + "c++", + "-std=c++17", + ], + header = "//tcmalloc:malloc_extension.h", + alwayslink = 1, +) + +rust_bindgen_library( + name = "tcmalloc_sys", + bindgen_flags = [ + "--allowlist-function=malloc", + "--allowlist-function=free", + "--allowlist-function=calloc", + "--allowlist-function=realloc", + "--allowlist-function=malloc_usable_size", + # 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 = ":tcmalloc", + clang_flags = [ + # "-v", # enable for debugging + "-x", + "c++", + "-std=c++17", + ], + header = "//tcmalloc:tcmalloc.h", +) + +rust_library( + name = "tcmalloc_rs", + srcs = ["lib.rs"], + visibility = ["//visibility:public"], + deps = [":tcmalloc_sys"], +) + +rust_library( + name = "tcmalloc_extensions_rs", + srcs = ["extension.rs"], + visibility = ["//visibility:public"], + deps = [":tcmalloc_sys_extensions"], +) + +# bazel test //tcmalloc_rs/... --extra_toolchains=@llvm_toolchain//:cc-toolchain-x86_64-linux +rust_test( + name = "test_malloc", + size = "small", + crate = ":tcmalloc_rs", +) + +rust_test( + name = "test_malloc_extension", + size = "small", + crate = ":tcmalloc_extensions_rs", + deps = [":tcmalloc_rs"], +) + +# bazel run //tcmalloc_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 = [ + ":tcmalloc_extensions_rs", + ":tcmalloc_rs", + ], +) diff --git a/tcmalloc_rs/extension.rs b/tcmalloc_rs/extension.rs new file mode 100644 index 000000000..0a8d49fd8 --- /dev/null +++ b/tcmalloc_rs/extension.rs @@ -0,0 +1,240 @@ +#![no_std] +use core; +use tcmalloc_sys_extensions; + +pub use tcmalloc_sys_extensions::tcmalloc_AddressRegionFactory as TCMallocRegionFactory; + +pub use tcmalloc_sys_extensions::absl_Duration_HiRep as AbslHiRep; +pub use tcmalloc_sys_extensions::absl_Duration as AbslDuration; + +pub use tcmalloc_sys_extensions::std_map as StdMap; + +pub use tcmalloc_sys_extensions::tcmalloc_Profile as Profile; + +pub use tcmalloc_sys_extensions::tcmalloc_MallocExtension_AllocationProfilingToken as AllocationProfilingToken; + + +fn convert_string_to_opaque_array(data: &str) -> tcmalloc_sys_extensions::absl_string_view { + tcmalloc_sys_extensions::__BindgenOpaqueArray([data.as_ptr() as u64, data.len() as u64]) +} + +fn convert_opaque_array_to_string(data: tcmalloc_sys_extensions::std_string) -> &'static str { + // Attempt to convert to a byte slice + let byte_slice: &[u8] = unsafe { + // data contains 3 elements + // 0 -> max size allocated + // 1 -> used size + // 2 -> starting address + core::slice::from_raw_parts(data.0[2] as *const u8, data.0[1] as usize) + }; + + // Attempt to convert the byte slice to a UTF-8 string + // unsafe may not be needed, but from_utf8_unchecked is a safety hatch incase a non utf8 character is in the stream + unsafe { core::str::from_utf8_unchecked(byte_slice) } +} + +pub fn get_stats() -> &'static str { + unsafe { convert_opaque_array_to_string(tcmalloc_sys_extensions::tcmalloc_MallocExtension_GetStats()) } +} + +pub fn get_numeric_property(property: &str) -> u8 { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_GetNumericProperty(convert_string_to_opaque_array(property)) } +} + +pub fn mark_thread_idle() { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_MarkThreadIdle() } +} + +pub fn mark_thread_busy() { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_MarkThreadBusy() } +} + +pub fn release_cpu_memory(cpu: i32) -> usize { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_ReleaseCpuMemory(cpu) } +} + +pub fn get_region_factory() -> *mut TCMallocRegionFactory { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_GetRegionFactory() } +} + +pub fn release_memory_to_system(num_bytes: usize) { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_ReleaseMemoryToSystem(num_bytes) } +} + +pub fn get_memory_limit(limit_kind: i32) -> usize { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_GetMemoryLimit(limit_kind) } +} + +pub fn set_memory_limit(limit: usize, limit_kind: i32) { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_SetMemoryLimit(limit, limit_kind) } +} + +pub fn get_profile_sampling_interval() -> i64 { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_GetProfileSamplingInterval() } +} + +pub fn set_profile_sampling_interval(interval: i64) { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_SetProfileSamplingInterval(interval) } +} + +pub fn get_guarded_sampling_interval() -> i64 { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_GetGuardedSamplingInterval() } +} + +pub fn set_guarded_sampling_interval(interval: i64) { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_SetGuardedSamplingInterval(interval) } +} + +pub fn activate_guarded_sampling() { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_ActivateGuardedSampling() } +} + +pub fn per_cpu_caches_active() -> bool { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_PerCpuCachesActive() } +} +pub fn get_max_per_cpu_cache_size() -> i32 { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_GetMaxPerCpuCacheSize() } +} + +pub fn set_max_per_cpu_cache_size(value: i32) { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_SetMaxPerCpuCacheSize(value) } +} + +pub fn get_max_total_thread_cache_bytes() -> i64 { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_GetMaxTotalThreadCacheBytes() } +} + +pub fn set_max_total_thread_cache_bytes(value: i64) { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_SetMaxTotalThreadCacheBytes(value) } +} + +pub fn get_background_process_actions_enabled() -> bool { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_GetBackgroundProcessActionsEnabled() } +} + +pub fn set_background_process_actions_enabled(value: bool) { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_SetBackgroundProcessActionsEnabled(value) } +} + +pub fn get_background_process_sleep_interval() -> AbslDuration { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_GetBackgroundProcessSleepInterval() } +} + +pub fn set_background_process_sleep_interval(value: AbslDuration) { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_SetBackgroundProcessSleepInterval(value) } +} + +pub fn get_skip_subrelease_short_interval() -> AbslDuration { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_GetSkipSubreleaseShortInterval() } +} + +pub fn set_skip_subrelease_short_interval(value: AbslDuration) { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_SetSkipSubreleaseShortInterval(value) } +} + +pub fn get_skip_subrelease_long_interval() -> AbslDuration { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_GetSkipSubreleaseLongInterval() } +} + +pub fn set_skip_subrelease_long_interval(value: AbslDuration) { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_SetSkipSubreleaseLongInterval(value) } +} + +pub fn get_cache_demand_based_release() -> bool { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_GetCacheDemandBasedRelease() } +} + +pub fn set_cache_demand_based_release(value: bool) { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_SetCacheDemandBasedRelease(value) } +} + +pub fn get_cache_demand_release_short_interval() -> AbslDuration { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_GetCacheDemandReleaseShortInterval() } +} + +pub fn set_cache_demand_release_short_interval(value: AbslDuration) { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_SetCacheDemandReleaseShortInterval(value) } +} + +pub fn get_cache_demand_release_long_interval() -> AbslDuration { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_GetCacheDemandReleaseLongInterval() } +} + +pub fn set_cache_demand_release_long_interval(value: AbslDuration) { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_SetCacheDemandReleaseLongInterval(value) } +} + +pub fn get_estimated_allocated_size(size: usize) -> usize { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_GetEstimatedAllocatedSize(size) } +} + +pub fn get_estimated_allocated_size_hot_cold(size: usize, hot_cold: u8) -> usize { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_GetEstimatedAllocatedSize1(size, hot_cold) } +} + +pub fn get_allocated_size(p: *const ::core::ffi::c_void) -> u8 { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_GetAllocatedSize(p) } +} + +pub fn get_ownership(p: *const ::core::ffi::c_void) -> i32 { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_GetOwnership(p) } +} + +pub fn get_properties() -> StdMap { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_GetProperties() } +} + +pub fn snapshot_current(r#type: i32) -> Profile { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_SnapshotCurrent(r#type) } +} + +pub fn start_allocation_profiling() -> AllocationProfilingToken { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_StartAllocationProfiling() } +} + +pub fn start_lifetime_profiling() -> AllocationProfilingToken { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_StartLifetimeProfiling() } +} + +pub fn process_background_actions() { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_ProcessBackgroundActions() } +} + +pub fn needs_process_background_actions() -> bool { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_NeedsProcessBackgroundActions() } +} + +pub fn get_background_release_rate() -> usize { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_GetBackgroundReleaseRate() } +} + +pub fn set_background_release_rate(rate: usize) { + unsafe { tcmalloc_sys_extensions::tcmalloc_MallocExtension_SetBackgroundReleaseRate(rate) } +} + +#[cfg(test)] +mod tests { + use super::{TCMallocRegionFactory, get_numeric_property, get_region_factory, get_stats}; + use tcmalloc_rs; + + #[global_allocator] + static GLOBAL: tcmalloc_rs::TCMalloc = tcmalloc_rs::TCMalloc; + + #[test] + fn test_get_stats() { + let stats = get_stats(); + assert!(!stats.is_empty()); + } + + #[test] + fn test_get_numeric_property() { + let value = get_numeric_property("tcmalloc.current_total_thread_cache_bytes"); + assert_ne!(0, value); + } + + #[test] + fn test_get_region_factory_internal_byts_allocated() { + let _value = get_region_factory(); + assert_ne!(0, unsafe { TCMallocRegionFactory::InternalBytesAllocated() }); + } +} \ No newline at end of file diff --git a/tcmalloc_rs/lib.rs b/tcmalloc_rs/lib.rs new file mode 100644 index 000000000..44de90457 --- /dev/null +++ b/tcmalloc_rs/lib.rs @@ -0,0 +1,205 @@ +#![no_std] +//! `tcmalloc_rs` provides bindings for [`google/tcmalloc`](https://github.com/google/tcmalloc) to make it usable as a global allocator for rust. +//! +//! To use `tcmalloc_rs` add it as a bazel dependency: +//! ```Starlark +//! //tcmalloc_rs +//! ``` +//! +//! To set `TCMalloc` as the global allocator add this to your project: +//! ```rust +//! use tcmalloc_rs; +//! +//! #[global_allocator] +//! static GLOBAL: tcmalloc_rs::TCMalloc = tcmalloc_rs::TCMalloc; +//! ``` + + +use core::{ + alloc::{GlobalAlloc, Layout}, + ptr::NonNull, +}; +use tcmalloc_sys; + +#[derive(Debug, Copy, Clone)] +#[repr(C)] +pub struct TCMalloc; + +unsafe impl Send for TCMalloc {} +unsafe impl Sync for TCMalloc {} + +impl TCMalloc { + #[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 { tcmalloc_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), + _ => NonNull::new(unsafe { tcmalloc_sys::malloc(layout.size()) }.cast()) + } + } +} + +unsafe impl GlobalAlloc for TCMalloc { + /// 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 constrains are not full-filled. + #[inline(always)] + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + match layout.size() { + 0 => layout.align() as *mut u8, + _ => tcmalloc_sys::malloc(layout.size()).cast() + } + } + + /// 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 constrains are not full-filled. + #[inline(always)] + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + if layout.size() != 0 { + tcmalloc_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 => tcmalloc_sys::calloc(layout.align(), size).cast() + } + } + + /// 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 remains 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 constrains 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())) + } + _ => tcmalloc_sys::realloc(ptr.cast(), new_size).cast() + } + } +} + +#[cfg(test)] +mod tests { + use super::TCMalloc; + use core::alloc::{GlobalAlloc, Layout}; + #[test] + fn allocation_lifecycle() { + let alloc = TCMalloc::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 = TCMalloc; + + 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 = TCMalloc; + + 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 = TCMalloc; + + 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 = TCMalloc; + + let ptr = alloc.alloc(layout); + alloc.dealloc(ptr, layout); + } + } + + #[test] + fn test_usable_size() { + let alloc = TCMalloc::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); + } + } +} \ No newline at end of file diff --git a/tcmalloc_rs/main.rs b/tcmalloc_rs/main.rs new file mode 100644 index 000000000..952c64798 --- /dev/null +++ b/tcmalloc_rs/main.rs @@ -0,0 +1,17 @@ +use tcmalloc_rs; +use tcmalloc_extensions_rs::get_stats; + +#[global_allocator] +static GLOBAL: tcmalloc_rs::TCMalloc = tcmalloc_rs::TCMalloc; + +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 stats = get_stats(); + println!("{}", stats) +} \ No newline at end of file