Enable heap profiling

proxy_components/proxy_ffi/src/jemalloc_utils.rs

@@ -23,55 +23,56 @@ extern "C" {
 #[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"))]
         {
             #[cfg(feature = "jemalloc")]
             {
                 // See NO_UNPREFIXED_MALLOC
                 #[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,
-                );
+                _rjem_mallctl(c_ptr, oldptr, oldsize, newptr, newsize);
                 #[cfg(not(any(
                     target_os = "android",
                     target_os = "dragonfly",
                     target_os = "macos"
                 )))]
-                mallctl(
-                    c_ptr,
-                    &mut ptr as *mut _ as *mut ::std::os::raw::c_void,
-                    &mut size as *mut u64,
-                    std::ptr::null_mut(),
-                    0,
-                );
+                mallctl(c_ptr, oldptr, oldsize, newptr, newsize);
             }
         }
 
         #[cfg(not(any(test, feature = "testexport")))]
         {
             // Must linked to tiflash.
             #[cfg(feature = "external-jemalloc")]
-            mallctl(
-                c_ptr,
-                &mut ptr as *mut _ as *mut ::std::os::raw::c_void,
-                &mut size as *mut u64,
-                std::ptr::null_mut(),
-                0,
-            );
+            let r = mallctl(c_ptr, 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

@@ -21,14 +21,14 @@ use lazy_static::lazy_static;
 use pprof::protos::Message;
 use regex::Regex;
 use tempfile::{NamedTempFile, TempDir};
-#[cfg(not(test))]
-use tikv_alloc::{activate_prof, deactivate_prof, dump_prof};
 use tokio::sync::{Mutex, MutexGuard};
 
 #[cfg(test)]
 pub use self::test_utils::TEST_PROFILE_MUTEX;
 #[cfg(test)]
 use self::test_utils::{activate_prof, deactivate_prof, dump_prof};
+#[cfg(not(test))]
+use super::vendored_utils::{activate_prof, deactivate_prof, dump_prof};
 
 // File name suffix for periodically dumped heap profiles.
 const HEAP_PROFILE_SUFFIX: &str = ".heap";

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 len = std::mem::size_of_val(&value) as u64;
+        issue_mallctl_args(
+            "prof.active",
+            std::ptr::null_mut(),
+            std::ptr::null_mut(),
+            &mut value as *mut _ as *mut _,
+            len,
+        );
+    }
+    Ok(())
+}
+
+pub fn deactivate_prof() -> ProfResult<()> {
+    {
+        let mut value: bool = false;
+        let len = std::mem::size_of_val(&value) as u64;
+        issue_mallctl_args(
+            "prof.active",
+            std::ptr::null_mut(),
+            std::ptr::null_mut(),
+            &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(())
+}