Add "override_allocator_on_supported_platforms" feature flag

Signed-off-by: Mads Marquart <[email protected]>

Author: madsmtm

ci/run.sh

@@ -32,7 +32,7 @@ cargo test --target "${TARGET}" --features stats
 cargo test --target "${TARGET}" --features 'debug profiling'
 
 cargo test --target "${TARGET}" \
-    --features unprefixed_malloc_on_supported_platforms
+    --features override_allocator_on_supported_platforms
 cargo test --target "${TARGET}" --no-default-features
 cargo test --target "${TARGET}" --no-default-features \
     --features background_threads_runtime_support
@@ -48,7 +48,7 @@ cargo test --target "${TARGET}" --release
 cargo test --target "${TARGET}" --manifest-path jemalloc-sys/Cargo.toml
 cargo test --target "${TARGET}" \
              --manifest-path jemalloc-sys/Cargo.toml \
-             --features unprefixed_malloc_on_supported_platforms
+             --features override_allocator_on_supported_platforms
 
 # FIXME: jemalloc-ctl fails in the following targets
 case "${TARGET}" in
@@ -78,13 +78,13 @@ cargo test --target "${TARGET}" \
 #     ${CARGO_CMD} test --target "${TARGET}" --features alloc_trait
 # fi
 
-# Test that unprefixed_malloc_on_supported_platforms works in dylibs.
+# Test that overriding works in dylibs.
 case "$TARGET" in
     "i686-unknown-linux-musl") ;;
     "x86_64-unknown-linux-musl") ;;
     *)
         cargo run --target "${TARGET}" \
             -p test-dylib \
-            --features unprefixed_malloc_on_supported_platforms
+            --features override_allocator_on_supported_platforms
         ;;
 esac

jemalloc-sys/Cargo.toml

@@ -46,6 +46,7 @@ background_threads_runtime_support = []
 background_threads = [ "background_threads_runtime_support" ]
 stats = []
 unprefixed_malloc_on_supported_platforms = []
+override_allocator_on_supported_platforms = [ "unprefixed_malloc_on_supported_platforms" ]
 disable_initial_exec_tls = []
 disable_cache_oblivious = []
 

jemalloc-sys/README.md

@@ -76,6 +76,14 @@ This crate provides following cargo feature flags:
   program to use `jemalloc` as well. On some platforms prefixes are always used
   because unprefixing is known to cause segfaults due to allocator mismatches.
 
+* `override_allocator_on_supported_platforms`: override the system allocator,
+  even outside Rust code.
+
+  This enables the `unprefixed_malloc_on_supported_platforms` feature, with the
+  addition that it forces overriding the allocator even if `malloc` and `free`
+  would not usually have been seen by the linker. It also overrides the
+  allocator on Apple platforms.
+
   Note that to use this, the `jemalloc-sys` crate must actually be visible to
   `rustc` (it is not enough to only declare it in `Cargo.toml`). This can be
   done by adding:

jemalloc-sys/build.rs

@@ -114,7 +114,8 @@ fn main() {
         // Apple targets don't support unprefixed, but they do support
         // overriding (if you do the `zone_register` trick), so no need to
         // warn there.
-        if !target.contains("apple") {
+        let override_ = env::var("CARGO_FEATURE_OVERRIDE_ALLOCATOR_ON_SUPPORTED_PLATFORMS").is_ok();
+        if !target.contains("apple") || !override_ {
             warning!(
                 "Unprefixed `malloc` requested on unsupported platform `{}` => using prefixed `malloc`",
                 target

jemalloc-sys/src/lib.rs

@@ -891,7 +891,7 @@ mod env;
 
 pub use env::*;
 
-// When using the `"unprefixed_malloc_on_supported_platforms"` feature flag,
+// When using the `"override_allocator_on_supported_platforms"` feature flag,
 // the user wants us to globally override the system allocator.
 //
 // However, since we build `jemalloc` as a static library (an archive), the
@@ -930,7 +930,7 @@ pub use env::*;
 // how this works:
 // <https://github.com/rust-lang/rust/pull/95604>
 
-#[cfg(not(prefixed))]
+#[cfg(all(feature = "override_allocator_on_supported_platforms", not(target_vendor = "apple")))]
 mod set_up_statics {
     use super::*;
 
@@ -957,10 +957,7 @@ mod set_up_statics {
 // as "used" when defined in an archive member in a static library, so we need
 // to explicitly reference the function via. Rust's `#[used]`.
 
-#[cfg(all(
-    feature = "unprefixed_malloc_on_supported_platforms",
-    target_vendor = "apple"
-))]
+#[cfg(all(feature = "override_allocator_on_supported_platforms", target_vendor = "apple"))]
 #[used]
 static USED_ZONE_REGISTER: unsafe extern "C" fn() = {
     extern "C" {

jemalloc-sys/tests/unprefixed_malloc.rs

@@ -10,13 +10,7 @@ fn malloc_is_overridden() {
     assert_eq!(tikv_jemalloc_sys::malloc as usize, libc::malloc as usize)
 }
 
-#[cfg(any(
-    not(prefixed),
-    all(
-        feature = "unprefixed_malloc_on_supported_platforms",
-        target_vendor = "apple"
-    ),
-))]
+#[cfg(any(not(prefixed), all(feature = "override_allocator_on_supported_platforms", target_vendor = "apple"),))]
 #[test]
 fn malloc_and_libc_are_interoperable_when_overridden() {
     let ptr = unsafe { tikv_jemalloc_sys::malloc(42) };

jemallocator/Cargo.toml

@@ -53,6 +53,7 @@ stats = ["tikv-jemalloc-sys/stats"]
 background_threads_runtime_support = ["tikv-jemalloc-sys/background_threads_runtime_support"]
 background_threads = ["tikv-jemalloc-sys/background_threads"]
 unprefixed_malloc_on_supported_platforms = ["tikv-jemalloc-sys/unprefixed_malloc_on_supported_platforms"]
+override_allocator_on_supported_platforms = ["tikv-jemalloc-sys/override_allocator_on_supported_platforms"]
 disable_initial_exec_tls = ["tikv-jemalloc-sys/disable_initial_exec_tls"]
 disable_cache_oblivious = ["tikv-jemalloc-sys/disable_cache_oblivious"]
 

test-dylib/Cargo.toml

@@ -14,8 +14,8 @@ tikv-jemalloc-sys = { path = "../jemalloc-sys" }
 cc = "^1.0.13"
 
 [features]
-unprefixed_malloc_on_supported_platforms = [
-    "tikv-jemalloc-sys/unprefixed_malloc_on_supported_platforms",
+override_allocator_on_supported_platforms = [
+    "tikv-jemalloc-sys/override_allocator_on_supported_platforms",
 ]
 
 [[bin]]

test-dylib/src/main.rs

@@ -32,7 +32,7 @@ fn main() {
 
     // If overidden, test that the same is true for `tikv_jemalloc_sys`'
     // symbols being interoperable with `free`.
-    if cfg!(feature = "unprefixed_malloc_on_supported_platforms") {
+    if cfg!(feature = "override_allocator_on_supported_platforms") {
         let ptr = unsafe { tikv_jemalloc_sys::malloc(10) };
         unsafe { dep_free(ptr) };
         let ptr = unsafe { tikv_jemalloc_sys::malloc(10) };