Bump jemalloc.

Signed-off-by: Jason Volk <jason@zemos.net>

Cargo.lock

@@ -2180,6 +2180,26 @@ version = "1.0.15"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c"
 
+[[package]]
+name = "jevmalloc"
+version = "0.0.0+5.3.0-1-g54eaed1d8b56b1aa528be3bdd1877e59c56fa90c"
+source = "git+https://github.com/matrix-construct/jevmalloc?rev=29fde02533ee2c6a5e9b87095324e6f7a5d43738#29fde02533ee2c6a5e9b87095324e6f7a5d43738"
+dependencies = [
+ "jevmalloc-sys",
+ "libc",
+ "paste",
+]
+
+[[package]]
+name = "jevmalloc-sys"
+version = "0.0.0+5.3.0-1-g54eaed1d8b56b1aa528be3bdd1877e59c56fa90c"
+source = "git+https://github.com/matrix-construct/jevmalloc?rev=29fde02533ee2c6a5e9b87095324e6f7a5d43738#29fde02533ee2c6a5e9b87095324e6f7a5d43738"
+dependencies = [
+ "cc",
+ "libc",
+ "rustflags",
+]
+
 [[package]]
 name = "jobserver"
 version = "0.1.34"
@@ -3851,24 +3871,24 @@ dependencies = [
 [[package]]
 name = "rust-librocksdb-sys"
 version = "0.40.0+10.7.5"
-source = "git+https://github.com/matrix-construct/rust-rocksdb?rev=2c0ad2ac6d933f61afb4c79ee2f965ffe2293774#2c0ad2ac6d933f61afb4c79ee2f965ffe2293774"
+source = "git+https://github.com/matrix-construct/rust-rocksdb?rev=31e6f8dd4a531b379c132e82773efc9315b0489a#31e6f8dd4a531b379c132e82773efc9315b0489a"
 dependencies = [
  "bindgen",
  "bzip2-sys",
  "cc",
  "glob",
+ "jevmalloc-sys",
  "libc",
  "libz-sys",
  "lz4-sys",
  "pkg-config",
- "tikv-jemalloc-sys",
  "zstd-sys",
 ]
 
 [[package]]
 name = "rust-rocksdb"
 version = "0.44.2"
-source = "git+https://github.com/matrix-construct/rust-rocksdb?rev=2c0ad2ac6d933f61afb4c79ee2f965ffe2293774#2c0ad2ac6d933f61afb4c79ee2f965ffe2293774"
+source = "git+https://github.com/matrix-construct/rust-rocksdb?rev=31e6f8dd4a531b379c132e82773efc9315b0489a#31e6f8dd4a531b379c132e82773efc9315b0489a"
 dependencies = [
  "libc",
  "parking_lot",
@@ -3897,6 +3917,12 @@ dependencies = [
  "semver",
 ]
 
+[[package]]
+name = "rustflags"
+version = "0.1.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a39e0e9135d7a7208ee80aa4e3e4b88f0f5ad7be92153ed70686c38a03db2e63"
+
 [[package]]
 name = "rusticata-macros"
 version = "4.1.0"
@@ -4679,34 +4705,6 @@ dependencies = [
  "zune-jpeg 0.4.21",
 ]
 
-[[package]]
-name = "tikv-jemalloc-ctl"
-version = "0.6.0"
-source = "git+https://github.com/matrix-construct/jemallocator?rev=03bed96afbbc898bef4d4f7d335c0519e3d1afad#03bed96afbbc898bef4d4f7d335c0519e3d1afad"
-dependencies = [
- "libc",
- "paste",
- "tikv-jemalloc-sys",
-]
-
-[[package]]
-name = "tikv-jemalloc-sys"
-version = "0.6.0+5.3.0-1-ge13ca993e8ccb9ba9847cc330696e02839f328f7"
-source = "git+https://github.com/matrix-construct/jemallocator?rev=03bed96afbbc898bef4d4f7d335c0519e3d1afad#03bed96afbbc898bef4d4f7d335c0519e3d1afad"
-dependencies = [
- "cc",
- "libc",
-]
-
-[[package]]
-name = "tikv-jemallocator"
-version = "0.6.0"
-source = "git+https://github.com/matrix-construct/jemallocator?rev=03bed96afbbc898bef4d4f7d335c0519e3d1afad#03bed96afbbc898bef4d4f7d335c0519e3d1afad"
-dependencies = [
- "libc",
- "tikv-jemalloc-sys",
-]
-
 [[package]]
 name = "time"
 version = "0.3.44"
@@ -5202,6 +5200,7 @@ dependencies = [
  "insta",
  "ipaddress",
  "itertools 0.14.0",
+ "jevmalloc",
  "jsonwebtoken",
  "libc",
  "libloading",
@@ -5223,9 +5222,6 @@ dependencies = [
  "smallstr",
  "smallvec",
  "thiserror 2.0.17",
- "tikv-jemalloc-ctl",
- "tikv-jemalloc-sys",
- "tikv-jemallocator",
  "tokio",
  "tokio-metrics",
  "toml 0.9.8",

Cargo.toml

@@ -231,6 +231,10 @@ version = "0.1"
 [workspace.dependencies.itertools]
 version = "0.14"
 
+[workspace.dependencies.jevmalloc]
+git = "https://github.com/matrix-construct/jevmalloc"
+rev = "29fde02533ee2c6a5e9b87095324e6f7a5d43738"
+
 [workspace.dependencies.jsonwebtoken]
 version = "9.3"
 default-features = false
@@ -368,7 +372,7 @@ default-features = false
 
 [workspace.dependencies.rust-rocksdb]
 git = "https://github.com/matrix-construct/rust-rocksdb"
-rev = "2c0ad2ac6d933f61afb4c79ee2f965ffe2293774"
+rev = "31e6f8dd4a531b379c132e82773efc9315b0489a"
 default-features = false
 features = [
 	"bzip2",
@@ -477,30 +481,6 @@ default-features = false
 version = "2.0"
 default-features = false
 
-[workspace.dependencies.tikv-jemallocator]
-git = "https://github.com/matrix-construct/jemallocator"
-rev = "03bed96afbbc898bef4d4f7d335c0519e3d1afad"
-default-features = false
-features = [
-	"background_threads_runtime_support",
-	"unprefixed_malloc_on_supported_platforms",
-]
-
-[workspace.dependencies.tikv-jemalloc-ctl]
-git = "https://github.com/matrix-construct/jemallocator"
-rev = "03bed96afbbc898bef4d4f7d335c0519e3d1afad"
-default-features = false
-features = ["use_std"]
-
-[workspace.dependencies.tikv-jemalloc-sys]
-git = "https://github.com/matrix-construct/jemallocator"
-rev = "03bed96afbbc898bef4d4f7d335c0519e3d1afad"
-default-features = false
-features = [
-	"background_threads_runtime_support",
-	"unprefixed_malloc_on_supported_platforms",
-]
-
 [workspace.dependencies.tokio]
 version = "1.48"
 default-features = false

src/core/Cargo.toml

@@ -26,18 +26,14 @@ gzip_compression = [
 	"reqwest/gzip",
 ]
 jemalloc = [
-	"dep:tikv-jemalloc-sys",
+	"dep:jevmalloc",
-	"dep:tikv-jemalloc-ctl",
-	"dep:tikv-jemallocator",
 ]
 jemalloc_conf = []
 jemalloc_prof = [
-	#"tikv-jemalloc-sys/profiling",
+	#"jevmalloc/profiling",
 ]
 jemalloc_stats = [
-    #"tikv-jemalloc-sys/stats",
+	#"jevmalloc/stats",
-    #"tikv-jemalloc-ctl/stats",
-    #"tikv-jemallocator/stats",
 ]
 perf_measurements = []
 release_max_log_level = [
@@ -96,12 +92,8 @@ serde.workspace = true
 smallvec.workspace = true
 smallstr.workspace = true
 thiserror.workspace = true
-tikv-jemallocator.optional = true
+jevmalloc.optional = true
-tikv-jemallocator.workspace = true
+jevmalloc.workspace = true
-tikv-jemalloc-ctl.optional = true
-tikv-jemalloc-ctl.workspace = true
-tikv-jemalloc-sys.optional = true
-tikv-jemalloc-sys.workspace = true
 tokio.workspace = true
 tokio-metrics.workspace = true
 toml.workspace = true

src/core/alloc/je.rs

@@ -1,20 +1,26 @@
 //! jemalloc allocator
 
 use std::{
+	alloc::Layout,
 	cell::OnceCell,
 	ffi::{CStr, c_char, c_void},
 	fmt::Debug,
-	sync::RwLock,
+	panic::catch_unwind,
+	process::abort,
+	sync::{
+		Mutex,
+		atomic::{AtomicBool, AtomicU64, Ordering},
+	},
 };
 
-use arrayvec::ArrayVec;
+use jevmalloc as jemalloc;
-use tikv_jemalloc_ctl as mallctl;
+use jevmalloc::{ctl as mallctl, ffi};
-use tikv_jemalloc_sys as ffi;
-use tikv_jemallocator as jemalloc;
 
 use crate::{
-	Result, err, is_equal_to, is_nonzero,
+	Result,
-	utils::{math, math::Tried},
+	arrayvec::ArrayVec,
+	err, is_equal_to, is_nonzero,
+	utils::{BoolExt, math, math::Tried},
 };
 
 #[cfg(feature = "jemalloc_conf")]
@@ -45,20 +51,28 @@ const _MALLOC_CONF_PROF: &str = ",prof_active:false";
 ))]
 const _MALLOC_CONF_PROF: &str = "";
 
-#[global_allocator]
-static JEMALLOC: jemalloc::Jemalloc = jemalloc::Jemalloc;
-static CONTROL: RwLock<()> = RwLock::new(());
-
 type Name = ArrayVec<u8, NAME_MAX>;
 type Key = ArrayVec<usize, KEY_SEGS>;
 
 const NAME_MAX: usize = 128;
 const KEY_SEGS: usize = 8;
 
+#[global_allocator]
+static JEMALLOC: jemalloc::Jemalloc = jemalloc::Jemalloc;
+static CONTROL: Mutex<()> = Mutex::new(());
+
+static GLOBAL_ALLOCS: AtomicU64 = AtomicU64::new(0);
+static COUNT_GLOBAL_ALLOCS: AtomicBool = AtomicBool::new(false);
+static TRACE_GLOBAL_ALLOCS: AtomicBool = AtomicBool::new(false);
+
 #[crate::ctor]
 fn _static_initialization() {
-	acq_epoch().expect("pre-initialization of jemalloc failed");
+	// SAFETY: Mutable static globals in jemalloc crate; must be initialized
-	acq_epoch().expect("pre-initialization of jemalloc failed");
+	// properly and uniquely.
+	unsafe {
+		jemalloc::hook::ALLOC = Some(global_alloc_hook);
+		jemalloc::hook::ALLOC_ZEROED = Some(global_alloc_zeroed_hook);
+	};
 }
 
 #[must_use]
@@ -139,6 +153,43 @@ fn handle_malloc_stats(opaque: *mut c_void, msg: *const c_char) {
 	res.push_str(msg.as_ref());
 }
 
+fn global_alloc_hook(layout: Layout) {
+	catch_unwind(move || handle_global_alloc(layout))
+		.map_err(|_| abort())
+		.ok();
+}
+
+fn global_alloc_zeroed_hook(layout: Layout) {
+	catch_unwind(move || handle_global_alloc(layout))
+		.map_err(|_| abort())
+		.ok();
+}
+
+fn handle_global_alloc(layout: Layout) {
+	use std::io::Write;
+
+	use libc::{STDOUT_FILENO, write};
+
+	let do_count = COUNT_GLOBAL_ALLOCS.load(Ordering::Relaxed);
+	let count = GLOBAL_ALLOCS.fetch_add(do_count.into(), Ordering::Relaxed);
+
+	if TRACE_GLOBAL_ALLOCS.load(Ordering::Relaxed) {
+		let mut buf = ArrayVec::<u8, 128>::new();
+		writeln!(&mut buf, "{count} align={} size={}", layout.align(), layout.size())
+			.expect("writeln! to buffer failed");
+
+		// SAFETY: Valid ptr and len from buf for writing to stdout.
+		unsafe { write(STDOUT_FILENO, buf.as_ptr().cast::<c_void>(), buf.len()) }
+			.ge(&0)
+			.into_result()
+			.expect("write(2) error");
+	}
+}
+
+#[inline]
+#[must_use]
+pub fn global_alloc_count() -> u64 { GLOBAL_ALLOCS.load(Ordering::Relaxed) }
+
 macro_rules! mallctl {
 	($name:expr_2021) => {{
 		thread_local! {
@@ -346,7 +397,7 @@ fn set<T>(key: &Key, val: T) -> Result<T>
 where
 	T: Copy + Debug,
 {
-	let _lock = CONTROL.write()?;
+	let _lock = CONTROL.lock()?;
 	let res = xchg(key, val)?;
 	inc_epoch()?;
 
@@ -364,7 +415,6 @@ where
 	T: Copy + Debug,
 {
 	acq_epoch()?;
-	acq_epoch()?;
 
 	// SAFETY: T must be perfectly valid to receive value.
 	unsafe { mallctl::raw::read_mib(key.as_slice()) }.map_err(map_err)
@@ -408,6 +458,4 @@ fn name(name: &str) -> Result<Name> {
 	Ok(buf)
 }
 
-fn map_err(error: tikv_jemalloc_ctl::Error) -> crate::Error {
+fn map_err(error: jemalloc::ctl::Error) -> crate::Error { err!("mallctl: {}", error.to_string()) }
-	err!("mallctl: {}", error.to_string())
-}

src/database/pool.rs

@@ -287,22 +287,6 @@ fn worker_init(&self, id: usize) {
 	// affinity is empty (no-op) if there's only one queue
 	set_affinity(affinity.clone());
 
-	#[cfg(all(not(target_env = "msvc"), feature = "jemalloc"))]
-	if affinity.clone().count() == 1 && tuwunel_core::alloc::je::is_affine_arena() {
-		use tuwunel_core::{
-			alloc::je::this_thread::{arena_id, set_arena},
-			result::LogDebugErr,
-		};
-
-		let id = affinity.clone().next().expect("at least one id");
-
-		if let Ok(arena) = arena_id() {
-			if arena != id {
-				set_arena(id).log_debug_err().ok();
-			}
-		}
-	}
-
 	trace!(
 		?group,
 		affinity = ?affinity.collect::<Vec<_>>(),

src/main/runtime.rs

@@ -174,15 +174,8 @@ fn set_worker_affinity() {
 }
 
 #[cfg(all(not(target_env = "msvc"), feature = "jemalloc"))]
-fn set_worker_mallctl(id: usize) {
+fn set_worker_mallctl(_id: usize) {
-	use tuwunel_core::alloc::je::{
+	use tuwunel_core::alloc::je::this_thread::set_muzzy_decay;
-		is_affine_arena,
-		this_thread::{set_arena, set_muzzy_decay},
-	};
-
-	if is_affine_arena() {
-		set_arena(id).log_debug_err().ok();
-	}
 
 	let muzzy_option = GC_MUZZY
 		.get()