add preferred jemalloc config

add muzzy/dirty configuration mallctl interface

add program argument for --gc-muzzy=false

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

src/core/Cargo.toml

@@ -36,6 +36,7 @@ jemalloc_stats = [
     "tikv-jemalloc-ctl/stats",
     "tikv-jemallocator/stats",
 ]
+jemalloc_conf = []
 hardened_malloc = [
 	"dep:hardened_malloc-rs"
 ]

src/core/alloc/default.rs

@@ -1,7 +1,7 @@
 //! Default allocator with no special features
 
 /// Always returns Ok
-pub fn trim() -> crate::Result { Ok(()) }
+pub fn trim<I: Into<Option<usize>>>(_: I) -> crate::Result { Ok(()) }
 
 /// Always returns None
 #[must_use]

src/core/alloc/hardened.rs

@@ -3,7 +3,7 @@
 #[global_allocator]
 static HMALLOC: hardened_malloc_rs::HardenedMalloc = hardened_malloc_rs::HardenedMalloc;
 
-pub fn trim() -> crate::Result { Ok(()) }
+pub fn trim<I: Into<Option<usize>>>(_: I) -> crate::Result { Ok(()) }
 
 #[must_use]
 //TODO: get usage

src/core/alloc/je.rs

@@ -2,8 +2,9 @@
 
 use std::{
 	cell::OnceCell,
-	ffi::{c_char, c_void},
+	ffi::{c_char, c_void, CStr},
 	fmt::Debug,
+	sync::RwLock,
 };
 
 use arrayvec::ArrayVec;
@@ -11,10 +12,14 @@ use tikv_jemalloc_ctl as mallctl;
 use tikv_jemalloc_sys as ffi;
 use tikv_jemallocator as jemalloc;
 
-use crate::{err, is_equal_to, utils::math::Tried, Result};
+use crate::{
+	err, is_equal_to, is_nonzero,
+	utils::{math, math::Tried},
+	Result,
+};
 
 #[cfg(feature = "jemalloc_conf")]
-#[no_mangle]
+#[unsafe(no_mangle)]
 pub static malloc_conf: &[u8] = b"\
 metadata_thp:always\
 ,percpu_arena:percpu\
@@ -22,19 +27,26 @@ metadata_thp:always\
 ,max_background_threads:-1\
 ,lg_extent_max_active_fit:4\
 ,oversize_threshold:33554432\
-,tcache_max:2097152\
+,tcache_max:1048576\
 ,dirty_decay_ms:16000\
 ,muzzy_decay_ms:144000\
 \0";
 
 #[global_allocator]
 static JEMALLOC: jemalloc::Jemalloc = jemalloc::Jemalloc;
+static CONTROL: RwLock<()> = RwLock::new(());
 
-type Key = ArrayVec<usize, KEY_SEGS>;
 type Name = ArrayVec<u8, NAME_MAX>;
+type Key = ArrayVec<usize, KEY_SEGS>;
 
-const KEY_SEGS: usize = 8;
 const NAME_MAX: usize = 128;
+const KEY_SEGS: usize = 8;
+
+#[crate::ctor]
+fn _static_initialization() {
+	acq_epoch().expect("pre-initialization of jemalloc failed");
+	acq_epoch().expect("pre-initialization of jemalloc failed");
+}
 
 #[must_use]
 #[cfg(feature = "jemalloc_stats")]
@@ -49,6 +61,9 @@ pub fn memory_usage() -> Option<String> {
 		kibs / 1024.0
 	};
 
+	// Acquire the epoch; ensure latest stats are pulled in
+	acq_epoch().ok()?;
+
 	let allocated = mibs(stats::allocated::read());
 	let active = mibs(stats::active::read());
 	let mapped = mibs(stats::mapped::read());
@@ -76,6 +91,9 @@ pub fn memory_stats(opts: &str) -> Option<String> {
 		.into_raw()
 		.cast_const();
 
+	// Acquire the epoch; ensure latest stats are pulled in
+	acq_epoch().ok()?;
+
 	// SAFETY: calls malloc_stats_print() with our string instance which must remain
 	// in this frame. https://docs.rs/tikv-jemalloc-sys/latest/tikv_jemalloc_sys/fn.malloc_stats_print.html
 	unsafe { ffi::malloc_stats_print(Some(malloc_stats_cb), opaque, opts_p) };
@@ -95,7 +113,7 @@ unsafe extern "C" fn malloc_stats_cb(opaque: *mut c_void, msg: *const c_char) {
 	};
 
 	// SAFETY: we have to trust the string is null terminated.
-	let msg = unsafe { std::ffi::CStr::from_ptr(msg) };
+	let msg = unsafe { CStr::from_ptr(msg) };
 
 	let msg = String::from_utf8_lossy(msg.to_bytes());
 	res.push_str(msg.as_ref());
@@ -114,58 +132,168 @@ macro_rules! mallctl {
 	}};
 }
 
-pub fn trim() -> Result { set(&mallctl!("arena.4096.purge"), ()) }
-
-pub fn decay() -> Result { set(&mallctl!("arena.4096.purge"), ()) }
-
-pub fn set_by_name<T: Copy + Debug>(name: &str, val: T) -> Result { set(&key(name)?, val) }
-
-pub fn get_by_name<T: Copy + Debug>(name: &str) -> Result<T> { get(&key(name)?) }
-
 pub mod this_thread {
-	use super::{get, key, set, Key, OnceCell, Result};
+	use super::{is_nonzero, key, math, Debug, Key, OnceCell, Result};
 
-	pub fn trim() -> Result {
-		let mut key = mallctl!("arena.0.purge");
-		key[1] = arena_id()?.try_into()?;
-		set(&key, ())
+	pub fn trim() -> Result { notify(mallctl!("arena.0.purge")) }
+
+	pub fn decay() -> Result { notify(mallctl!("arena.0.decay")) }
+
+	pub fn flush() -> Result { super::notify(&mallctl!("thread.tcache.flush")) }
+
+	pub fn set_muzzy_decay(decay_ms: isize) -> Result<isize> {
+		set(mallctl!("arena.0.muzzy_decay_ms"), decay_ms)
 	}
 
-	pub fn decay() -> Result {
-		let mut key = mallctl!("arena.0.decay");
-		key[1] = arena_id()?.try_into()?;
-		set(&key, ())
+	pub fn get_muzzy_decay() -> Result<isize> { get(mallctl!("arena.0.muzzy_decay_ms")) }
+
+	pub fn set_dirty_decay(decay_ms: isize) -> Result<isize> {
+		set(mallctl!("arena.0.dirty_decay_ms"), decay_ms)
 	}
 
-	pub fn cache(enable: bool) -> Result {
-		set(&mallctl!("thread.tcache.enabled"), u8::from(enable))
+	pub fn get_dirty_decay() -> Result<isize> { get(mallctl!("arena.0.dirty_decay_ms")) }
+
+	pub fn enable_cache(enable: bool) -> Result<bool> {
+		super::set::<u8>(&mallctl!("thread.tcache.enabled"), enable.into()).map(is_nonzero!())
 	}
 
-	pub fn flush() -> Result { set(&mallctl!("thread.tcache.flush"), ()) }
+	pub fn is_cache_enabled() -> Result<bool> {
+		super::get::<u8>(&mallctl!("thread.tcache.enabled")).map(is_nonzero!())
+	}
 
-	pub fn allocated() -> Result<u64> { get::<u64>(&mallctl!("thread.allocated")) }
+	pub fn set_arena(id: usize) -> Result<usize> {
+		super::set::<u32>(&mallctl!("thread.arena"), id.try_into()?).and_then(math::try_into)
+	}
 
-	pub fn deallocated() -> Result<u64> { get::<u64>(&mallctl!("thread.deallocated")) }
+	pub fn arena_id() -> Result<usize> {
+		super::get::<u32>(&mallctl!("thread.arena")).and_then(math::try_into)
+	}
 
-	pub fn arena_id() -> Result<u32> { get::<u32>(&mallctl!("thread.arena")) }
+	pub fn allocated() -> Result<u64> { super::get(&mallctl!("thread.allocated")) }
+
+	pub fn deallocated() -> Result<u64> { super::get(&mallctl!("thread.deallocated")) }
+
+	fn notify(key: Key) -> Result { super::notify_by_arena(Some(arena_id()?), key) }
+
+	fn set<T>(key: Key, val: T) -> Result<T>
+	where
+		T: Copy + Debug,
+	{
+		super::set_by_arena(Some(arena_id()?), key, val)
+	}
+
+	fn get<T>(key: Key) -> Result<T>
+	where
+		T: Copy + Debug,
+	{
+		super::get_by_arena(Some(arena_id()?), key)
+	}
 }
 
-fn set<T>(key: &Key, val: T) -> Result
+pub fn trim<I: Into<Option<usize>>>(arena: I) -> Result {
+	notify_by_arena(arena.into(), mallctl!("arena.4096.purge"))
+}
+
+pub fn decay<I: Into<Option<usize>>>(arena: I) -> Result {
+	notify_by_arena(arena.into(), mallctl!("arena.4096.decay"))
+}
+
+pub fn set_muzzy_decay<I: Into<Option<usize>>>(arena: I, decay_ms: isize) -> Result<isize> {
+	if let Some(arena) = arena.into() {
+		set_by_arena(Some(arena), mallctl!("arena.4096.muzzy_decay_ms"), decay_ms)
+	} else {
+		set(&mallctl!("arenas.muzzy_decay_ms"), decay_ms)
+	}
+}
+
+pub fn set_dirty_decay<I: Into<Option<usize>>>(arena: I, decay_ms: isize) -> Result<isize> {
+	if let Some(arena) = arena.into() {
+		set_by_arena(Some(arena), mallctl!("arena.4096.dirty_decay_ms"), decay_ms)
+	} else {
+		set(&mallctl!("arenas.dirty_decay_ms"), decay_ms)
+	}
+}
+
+#[inline]
+#[must_use]
+pub fn is_affine_arena() -> bool { is_percpu_arena() || is_phycpu_arena() }
+
+#[inline]
+#[must_use]
+pub fn is_percpu_arena() -> bool { percpu_arenas().is_ok_and(is_equal_to!("percpu")) }
+
+#[inline]
+#[must_use]
+pub fn is_phycpu_arena() -> bool { percpu_arenas().is_ok_and(is_equal_to!("phycpu")) }
+
+pub fn percpu_arenas() -> Result<&'static str> {
+	let ptr = get::<*const c_char>(&mallctl!("opt.percpu_arena"))?;
+	//SAFETY: ptr points to a null-terminated string returned for opt.percpu_arena.
+	let cstr = unsafe { CStr::from_ptr(ptr) };
+	cstr.to_str().map_err(Into::into)
+}
+
+pub fn arenas() -> Result<usize> {
+	get::<u32>(&mallctl!("arenas.narenas")).and_then(math::try_into)
+}
+
+pub fn inc_epoch() -> Result<u64> { xchg(&mallctl!("epoch"), 1_u64) }
+
+pub fn acq_epoch() -> Result<u64> { xchg(&mallctl!("epoch"), 0_u64) }
+
+fn notify_by_arena(id: Option<usize>, mut key: Key) -> Result {
+	key[1] = id.unwrap_or(4096);
+	notify(&key)
+}
+
+fn set_by_arena<T>(id: Option<usize>, mut key: Key, val: T) -> Result<T>
 where
 	T: Copy + Debug,
 {
-	// SAFETY: T must be the exact expected type.
-	unsafe { mallctl::raw::write_mib(key.as_slice(), val) }.map_err(map_err)
+	key[1] = id.unwrap_or(4096);
+	set(&key, val)
+}
+
+fn get_by_arena<T>(id: Option<usize>, mut key: Key) -> Result<T>
+where
+	T: Copy + Debug,
+{
+	key[1] = id.unwrap_or(4096);
+	get(&key)
+}
+
+fn notify(key: &Key) -> Result { xchg(key, ()) }
+
+fn set<T>(key: &Key, val: T) -> Result<T>
+where
+	T: Copy + Debug,
+{
+	let _lock = CONTROL.write()?;
+	let res = xchg(key, val)?;
+	inc_epoch()?;
+
+	Ok(res)
 }
 
 fn get<T>(key: &Key) -> Result<T>
 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)
 }
 
+fn xchg<T>(key: &Key, val: T) -> Result<T>
+where
+	T: Copy + Debug,
+{
+	// SAFETY: T must be the exact expected type.
+	unsafe { mallctl::raw::update_mib(key.as_slice(), val) }.map_err(map_err)
+}
+
 fn key(name: &str) -> Result<Key> {
 	// tikv asserts the output buffer length is tight to the number of required mibs
 	// so we slice that down here.