Support arena.i.purge and malloc_stats_print when jemalloc is used (#402

)

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 }

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"

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 }

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"] }

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"

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 }

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<unsafe extern "C" fn(*mut c_void, *const i8)>,
+        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::<PtrUnderlying>() 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<String>,
+}
+
+#[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()
+}

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<Body>) -> hyper::Result<Response<Body>> {
+        jeprof_purge_arena();
+        Ok(make_response(StatusCode::OK, "purge OK"))
+    }
+
+    async fn memory_status(_: Request<Body>) -> hyper::Result<Response<Body>> {
+        let s = jeprof_memory_status();
+        Ok(make_response(StatusCode::OK, s))
+    }
+
     #[allow(dead_code)]
     async fn dump_heap_prof_to_resp(req: Request<Body>) -> hyper::Result<Response<Body>> {
         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

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()
+}