Make unprefixed consistently override the system allocator

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

Author: madsmtm

jemalloc-sys/README.md

@@ -48,11 +48,11 @@ This crate provides following cargo feature flags:
 * `stats` (configure `jemalloc` with `--enable-stats`): Enable statistics
   gathering functionality. See the `jemalloc`'s "`opt.stats_print`" option
   documentation for usage details.
-  
+
 * `debug` (configure `jemalloc` with `--enable-debug`): Enable assertions and
   validation code. This incurs a substantial performance hit, but is very useful
   during application development.
-  
+
 * `background_threads_runtime_support` (enabled by default): enables
   background-threads run-time support when building `jemalloc-sys` on some POSIX
   targets supported by `jemalloc`. Background threads are disabled at run-time
@@ -72,16 +72,36 @@ This crate provides following cargo feature flags:
 * `unprefixed_malloc_on_supported_platforms`: when disabled, configure
   `jemalloc` with `--with-jemalloc-prefix=_rjem_`. Enabling this causes symbols
   like `malloc` to be emitted without a prefix, overriding the ones defined by
-  libc. This usually causes C and C++ code linked in the same program to use
-  `jemalloc` as well. On some platforms prefixes are always used because
-  unprefixing is known to cause segfaults due to allocator mismatches.
-  
+  libc. This usually causes C, Objective-C and C++ code linked in the same
+  program to use `jemalloc` as well. On some platforms prefixes are always used
+  because unprefixing is known to cause segfaults due to allocator mismatches.
+
+  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 with:
+  ```rust
+  use jemalloc_sys as _;
+  ```
+
+  On macOS, you'll need to do a bit more to work around a bug in `rustc`:
+  ```rust
+  // Workaround for https://github.com/rust-lang/rust/issues/133491
+  #[cfg(target_vendor = "apple")]
+  #[used]
+  static USED_ZONE_REGISTER: unsafe extern "C" fn() = {
+      extern "C" {
+          fn _rjem_je_zone_register();
+      }
+      _rjem_je_zone_register
+  };
+  ```
+
 * `disable_initial_exec_tls` (disabled by default): when enabled, jemalloc is
-  built with the `--disable-initial-exec-tls` option. It disables the 
-  initial-exec TLS model for jemalloc's internal thread-local storage (on those 
-  platforms that support explicit settings). This can allow jemalloc to be 
+  built with the `--disable-initial-exec-tls` option. It disables the
+  initial-exec TLS model for jemalloc's internal thread-local storage (on those
+  platforms that support explicit settings). This can allow jemalloc to be
   dynamically loaded after program startup (e.g. using dlopen). If you encounter
-  the error `yourlib.so: cannot allocate memory in static TLS block`, you'll 
+  the error `yourlib.so: cannot allocate memory in static TLS block`, you'll
   likely want to enable this.
 
 * `disable_cache_oblivious` (disabled by default): when enabled, jemalloc is
@@ -104,7 +124,7 @@ hyphens `-` are replaced with underscores `_`(see
   variable, the `/etc/malloc.conf` symlink, and the `MALLOC_CONF` environment
   variable (note: this variable might be prefixed as `_RJEM_MALLOC_CONF`). For
   example, to change the default decay time for dirty pages to 30 seconds:
-  
+
   ```
   JEMALLOC_SYS_WITH_MALLOC_CONF=dirty_decay_ms:30000
   ```
@@ -115,17 +135,17 @@ hyphens `-` are replaced with underscores `_`(see
   allocator page size equal to the system page size, so this option need not be
   specified unless the system page size may change between configuration and
   execution, e.g. when cross compiling.
-  
+
 * `JEMALLOC_SYS_WITH_LG_HUGEPAGE=<lg-hugepage>`: Specify the base 2 log of the
   system huge page size. This option is useful when cross compiling, or when
   overriding the default for systems that do not explicitly support huge pages.
-  
-  
+
+
 * `JEMALLOC_SYS_WITH_LG_QUANTUM=<lg-quantum>`: Specify the base 2 log of the
   minimum allocation alignment. jemalloc needs to know the minimum alignment
   that meets the following C standard requirement (quoted from the April 12,
   2011 draft of the C11 standard):
-  
+
   > The pointer returned if the allocation succeeds is suitably aligned so that
   > it may be assigned to a pointer to any type of object with a fundamental
   > alignment requirement and then used to access such an object or an array of

jemalloc-sys/build.rs

@@ -111,10 +111,15 @@ fn main() {
             .iter()
             .any(|i| target.contains(i))
     {
-        warning!(
-            "Unprefixed `malloc` requested on unsupported platform `{}` => using prefixed `malloc`",
-            target
-        );
+        // 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") {
+            warning!(
+                "Unprefixed `malloc` requested on unsupported platform `{}` => using prefixed `malloc`",
+                target
+            );
+        }
         use_prefix = true;
     }
 

jemalloc-sys/src/env.rs

@@ -21,4 +21,4 @@ pub static NO_BG_THREAD_TARGETS: &[&str] = &["musl"];
 // https://github.com/rust-lang/rust/commit/e3b414d8612314e74e2b0ebde1ed5c6997d28e8d
 // https://github.com/rust-lang/rust/commit/9f3de647326fbe50e0e283b9018ab7c41abccde3
 // https://github.com/rust-lang/rust/commit/ed015456a114ae907a36af80c06f81ea93182a24
-pub static NO_UNPREFIXED_MALLOC_TARGETS: &[&str] = &["android", "dragonfly", "darwin"];
+pub static NO_UNPREFIXED_MALLOC_TARGETS: &[&str] = &["android", "dragonfly", "apple"];

jemalloc-sys/src/lib.rs

@@ -890,3 +890,58 @@ pub type extent_merge_t = unsafe extern "C" fn(
 mod env;
 
 pub use env::*;
+
+// When using the `"unprefixed_malloc_on_supported_platforms"` feature flag to
+// override the system allocator, make sure that all allocator functions are
+// visible to the linker, such that it will override all of them.
+//
+// For example, this would fail horribly if one used a dynamic library that
+// calls `malloc`, and then later `free`'d that in Rust code that used
+// jemalloc-sys, since then the linker might think that only `free` from
+// jemalloc is needed, and wouldn't override `malloc`.
+#[cfg(not(prefixed))]
+mod set_up_statics {
+    use super::*;
+
+    #[used]
+    static USED_MALLOC: unsafe extern "C" fn(usize) -> *mut c_void = malloc;
+    #[used]
+    static USED_CALLOC: unsafe extern "C" fn(usize, usize) -> *mut c_void = calloc;
+    #[used]
+    static USED_POSIX_MEMALIGN: unsafe extern "C" fn(*mut *mut c_void, usize, usize) -> c_int =
+        posix_memalign;
+    #[used]
+    static USED_ALLOC: unsafe extern "C" fn(usize, usize) -> *mut c_void = aligned_alloc;
+    #[used]
+    static USED_REALLOC: unsafe extern "C" fn(*mut c_void, usize) -> *mut c_void = realloc;
+    #[used]
+    static USED_FREE: unsafe extern "C" fn(*mut c_void) = free;
+}
+
+// On macOS, jemalloc doesn't directly override malloc/free, but instead
+// registers itself with the allocator's zone APIs in a ctor. However, ld64
+// doesn't consider ctors as "used" when defined in an object file / archive
+// member in a static library, so we need to explicitly depend on the function
+// using Rust's `#[used]`.
+//
+// NOTE: `#[used]` currently doesn't actually work on macOS, but that's an
+// upstream rustc issue, see https://github.com/rust-lang/rust/issues/133491.
+#[cfg(all(
+    feature = "unprefixed_malloc_on_supported_platforms",
+    target_vendor = "apple"
+))]
+#[used]
+static USED_ZONE_REGISTER: unsafe extern "C" fn() = {
+    extern "C" {
+        #[cfg_attr(prefixed, link_name = "_rjem_je_zone_register")]
+        #[cfg_attr(not(prefixed), link_name = "je_zone_register")]
+        fn zone_register();
+    }
+
+    // NOTE: This internal function is marked as __attribute__((constructor)),
+    // and thus we neither can nor should call it ourselves.
+    //
+    // Instead, we simply need to reference it in a `#[used]` static to get
+    // the linker to see that it should be registering it as a constructor.
+    zone_register
+};