diff --git a/src/admin/debug/commands.rs b/src/admin/debug/commands.rs index 07daaf0a..d027fa73 100644 --- a/src/admin/debug/commands.rs +++ b/src/admin/debug/commands.rs @@ -923,3 +923,12 @@ pub(super) async fn database_stats( Ok(RoomMessageEventContent::notice_markdown(out)) } + +#[admin_command] +pub(super) async fn trim_memory(&self) -> Result { + conduwuit::alloc::trim()?; + + writeln!(self, "done").await?; + + Ok(RoomMessageEventContent::notice_plain("")) +} diff --git a/src/admin/debug/mod.rs b/src/admin/debug/mod.rs index c87dbb0a..cc2a8ddd 100644 --- a/src/admin/debug/mod.rs +++ b/src/admin/debug/mod.rs @@ -207,6 +207,9 @@ pub(super) enum DebugCommand { map: Option, }, + /// - Trim memory usage + TrimMemory, + /// - Developer test stubs #[command(subcommand)] #[allow(non_snake_case)] diff --git a/src/core/alloc/default.rs b/src/core/alloc/default.rs index 83bfca7d..5db02884 100644 --- a/src/core/alloc/default.rs +++ b/src/core/alloc/default.rs @@ -1,5 +1,8 @@ //! Default allocator with no special features +/// Always returns Ok +pub fn trim() -> crate::Result { Ok(()) } + /// Always returns None #[must_use] pub fn memory_stats() -> Option { None } diff --git a/src/core/alloc/hardened.rs b/src/core/alloc/hardened.rs index 335a3307..e2d9b28e 100644 --- a/src/core/alloc/hardened.rs +++ b/src/core/alloc/hardened.rs @@ -3,6 +3,8 @@ #[global_allocator] static HMALLOC: hardened_malloc_rs::HardenedMalloc = hardened_malloc_rs::HardenedMalloc; +pub fn trim() -> crate::Result { Ok(()) } + #[must_use] //TODO: get usage pub fn memory_usage() -> Option { None } diff --git a/src/core/alloc/je.rs b/src/core/alloc/je.rs index 423f5408..b2c1fe85 100644 --- a/src/core/alloc/je.rs +++ b/src/core/alloc/je.rs @@ -1,18 +1,45 @@ //! jemalloc allocator -use std::ffi::{c_char, c_void}; +use std::{ + cell::OnceCell, + ffi::{c_char, c_void}, + fmt::{Debug, Write}, +}; +use arrayvec::ArrayVec; +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}; + +#[cfg(feature = "jemalloc_conf")] +#[no_mangle] +pub static malloc_conf: &[u8] = b"\ +metadata_thp:always\ +,percpu_arena:percpu\ +,background_thread:true\ +,max_background_threads:-1\ +,lg_extent_max_active_fit:4\ +,oversize_threshold:33554432\ +,tcache_max:2097152\ +,dirty_decay_ms:16000\ +,muzzy_decay_ms:144000\ +\0"; + #[global_allocator] static JEMALLOC: jemalloc::Jemalloc = jemalloc::Jemalloc; +type Key = ArrayVec; +type Name = ArrayVec; + +const KEY_SEGS: usize = 8; +const NAME_MAX: usize = 128; + #[must_use] #[cfg(feature = "jemalloc_stats")] pub fn memory_usage() -> Option { use mallctl::stats; - use tikv_jemalloc_ctl as mallctl; let mibs = |input: Result| { let input = input.unwrap_or_default(); @@ -62,7 +89,12 @@ pub fn memory_stats() -> Option { unsafe extern "C" fn malloc_stats_cb(opaque: *mut c_void, msg: *const c_char) { // SAFETY: we have to trust the opaque points to our String - let res: &mut String = unsafe { opaque.cast::().as_mut().unwrap() }; + let res: &mut String = unsafe { + opaque + .cast::() + .as_mut() + .expect("failed to cast void* to &mut String") + }; // SAFETY: we have to trust the string is null terminated. let msg = unsafe { std::ffi::CStr::from_ptr(msg) }; @@ -70,3 +102,92 @@ unsafe extern "C" fn malloc_stats_cb(opaque: *mut c_void, msg: *const c_char) { let msg = String::from_utf8_lossy(msg.to_bytes()); res.push_str(msg.as_ref()); } + +macro_rules! mallctl { + ($name:literal) => {{ + thread_local! { + static KEY: OnceCell = OnceCell::default(); + }; + + KEY.with(|once| { + once.get_or_init(move || key($name).expect("failed to translate name into mib key")) + .clone() + }) + }}; +} + +pub fn trim() -> Result { set(&mallctl!("arena.4096.purge"), ()) } + +pub fn decay() -> Result { set(&mallctl!("arena.4096.purge"), ()) } + +pub fn set_by_name(name: &str, val: T) -> Result { set(&key(name)?, val) } + +pub fn get_by_name(name: &str) -> Result { get(&key(name)?) } + +pub mod this_thread { + use super::{get, key, set, Key, OnceCell, Result}; + + pub fn trim() -> Result { + let mut key = mallctl!("arena.0.purge"); + key[1] = arena_id()?.try_into()?; + set(&key, ()) + } + + pub fn decay() -> Result { + let mut key = mallctl!("arena.0.decay"); + key[1] = arena_id()?.try_into()?; + set(&key, ()) + } + + pub fn cache(enable: bool) -> Result { + set(&mallctl!("thread.tcache.enabled"), u8::from(enable)) + } + + pub fn flush() -> Result { set(&mallctl!("thread.tcache.flush"), ()) } + + pub fn allocated() -> Result { get::(&mallctl!("thread.allocated")) } + + pub fn deallocated() -> Result { get::(&mallctl!("thread.deallocated")) } + + pub fn arena_id() -> Result { get::(&mallctl!("thread.arena")) } +} + +fn set(key: &Key, val: T) -> Result +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) +} + +fn get(key: &Key) -> Result +where + T: Copy + Debug, +{ + // SAFETY: T must be perfectly valid to receive value. + unsafe { mallctl::raw::read_mib(key.as_slice()) }.map_err(map_err) +} + +fn key(name: &str) -> Result { + // tikv asserts the output buffer length is tight to the number of required mibs + // so we slice that down here. + let segs = name.chars().filter(is_equal_to!(&'.')).count().try_add(1)?; + + let name = self::name(name)?; + let mut buf = [0_usize; KEY_SEGS]; + mallctl::raw::name_to_mib(name.as_slice(), &mut buf[0..segs]) + .map_err(map_err) + .map(move |()| buf.into_iter().take(segs).collect()) +} + +fn name(name: &str) -> Result { + let mut buf = Name::new(); + buf.try_extend_from_slice(name.as_bytes())?; + buf.try_extend_from_slice(b"\0")?; + + Ok(buf) +} + +fn map_err(error: tikv_jemalloc_ctl::Error) -> crate::Error { + err!("mallctl: {}", error.to_string()) +} diff --git a/src/core/alloc/mod.rs b/src/core/alloc/mod.rs index 31eb033c..0ed1b1a6 100644 --- a/src/core/alloc/mod.rs +++ b/src/core/alloc/mod.rs @@ -4,7 +4,7 @@ #[cfg(all(not(target_env = "msvc"), feature = "jemalloc"))] pub mod je; #[cfg(all(not(target_env = "msvc"), feature = "jemalloc"))] -pub use je::{memory_stats, memory_usage}; +pub use je::{memory_stats, memory_usage, trim}; #[cfg(all(not(target_env = "msvc"), feature = "hardened_malloc", not(feature = "jemalloc")))] pub mod hardened; @@ -13,7 +13,7 @@ pub mod hardened; feature = "hardened_malloc", not(feature = "jemalloc") ))] -pub use hardened::{memory_stats, memory_usage}; +pub use hardened::{memory_stats, memory_usage, trim}; #[cfg(any( target_env = "msvc", @@ -24,4 +24,4 @@ pub mod default; target_env = "msvc", all(not(feature = "hardened_malloc"), not(feature = "jemalloc")) ))] -pub use default::{memory_stats, memory_usage}; +pub use default::{memory_stats, memory_usage, trim};