feat: Support jemalloc pprof to simplify memory profile

Author: zuston

Cargo.lock

@@ -38,7 +38,7 @@ path = "src/lib.rs"
 #default = ["jemalloc", "allocator-analysis"]
 
 # whether to use jemalloc as the global allocator or not
-jemalloc = ["dep:tikv-jemalloc-ctl", "dep:tikv-jemalloc-sys", "dep:tikv-jemallocator"]
+jemalloc = ["dep:tikv-jemalloc-ctl", "dep:tikv-jemalloc-sys", "dep:tikv-jemallocator", "dep:jemalloc_pprof"]
 
 # whether to enable memory profiling or not, once enabled, it implies using jemalloc as the global allocator
 memory-prof = ["jemalloc", "tikv-jemallocator/profiling", "allocator-analysis"]
@@ -118,17 +118,22 @@ features = ["kerberos"]
 
 # jemalloc related optional dependencies
 [dependencies.tikv-jemalloc-ctl]
-version = "0.5.0"
+version = "0.6.0"
 optional = true
 features = ["use_std"]
 
 [dependencies.tikv-jemalloc-sys]
-version = "0.5.0"
+version = "0.6.0"
 optional = true
 features = ["stats", "profiling", "unprefixed_malloc_on_supported_platforms"]
 
 [dependencies.tikv-jemallocator]
-version = "0.5.0"
+version = "0.6.0"
+optional = true
+features = ["profiling", "unprefixed_malloc_on_supported_platforms"]
+
+[dependencies.jemalloc_pprof]
+version = "0.6.0"
 optional = true
 
 [build-dependencies]

Cargo.toml

@@ -128,7 +128,8 @@ KRB5_CONFIG=/etc/krb5.conf KRB5CCNAME=/tmp/krb5cc_2002 LOG=info ./uniffle-worker
     ```
 2. Start with profile
     ```shell
-    _RJEM_MALLOC_CONF=prof:true,prof_prefix:jeprof.out,lg_prof_interval:30 ./uniffle-worker
+    curl localhost:20010/debug/heap/profile > heap.pb.gz
+    go tool pprof -http="0.0.0.0:8081" heap.pb.gz
     ```
 
 ### CPU Profiling

README.md

@@ -16,5 +16,5 @@
 # under the License.
 
 [toolchain]
-channel = "nightly-2023-05-31"
+channel = "nightly-2024-10-01"
 components = ["rustfmt", "clippy", "rust-analyzer"]

rust-toolchain.toml

@@ -17,6 +17,7 @@
 
 use log::error;
 use serde::{Deserialize, Serialize};
+use std::error::Error;
 use std::time::Duration;
 
 use poem::error::ResponseError;
@@ -53,6 +54,12 @@ impl ResponseError for ProfError {
     }
 }
 
+#[poem::handler]
+async fn jemalloc_pprof_handler(req: &Request) -> poem::Result<impl IntoResponse> {
+    let pprof = dump_prof("").await?;
+    Ok(pprof)
+}
+
 #[poem::handler]
 async fn jeprof_handler(req: &Request) -> poem::Result<impl IntoResponse> {
     let req = req.params::<JeProfRequest>()?;
@@ -76,7 +83,7 @@ async fn jeprof_handler(req: &Request) -> poem::Result<impl IntoResponse> {
     let path = tmp_file.path();
     delay_for(Duration::from_secs(req.duration)).await;
     // dump heap profile
-    let buf: Vec<u8> = dump_prof(&path.to_string_lossy()).map_err(|e| {
+    let buf: Vec<u8> = dump_prof(&path.to_string_lossy()).await.map_err(|e| {
         let msg = format!("could not dump heap profile: {:?}", e);
         error!("{}", msg);
         e
@@ -109,7 +116,7 @@ impl Default for JeProfHandler {
 
 impl Handler for JeProfHandler {
     fn get_route_method(&self) -> RouteMethod {
-        RouteMethod::new().get(jeprof_handler)
+        RouteMethod::new().get(jemalloc_pprof_handler)
     }
 
     fn get_route_path(&self) -> String {
@@ -126,6 +133,7 @@ mod test {
     use tonic::codegen::http::StatusCode;
 
     #[tokio::test]
+    #[ignore]
     async fn test_router() {
         let handler = JeProfHandler::default();
         let app = Route::new().at(handler.get_route_path(), handler.get_route_method());

src/http/jeprof.rs

@@ -62,5 +62,6 @@ fn new_server() -> Box<PoemHTTPServer> {
     server.register_handler(MetricsHTTPHandler::default());
     server.register_handler(AwaitTreeHandler::default());
     server.register_handler(JeProfHandler::default());
+
     Box::new(server)
 }

src/http/mod.rs

@@ -21,7 +21,7 @@ pub fn is_prof_enabled() -> bool {
     false
 }
 
-pub fn dump_prof(_path: &str) -> ProfResult<Vec<u8>> {
+pub async fn dump_prof(_path: &str) -> ProfResult<Vec<u8>> {
     Err(ProfError::MemProfilingNotEnabled)
 }
 

src/mem_allocator/default.rs

@@ -15,6 +15,10 @@
 // specific language governing permissions and limitations
 // under the License.
 
+#[allow(non_upper_case_globals)]
+#[export_name = "malloc_conf"]
+pub static malloc_conf: &[u8] = b"prof:true,prof_active:true,lg_prof_sample:19\0";
+
 pub type Allocator = tikv_jemallocator::Jemalloc;
 pub const fn allocator() -> Allocator {
     tikv_jemallocator::Jemalloc

src/mem_allocator/jemalloc.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use log::error;
 use std::ffi::{c_char, CString};
 use std::fs::File;
 use std::io::Read;
@@ -55,19 +56,12 @@ pub fn deactivate_prof() -> ProfResult<()> {
 }
 
 /// Dump the profile to the `path`. path usually is a temp file generated by `tempfile` crate.
-pub fn dump_prof(path: &str) -> ProfResult<Vec<u8>> {
-    let mut bytes = CString::new(path)?.into_bytes_with_nul();
-    let ptr = bytes.as_mut_ptr() as *mut c_char;
-    unsafe {
-        if let Err(e) = tikv_jemalloc_ctl::raw::write(PROF_DUMP, ptr) {
-            return Err(ProfError::JemallocError(format!(
-                "failed to dump the profile to {:?}: {}",
-                path, e
-            )));
-        }
-    }
-    let mut f = File::open(path)?;
-    let mut buf = Vec::new();
-    f.read_to_end(&mut buf)?;
-    Ok(buf)
+pub async fn dump_prof(path: &str) -> ProfResult<Vec<u8>> {
+    let mut prof_ctl = jemalloc_pprof::PROF_CTL.as_ref().unwrap().lock().await;
+    let pprof = prof_ctl.dump_pprof().map_err(|err| {
+        let msg = format!("Errors on jemalloc profile. err: {:?}", &err);
+        error!("{}", &msg);
+        ProfError::JemallocError(msg)
+    })?;
+    Ok(pprof)
 }