From 50b8d2e0facdd8d7940ad6256a09945318cbaf6b Mon Sep 17 00:00:00 2001 From: Calvin Neo Date: Wed, 11 Dec 2024 15:13:08 +0800 Subject: [PATCH] Support arena.i.purge and malloc_stats_print when jemalloc is used (#402) --- components/engine_panic/Cargo.toml | 2 +- components/engine_rocks/Cargo.toml | 2 +- components/engine_traits/Cargo.toml | 2 +- components/hybrid_engine/Cargo.toml | 26 ++--- components/in_memory_engine/Cargo.toml | 42 +++---- components/raft_log_engine/Cargo.toml | 2 +- .../proxy_ffi/src/jemalloc_utils.rs | 103 +++++++++++++++++- .../proxy_server/src/status_server/mod.rs | 17 +++ .../src/status_server/vendored_utils.rs | 25 ++++- 9 files changed, 181 insertions(+), 40 deletions(-) diff --git a/components/engine_panic/Cargo.toml b/components/engine_panic/Cargo.toml index 3ee5d82ad14..c125ef7ccf2 100644 --- a/components/engine_panic/Cargo.toml +++ b/components/engine_panic/Cargo.toml @@ -10,9 +10,9 @@ license = "Apache-2.0" testexport = [] [dependencies] +encryption = { workspace = true } engine_traits = { workspace = true } kvproto = { workspace = true } -encryption = { workspace = true } raft = { workspace = true } tracker = { workspace = true } txn_types = { workspace = true } diff --git a/components/engine_rocks/Cargo.toml b/components/engine_rocks/Cargo.toml index d3893da020a..26ca70e1606 100644 --- a/components/engine_rocks/Cargo.toml +++ b/components/engine_rocks/Cargo.toml @@ -61,6 +61,6 @@ package = "rocksdb" features = ["encryption"] [dev-dependencies] +proptest = "1.0.0" rand = "0.8" toml = "0.5" -proptest = "1.0.0" diff --git a/components/engine_traits/Cargo.toml b/components/engine_traits/Cargo.toml index 2e86822ceac..85b8337ef0d 100644 --- a/components/engine_traits/Cargo.toml +++ b/components/engine_traits/Cargo.toml @@ -11,6 +11,7 @@ testexport = [] [dependencies] collections = { workspace = true } +encryption = { workspace = true } error_code = { workspace = true } fail = "0.5" file_system = { workspace = true } @@ -19,7 +20,6 @@ kvproto = { workspace = true } log_wrappers = { workspace = true } protobuf = "2" raft = { workspace = true } -encryption = { workspace = true } serde = "1.0" slog = { workspace = true } slog-global = { workspace = true } diff --git a/components/hybrid_engine/Cargo.toml b/components/hybrid_engine/Cargo.toml index 3a23f9927da..2938a329f97 100644 --- a/components/hybrid_engine/Cargo.toml +++ b/components/hybrid_engine/Cargo.toml @@ -6,27 +6,27 @@ publish = false license = "Apache-2.0" [dependencies] -engine_traits = { workspace = true } -txn_types = { workspace = true } -tikv_util = { workspace = true } +crossbeam = { workspace = true } engine_rocks = { workspace = true } -online_config = { workspace = true } +engine_traits = { workspace = true } in_memory_engine = { workspace = true } -slog = { workspace = true } -slog-global = { workspace = true } -tempfile = "3.0" +keys = { workspace = true } +kvproto = { workspace = true } +lazy_static = "1.4.0" +online_config = { workspace = true } prometheus = { version = "0.13", default-features = false, features = [ "nightly", ] } prometheus-static-metric = "0.5" -lazy_static = "1.4.0" -crossbeam = { workspace = true } -raftstore = { workspace = true } raft = { workspace = true } -kvproto = { workspace = true } -keys = { workspace = true } +raftstore = { workspace = true } +slog = { workspace = true } +slog-global = { workspace = true } +tempfile = "3.0" +tikv_util = { workspace = true } +txn_types = { workspace = true } [dev-dependencies] +fail = { version = "0.5", features = ["failpoints"] } tempfile = "3.0" test_util = { workspace = true } -fail = { version = "0.5", features = ["failpoints"] } diff --git a/components/in_memory_engine/Cargo.toml b/components/in_memory_engine/Cargo.toml index a29ad0d6daf..00052d4542d 100644 --- a/components/in_memory_engine/Cargo.toml +++ b/components/in_memory_engine/Cargo.toml @@ -20,44 +20,44 @@ path = "benches/load_region.rs" harness = false [dependencies] -engine_traits = { workspace = true } -collections = { workspace = true } -crossbeam-skiplist = { workspace = true } bytes = "1.0" +collections = { workspace = true } crossbeam = { workspace = true } +crossbeam-skiplist = { workspace = true } +dashmap = "5.1" +engine_rocks = { workspace = true } +engine_traits = { workspace = true } +fail = "0.5" futures = { version = "0.3", features = ["compat"] } -tikv_util = { workspace = true } -txn_types = { workspace = true } +hex = "0.4" +keys = { workspace = true } kvproto = { workspace = true } +lazy_static = "1.4.0" +libc = "0.2" log_wrappers = { workspace = true } +online_config = { workspace = true } +parking_lot = "0.12" pd_client = { workspace = true } +prometheus = { version = "0.13", default-features = false, features = ["nightly"] } +prometheus-static-metric = "0.5" raftstore = { workspace = true } -dashmap = "5.1" +rand = "0.8" security = { workspace = true } serde = "1.0" serde_json = "1.0" -slog-global = { workspace = true } slog = { workspace = true } +slog-global = { workspace = true } +smallvec = "1.4" strum = { version = "0.20", features = ["derive"] } -engine_rocks = { workspace = true } -fail = "0.5" -yatp = { workspace = true } -parking_lot = "0.12" -keys = { workspace = true } -prometheus = { version = "0.13", default-features = false, features = ["nightly"] } -prometheus-static-metric = "0.5" -lazy_static = "1.4.0" -hex = "0.4" thiserror = "1.0" -online_config = { workspace = true } -libc = "0.2" -rand = "0.8" +tikv_util = { workspace = true } tokio = { version = "1.5", features = ["rt-multi-thread"] } -smallvec = "1.4" +txn_types = { workspace = true } +yatp = { workspace = true } [dev-dependencies] criterion = "0.3" +proptest = "1.0.0" tempfile = "3.0" test_pd = { workspace = true } test_util = { workspace = true } -proptest = "1.0.0" diff --git a/components/raft_log_engine/Cargo.toml b/components/raft_log_engine/Cargo.toml index c1c48988f44..15f2388b460 100644 --- a/components/raft_log_engine/Cargo.toml +++ b/components/raft_log_engine/Cargo.toml @@ -9,9 +9,9 @@ license = "Apache-2.0" failpoints = ["raft-engine/failpoints"] [dependencies] +codec = { workspace = true } encryption = { workspace = true } engine_traits = { workspace = true } -codec = { workspace = true } file_system = { workspace = true } kvproto = { workspace = true } raft = { workspace = true } diff --git a/proxy_components/proxy_ffi/src/jemalloc_utils.rs b/proxy_components/proxy_ffi/src/jemalloc_utils.rs index 763669fffd5..ce07031fef3 100644 --- a/proxy_components/proxy_ffi/src/jemalloc_utils.rs +++ b/proxy_components/proxy_ffi/src/jemalloc_utils.rs @@ -1,4 +1,5 @@ // Copyright 2024 TiKV Project Authors. Licensed under Apache-2.0. +use std::sync::Mutex; extern "C" { // External jemalloc @@ -18,6 +19,12 @@ extern "C" { newp: *mut ::std::os::raw::c_void, newlen: u64, ) -> ::std::os::raw::c_int; + + pub fn malloc_stats_print( + write_cb: Option, + cbopaque: *mut c_void, + opts: *const i8, + ); } #[allow(unused_variables)] @@ -80,7 +87,7 @@ pub fn issue_mallctl_args( #[allow(unused_variables)] #[allow(unused_mut)] #[allow(unused_unsafe)] -fn issue_mallctl(command: &str) -> u64 { +pub fn issue_mallctl(command: &str) -> u64 { type PtrUnderlying = u64; let mut ptr: PtrUnderlying = 0; let mut size = std::mem::size_of::() as u64; @@ -109,3 +116,97 @@ pub fn get_allocate() -> u64 { pub fn get_deallocate() -> u64 { issue_mallctl("thread.deallocated") } + +use std::ffi::{c_char, c_void, CStr}; +struct CaptureContext { + buffer: Mutex, +} + +#[allow(dead_code)] +extern "C" fn write_to_string(ctx: *mut c_void, message: *const c_char) { + if ctx.is_null() || message.is_null() { + return; + } + + let context = unsafe { &*(ctx as *mut CaptureContext) }; + + let c_str = unsafe { CStr::from_ptr(message) }; + if let Ok(str_slice) = c_str.to_str() { + let mut buffer = context.buffer.lock().unwrap(); + buffer.push_str(str_slice); + } +} + +#[allow(unused_variables)] +#[allow(unused_mut)] +#[allow(unused_unsafe)] +pub fn get_malloc_stats() -> String { + let context = CaptureContext { + buffer: Mutex::new(String::new()), + }; + + // Use Json format. + let c_str = std::ffi::CString::new("J").unwrap(); + let ops_str = c_str.as_ptr() as *const std::os::raw::c_char; + + unsafe { + // See unprefixed_malloc_on_supported_platforms in tikv-jemalloc-sys. + #[cfg(any(test, feature = "testexport"))] + { + // Test part + #[cfg(feature = "jemalloc")] + { + // See NO_UNPREFIXED_MALLOC + #[cfg(any(target_os = "android", target_os = "dragonfly", target_os = "macos"))] + _rjem_malloc_stats_print( + Some(write_to_string), + &context as *const _ as *mut c_void, + ops_str, + ); + #[cfg(not(any( + target_os = "android", + target_os = "dragonfly", + target_os = "macos" + )))] + malloc_stats_print( + Some(write_to_string), + &context as *const _ as *mut c_void, + ops_str, + ); + } + } + + #[cfg(not(any(test, feature = "testexport")))] + { + // No test part + #[cfg(feature = "external-jemalloc")] + { + // Must linked to tiflash. + malloc_stats_print( + Some(write_to_string), + &context as *const _ as *mut c_void, + ops_str, + ); + } + #[cfg(not(feature = "external-jemalloc"))] + { + // Happens only with `raftstore-proxy-main` + #[cfg(not(any( + target_os = "android", + target_os = "dragonfly", + target_os = "macos" + )))] + { + malloc_stats_print( + Some(write_to_string), + &context as *const _ as *mut c_void, + ops_str, + ); + } + } + } + } + + let buffer = context.buffer.lock().unwrap(); + buffer.clone() +} diff --git a/proxy_components/proxy_server/src/status_server/mod.rs b/proxy_components/proxy_server/src/status_server/mod.rs index fc0f7f6afde..036960127e9 100644 --- a/proxy_components/proxy_server/src/status_server/mod.rs +++ b/proxy_components/proxy_server/src/status_server/mod.rs @@ -63,6 +63,7 @@ use tokio::{ sync::oneshot::{self, Receiver, Sender}, }; use tokio_openssl::SslStream; +use vendored_utils::{jeprof_memory_status, jeprof_purge_arena}; use crate::status_server::profile::set_prof_active; @@ -243,6 +244,16 @@ where } } + async fn arena_purge(_: Request) -> hyper::Result> { + jeprof_purge_arena(); + Ok(make_response(StatusCode::OK, "purge OK")) + } + + async fn memory_status(_: Request) -> hyper::Result> { + let s = jeprof_memory_status(); + Ok(make_response(StatusCode::OK, s)) + } + #[allow(dead_code)] async fn dump_heap_prof_to_resp(req: Request) -> hyper::Result> { let query = req.uri().query().unwrap_or(""); @@ -796,6 +807,12 @@ where (Method::GET, "/debug/pprof/heap") => { Self::dump_heap_prof_to_resp(req).await } + (Method::GET, "/debug/pprof/arena_purge") => { + Self::arena_purge(req).await + } + (Method::GET, "/debug/pprof/memory_status") => { + Self::memory_status(req).await + } (Method::GET, "/config") => { Self::get_config(req, &cfg_controller, engine_store_server_helper) .await diff --git a/proxy_components/proxy_server/src/status_server/vendored_utils.rs b/proxy_components/proxy_server/src/status_server/vendored_utils.rs index 14362abd665..d298a0dc3f2 100644 --- a/proxy_components/proxy_server/src/status_server/vendored_utils.rs +++ b/proxy_components/proxy_server/src/status_server/vendored_utils.rs @@ -1,6 +1,6 @@ // Copyright 2024 TiKV Project Authors. Licensed under Apache-2.0. -use proxy_ffi::jemalloc_utils::issue_mallctl_args; +use proxy_ffi::jemalloc_utils::{get_malloc_stats, issue_mallctl, issue_mallctl_args}; use tikv_alloc::error::ProfResult; pub fn activate_prof() -> ProfResult<()> { @@ -107,3 +107,26 @@ pub fn adhoc_dump(path: &str) -> tikv_alloc::error::ProfResult<()> { } Ok(()) } + +pub fn jeprof_purge_arena() { + let narenas = issue_mallctl("arenas.narenas"); + info!("jeprof_purge_arena purge {} arenas", narenas); + for i in 0..narenas { + let a_string = format!("arena.{}.purge", i); + let r = issue_mallctl_args( + &a_string, + std::ptr::null_mut(), + std::ptr::null_mut(), + std::ptr::null_mut(), + 0, + ); + if r != 0 { + warn!("jeprof_purge_arena purge {} return {}", a_string, r); + } + } + info!("jeprof_purge_arena purge {} arenas done", narenas); +} + +pub fn jeprof_memory_status() -> String { + get_malloc_stats() +}