start mallctl suite w/ jemalloc stats

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

Cargo.lock

@@ -606,6 +606,7 @@ dependencies = [
  "thiserror",
  "thread_local",
  "tikv-jemalloc-ctl",
+ "tikv-jemalloc-sys",
  "tikv-jemallocator",
  "tokio",
  "tower",

Cargo.toml

@@ -215,11 +215,16 @@ version = "0.32.3"
 optional = true
 
 # optional jemalloc usage
+[dependencies.tikv-jemalloc-sys]
+version = "0.5.4"
+optional = true
+default-features = false
+features = ["stats", "unprefixed_malloc_on_supported_platforms"]
 [dependencies.tikv-jemallocator]
 version = "0.5.4"
 optional = true
 default-features = false
-features = ["unprefixed_malloc_on_supported_platforms"]
+features = ["stats", "unprefixed_malloc_on_supported_platforms"]
 [dependencies.tikv-jemalloc-ctl]
 version = "0.5.4"
 optional = true
@@ -364,7 +369,8 @@ default = [
 backend_sqlite = ["sqlite"]
 backend_rocksdb = ["rocksdb"]
 rocksdb = ["rust-rocksdb"]
-jemalloc = ["tikv-jemalloc-ctl", "tikv-jemallocator", "rust-rocksdb/jemalloc"]
+jemalloc = ["tikv-jemalloc-sys", "tikv-jemalloc-ctl", "tikv-jemallocator", "rust-rocksdb/jemalloc"]
+jemalloc_prof = ["tikv-jemalloc-sys/profiling"]
 sqlite = ["rusqlite", "parking_lot", "thread_local"]
 systemd = ["sd-notify"]
 sentry_telemetry = ["sentry", "sentry-tracing", "sentry-tower"]

src/alloc/hardened.rs

@@ -0,0 +1,10 @@
+#![cfg(all(not(target_env = "msvc"), feature = "hardened_malloc", target_os = "linux", not(feature = "jemalloc")))]
+
+#[global_allocator]
+static HMALLOC: hardened_malloc_rs::HardenedMalloc = hardened_malloc_rs::HardenedMalloc;
+
+pub(crate) fn memory_usage() -> String {
+	String::default() //TODO: get usage
+}
+
+pub(crate) fn memory_stats() -> String { "Extended statistics are not available from hardened_malloc.".to_owned() }

src/alloc/je.rs

@@ -0,0 +1,48 @@
+#![cfg(all(not(target_env = "msvc"), feature = "jemalloc", not(feature = "hardened_malloc")))]
+#![allow(dead_code)]
+
+use std::ffi::{c_char, c_void};
+
+use tikv_jemalloc_ctl as mallctl;
+use tikv_jemalloc_sys as ffi;
+use tikv_jemallocator as jemalloc;
+
+#[global_allocator]
+static JEMALLOC: jemalloc::Jemalloc = jemalloc::Jemalloc;
+
+pub(crate) fn version() -> &'static str { mallctl::version::read().expect("version string") }
+
+pub(crate) fn memory_usage() -> String {
+	use mallctl::stats;
+	let allocated = stats::allocated::read().unwrap_or_default() as f64 / 1024.0 / 1024.0;
+	let active = stats::active::read().unwrap_or_default() as f64 / 1024.0 / 1024.0;
+	let mapped = stats::mapped::read().unwrap_or_default() as f64 / 1024.0 / 1024.0;
+	let metadata = stats::metadata::read().unwrap_or_default() as f64 / 1024.0 / 1024.0;
+	let resident = stats::resident::read().unwrap_or_default() as f64 / 1024.0 / 1024.0;
+	let retained = stats::retained::read().unwrap_or_default() as f64 / 1024.0 / 1024.0;
+	format!(
+		" allocated: {allocated:.2} MiB\n active: {active:.2} MiB\n mapped: {mapped:.2} MiB\n metadata: {metadata:.2} \
+		 MiB\n resident: {resident:.2} MiB\n retained: {retained:.2} MiB\n "
+	)
+}
+
+pub(crate) fn memory_stats() -> String {
+	const MAX_LENGTH: usize = 65536 - 4096;
+
+	let opts_s = "d";
+	let mut str: String = String::new();
+
+	let opaque: *mut c_void = &mut str as *mut _ as *mut c_void;
+	let opts_p: *const c_char = std::ffi::CString::new(opts_s).expect("cstring").into_raw() as *const c_char;
+	unsafe { ffi::malloc_stats_print(Some(malloc_stats_cb), opaque, opts_p) };
+
+	str.truncate(MAX_LENGTH);
+	format!("<code>{str}</code>")
+}
+
+extern "C" fn malloc_stats_cb(opaque: *mut c_void, msg: *const c_char) {
+	let res: &mut String = unsafe { std::mem::transmute::<*mut c_void, &mut String>(opaque) };
+	let msg = unsafe { std::ffi::CStr::from_ptr(msg) };
+	let msg = String::from_utf8_lossy(msg.to_bytes());
+	res.push_str(msg.as_ref());
+}

src/alloc/mod.rs

@@ -0,0 +1,20 @@
+pub(crate) mod hardened;
+pub(crate) mod je;
+
+#[cfg(all(not(target_env = "msvc"), feature = "hardened_malloc", target_os = "linux", not(feature = "jemalloc")))]
+pub(crate) fn memory_usage() -> String { hardened::memory_usage() }
+
+#[cfg(all(not(target_env = "msvc"), feature = "jemalloc", not(feature = "hardened_malloc")))]
+pub(crate) fn memory_usage() -> String { je::memory_usage() }
+
+#[cfg(any(target_env = "msvc", all(not(feature = "jemalloc"), not(feature = "hardened_malloc"))))]
+pub(crate) fn memory_usage() -> String { String::default() }
+
+#[cfg(all(not(target_env = "msvc"), feature = "hardened_malloc", target_os = "linux", not(feature = "jemalloc")))]
+pub(crate) fn memory_stats() -> String { hardened::memory_stats() }
+
+#[cfg(all(not(target_env = "msvc"), feature = "jemalloc", not(feature = "hardened_malloc")))]
+pub(crate) fn memory_stats() -> String { je::memory_stats() }
+
+#[cfg(any(target_env = "msvc", all(not(feature = "jemalloc"), not(feature = "hardened_malloc"))))]
+pub(crate) fn memory_stats() -> String { String::default() }

src/main.rs

@@ -49,6 +49,7 @@ use utils::{
 	error::{Error, Result},
 };
 
+mod alloc;
 mod api;
 mod config;
 mod database;
@@ -66,14 +67,6 @@ pub(crate) fn services() -> &'static Services<'static> {
 		.expect("SERVICES should be initialized when this is called")
 }
 
-#[cfg(all(not(target_env = "msvc"), feature = "jemalloc", not(feature = "hardened_malloc")))]
-#[global_allocator]
-static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
-
-#[cfg(all(not(target_env = "msvc"), feature = "hardened_malloc", target_os = "linux", not(feature = "jemalloc")))]
-#[global_allocator]
-static GLOBAL: hardened_malloc_rs::HardenedMalloc = hardened_malloc_rs::HardenedMalloc;
-
 struct Server {
 	config: Config,
 

src/service/admin/debug/debug_commands.rs

@@ -453,3 +453,7 @@ pub(crate) async fn resolve_true_destination(
 		"Actual destination: {actual_dest:?} | Hostname URI: {hostname_uri}"
 	)))
 }
+
+pub(crate) fn memory_stats() -> RoomMessageEventContent {
+	RoomMessageEventContent::text_html("HTML only".to_owned(), crate::alloc::memory_stats())
+}

src/service/admin/debug/mod.rs

@@ -3,7 +3,7 @@ use ruma::{events::room::message::RoomMessageEventContent, EventId, RoomId, Serv
 
 use self::debug_commands::{
 	change_log_level, force_device_list_updates, get_auth_chain, get_pdu, get_remote_pdu, get_remote_pdu_list,
-	get_room_state, parse_pdu, ping, resolve_true_destination, sign_json, verify_json,
+	get_room_state, memory_stats, parse_pdu, ping, resolve_true_destination, sign_json, verify_json,
 };
 use crate::Result;
 
@@ -117,6 +117,9 @@ pub(crate) enum DebugCommand {
 		#[arg(short, long)]
 		no_cache: bool,
 	},
+
+	/// - Print extended memory usage
+	MemoryStats,
 }
 
 pub(crate) async fn process(command: DebugCommand, body: Vec<&str>) -> Result<RoomMessageEventContent> {
@@ -153,5 +156,6 @@ pub(crate) async fn process(command: DebugCommand, body: Vec<&str>) -> Result<Ro
 			server_name,
 			no_cache,
 		} => resolve_true_destination(body, server_name, no_cache).await?,
+		DebugCommand::MemoryStats => memory_stats(),
 	})
 }

src/service/admin/server/server_commands.rs

@@ -26,11 +26,12 @@ pub(crate) async fn show_config(_body: Vec<&str>) -> Result<RoomMessageEventCont
 }
 
 pub(crate) async fn memory_usage(_body: Vec<&str>) -> Result<RoomMessageEventContent> {
-	let response1 = services().memory_usage().await;
-	let response2 = services().globals.db.memory_usage();
+	let response0 = services().memory_usage().await;
+	let response1 = services().globals.db.memory_usage();
+	let response2 = crate::alloc::memory_usage();
 
 	Ok(RoomMessageEventContent::text_plain(format!(
-		"Services:\n{response1}\n\nDatabase:\n{response2}"
+		"Services:\n{response0}\n\nDatabase:\n{response1}\nAllocator:\n{response2}"
 	)))
 }