Enable heap profiling

proxy_components/proxy_ffi/src/jemalloc_utils.rs

@@ -20,16 +20,14 @@ extern "C" {
     ) -> ::std::os::raw::c_int;
 }
 
+
 #[allow(unused_variables)]
 #[allow(unused_mut)]
 #[allow(unused_unsafe)]
-fn issue_mallctl(command: &str) -> u64 {
-    type PtrUnderlying = u64;
-    let mut ptr: PtrUnderlying = 0;
-    let mut size = std::mem::size_of::<PtrUnderlying>() as u64;
-    let c_str = std::ffi::CString::new(command).unwrap();
-    let c_ptr: *const ::std::os::raw::c_char = c_str.as_ptr() as *const ::std::os::raw::c_char;
+pub fn issue_mallctl_args(command: &str, oldptr: *mut ::std::os::raw::c_void, oldsize: *mut u64, newptr: *mut ::std::os::raw::c_void, newsize: u64) {
     unsafe {
+        let c_str = std::ffi::CString::new(command).unwrap();
+        let c_ptr: *const ::std::os::raw::c_char = c_str.as_ptr() as *const ::std::os::raw::c_char;
         // See unprefixed_malloc_on_supported_platforms in tikv-jemalloc-sys.
         #[cfg(any(test, feature = "testexport"))]
         {
@@ -39,10 +37,10 @@ fn issue_mallctl(command: &str) -> u64 {
                 #[cfg(any(target_os = "android", target_os = "dragonfly", target_os = "macos"))]
                 _rjem_mallctl(
                     c_ptr,
-                    &mut ptr as *mut _ as *mut ::std::os::raw::c_void,
-                    &mut size as *mut u64,
-                    std::ptr::null_mut(),
-                    0,
+                    oldptr,
+                    oldsize,
+                    newptr,
+                    newsize,
                 );
                 #[cfg(not(any(
                     target_os = "android",
@@ -51,10 +49,10 @@ fn issue_mallctl(command: &str) -> u64 {
                 )))]
                 mallctl(
                     c_ptr,
-                    &mut ptr as *mut _ as *mut ::std::os::raw::c_void,
-                    &mut size as *mut u64,
-                    std::ptr::null_mut(),
-                    0,
+                    oldptr,
+                    oldsize,
+                    newptr,
+                    newsize,
                 );
             }
         }
@@ -63,15 +61,25 @@ fn issue_mallctl(command: &str) -> u64 {
         {
             // Must linked to tiflash.
             #[cfg(feature = "external-jemalloc")]
-            mallctl(
+            let r = mallctl(
                 c_ptr,
-                &mut ptr as *mut _ as *mut ::std::os::raw::c_void,
-                &mut size as *mut u64,
-                std::ptr::null_mut(),
-                0,
+                oldptr,
+                oldsize,
+                newptr,
+                newsize,
             );
         }
     }
+}
+
+#[allow(unused_variables)]
+#[allow(unused_mut)]
+#[allow(unused_unsafe)]
+fn issue_mallctl(command: &str) -> u64 {
+    type PtrUnderlying = u64;
+    let mut ptr: PtrUnderlying = 0;
+    let mut size = std::mem::size_of::<PtrUnderlying>() as u64;
+    issue_mallctl_args(command, &mut ptr as *mut _ as *mut ::std::os::raw::c_void, &mut size as *mut u64, std::ptr::null_mut(), 0);
     ptr
 }
 

proxy_components/proxy_ffi/src/lib.rs

@@ -1,5 +1,5 @@
 // Copyright 2022 TiKV Project Authors. Licensed under Apache-2.0.
-
+#![feature(rustc_private)]
 #![feature(thread_id_value)]
 
 /// All mods end up with `_impls` impl structs defined in interface.

proxy_components/proxy_server/Cargo.toml

@@ -55,6 +55,7 @@ encryption_export = { workspace = true, default-features = false }
 engine_rocks = { workspace = true, default-features = false }
 engine_rocks_helper = { workspace = true }
 service = { workspace = true }
+proxy_ffi = { workspace = true, default-features = false }
 engine_store_ffi = { workspace = true, default-features = false }
 engine_tiflash = { workspace = true, default-features = false }
 engine_traits = { workspace = true, default-features = false }

proxy_components/proxy_server/src/status_server/mod.rs

@@ -1,6 +1,7 @@
 // Copyright 2018 TiKV Project Authors. Licensed under Apache-2.0.
 
 mod profile;
+mod vendored_utils;
 
 use std::{
     error::Error as StdError,
@@ -690,9 +691,9 @@ where
                             (Method::GET, "/debug/pprof/heap_deactivate") => {
                                 Self::deactivate_heap_prof(req)
                             }
-                            // (Method::GET, "/debug/pprof/heap") => {
-                            //     Self::dump_heap_prof_to_resp(req).await
-                            // }
+                            (Method::GET, "/debug/pprof/heap") => {
+                                Self::dump_heap_prof_to_resp(req).await
+                            }
                             (Method::GET, "/config") => {
                                 Self::get_config(req, &cfg_controller, engine_store_server_helper)
                                     .await

proxy_components/proxy_server/src/status_server/profile.rs

@@ -22,7 +22,7 @@ use pprof::protos::Message;
 use regex::Regex;
 use tempfile::{NamedTempFile, TempDir};
 #[cfg(not(test))]
-use tikv_alloc::{activate_prof, deactivate_prof, dump_prof};
+use super::vendored_utils::{activate_prof, deactivate_prof, dump_prof};
 use tokio::sync::{Mutex, MutexGuard};
 
 #[cfg(test)]

proxy_components/proxy_server/src/status_server/vendored_utils.rs

@@ -0,0 +1,50 @@
+// Copyright 2024 TiKV Project Authors. Licensed under Apache-2.0.
+
+use proxy_ffi::jemalloc_utils::issue_mallctl_args;
+use tikv_alloc::error::ProfResult;
+
+pub fn activate_prof() -> ProfResult<()> {
+    {
+        let mut value: bool = true;
+        let mut len = std::mem::size_of_val(&value) as u64;
+        issue_mallctl_args(
+            "prof.active",
+            &mut value as *mut _ as *mut _,
+            &mut len,
+            &mut value as *mut _ as *mut _,
+            len,
+        );
+    }
+    Ok(())
+}
+
+pub fn deactivate_prof() -> ProfResult<()> {
+    {
+        let mut value: bool = false;
+        let mut len = std::mem::size_of_val(&value) as u64;
+        issue_mallctl_args(
+            "prof.active",
+            &mut value as *mut _ as *mut _,
+            &mut len,
+            &mut value as *mut _ as *mut _,
+            len,
+        );
+    }
+    Ok(())
+}
+
+pub fn dump_prof(path: &str) -> tikv_alloc::error::ProfResult<()> {
+    {
+        let mut bytes = std::ffi::CString::new(path)?.into_bytes_with_nul();
+        let mut ptr = bytes.as_mut_ptr() as *mut ::std::os::raw::c_char;
+        let len = std::mem::size_of_val(&ptr) as u64;
+        issue_mallctl_args(
+            "prof.dump",
+            std::ptr::null_mut(),
+            std::ptr::null_mut(),
+            &mut ptr as *mut _ as *mut _,
+            len,
+        );
+    }
+    Ok(())
+}