Bump jemalloc.

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

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 tikv_jemalloc_ctl as mallctl;
-use tikv_jemalloc_sys as ffi;
-use tikv_jemallocator as jemalloc;
+use jevmalloc as jemalloc;
+use jevmalloc::{ctl as mallctl, ffi};
 
 use crate::{
-	Result, err, is_equal_to, is_nonzero,
-	utils::{math, math::Tried},
+	Result,
+	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");
-	acq_epoch().expect("pre-initialization of jemalloc failed");
+	// SAFETY: Mutable static globals in jemalloc crate; must be initialized
+	// 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 {
-	err!("mallctl: {}", error.to_string())
-}
+fn map_err(error: jemalloc::ctl::Error) -> crate::Error { err!("mallctl: {}", error.to_string()) }