diff --git a/Cargo.lock b/Cargo.lock index d94142d2..1d4a0e9d 100644 --- a/Cargo.lock +++ b/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", diff --git a/Cargo.toml b/Cargo.toml index df012a6f..614cc7a1 100644 --- a/Cargo.toml +++ b/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 diff --git a/src/core/Cargo.toml b/src/core/Cargo.toml index 95c8f1bc..90ea8bf2 100644 --- a/src/core/Cargo.toml +++ b/src/core/Cargo.toml @@ -26,18 +26,14 @@ gzip_compression = [ "reqwest/gzip", ] jemalloc = [ - "dep:tikv-jemalloc-sys", - "dep:tikv-jemalloc-ctl", - "dep:tikv-jemallocator", + "dep:jevmalloc", ] jemalloc_conf = [] jemalloc_prof = [ - #"tikv-jemalloc-sys/profiling", + #"jevmalloc/profiling", ] jemalloc_stats = [ - #"tikv-jemalloc-sys/stats", - #"tikv-jemalloc-ctl/stats", - #"tikv-jemallocator/stats", + #"jevmalloc/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 -tikv-jemallocator.workspace = true -tikv-jemalloc-ctl.optional = true -tikv-jemalloc-ctl.workspace = true -tikv-jemalloc-sys.optional = true -tikv-jemalloc-sys.workspace = true +jevmalloc.optional = true +jevmalloc.workspace = true tokio.workspace = true tokio-metrics.workspace = true toml.workspace = true diff --git a/src/core/alloc/je.rs b/src/core/alloc/je.rs index 483a46f7..ff7eb1b1 100644 --- a/src/core/alloc/je.rs +++ b/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; type Key = ArrayVec; 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::::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::(), 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(key: &Key, val: T) -> Result 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 { 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()) } diff --git a/src/database/pool.rs b/src/database/pool.rs index d7160e11..ae77c122 100644 --- a/src/database/pool.rs +++ b/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::>(), diff --git a/src/main/runtime.rs b/src/main/runtime.rs index 34a03f60..294e53af 100644 --- a/src/main/runtime.rs +++ b/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) { - use tuwunel_core::alloc::je::{ - is_affine_arena, - this_thread::{set_arena, set_muzzy_decay}, - }; - - if is_affine_arena() { - set_arena(id).log_debug_err().ok(); - } +fn set_worker_mallctl(_id: usize) { + use tuwunel_core::alloc::je::this_thread::set_muzzy_decay; let muzzy_option = GC_MUZZY .get()