Use jemalloc as the global allocator

Using jemalloc provides a 20% max RSS reduction
and 15% speed up in total indexing + resolution
time

Author: vinistock

ext/rubydex/extconf.rb

@@ -78,7 +78,7 @@
   ENV["RUSTFLAGS"] = "-Zsanitizer=#{ENV["SANITIZER"]}"
   "cargo +nightly build -Zbuild-std #{cargo_args.join(" ")}".strip
 else
-  "cargo build #{cargo_args.join(" ")}".strip
+  "cargo build --features rubydex/jemalloc_dylib #{cargo_args.join(" ")}".strip
 end
 
 lib_dir = gem_dir.join("lib").join("rubydex")

rust/rubydex/Cargo.toml

@@ -19,6 +19,11 @@ path = "src/main.rs"
 crate-type = ["rlib"]
 
 [features]
+jemalloc = ["dep:tikv-jemallocator"]
+# The Ruby native extension uses a dlopen to load the rubydex shared object. In these scenarios, jemalloc must be linked
+# dynamically or else the process crashes. This has a small performance cost, which we don't need to pay when using the
+# Rust crate by itself, so we gate it with a feature
+jemalloc_dylib = ["jemalloc", "tikv-jemallocator/disable_initial_exec_tls"]
 test_utils = ["dep:tempfile"]
 
 [dependencies]
@@ -37,8 +42,11 @@ crossbeam-deque = "0.8"
 crossbeam-utils = "0.8"
 crossbeam-channel = "0.5"
 
+[target.'cfg(not(target_os = "windows"))'.dependencies]
+tikv-jemallocator = { version = "0.7.0", optional = true }
+
 [dev-dependencies]
-rubydex = { path = ".", features = ["test_utils"] }
+rubydex = { path = ".", features = ["test_utils", "jemalloc"] }
 tempfile = "3.0"
 assert_cmd = "2.0"
 predicates = "3.1"

rust/rubydex/benches/graph_memory.rs

@@ -1,6 +1,4 @@
-//! Jemalloc does not compile on Windows, so this memory benchmark executable is a no-op
-
-#[cfg(not(target_os = "windows"))]
+#[cfg(all(feature = "jemalloc", not(target_os = "windows")))]
 mod imp {
     use std::collections::HashSet;
 
@@ -12,10 +10,6 @@ mod imp {
     };
     use tikv_jemalloc_ctl::{epoch, stats};
 
-    // Substitute the global allocator by jemalloc, so that we can track all allocations and measure the graph usage
-    #[global_allocator]
-    static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
-
     /// Advance the jemalloc epoch (stats are cached between epochs) and read the
     /// number of bytes currently allocated by the application.
     fn allocated_bytes() -> usize {
@@ -47,6 +41,6 @@ mod imp {
 }
 
 fn main() {
-    #[cfg(not(target_os = "windows"))]
+    #[cfg(all(feature = "jemalloc", not(target_os = "windows")))]
     imp::run();
 }

rust/rubydex/src/lib.rs

@@ -1,3 +1,9 @@
+// Setting the global allocator needs to happen during linking and consumers may not always want to change the global
+// allocator. We gate the usage of jemalloc with a cargo feature, so that consumers can decide if they want to use it
+#[cfg(all(feature = "jemalloc", not(target_os = "windows")))]
+#[global_allocator]
+static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
 pub mod compile_assertions;
 pub mod config;
 pub mod diagnostic;