add memory status

Signed-off-by: Calvin Neo <[email protected]>
pingcap · Dec 9, 2024 · 2dcda8a · 2dcda8a
1 parent da7e90c
commit 2dcda8a
Show file tree

Hide file tree

Showing 3 changed files with 110 additions and 2 deletions.
diff --git 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<unsafe extern "C" fn(*mut c_void, *const i8)>,
+        cbopaque: *mut c_void,
+        opts: *const i8,
+    );
 }
 
 #[allow(unused_variables)]
@@ -109,3 +116,92 @@ pub fn get_allocate() -> u64 {
 pub fn get_deallocate() -> u64 {
     issue_mallctl("thread.deallocated")
 }
+
+use std::{
+    ffi::{c_char, c_void, CStr},
+    ptr,
+};
+struct CaptureContext {
+    buffer: Mutex<String>,
+}
+
+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);
+    }
+}
+
+pub fn get_malloc_stats() -> String {
+    let context = CaptureContext {
+        buffer: Mutex::new(String::new()),
+    };
+
+    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,
+                    ptr::null(),
+                );
+                #[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,
+                    ptr::null(),
+                );
+            }
+        }
+
+        #[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,
+                    ptr::null(),
+                );
+            }
+            #[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,
+                        ptr::null(),
+                    );
+                }
+            }
+        }
+    }
+
+    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
@@ -63,7 +63,7 @@ use tokio::{
     sync::oneshot::{self, Receiver, Sender},
 };
 use tokio_openssl::SslStream;
-use vendored_utils::jeprof_purge_arena;
+use vendored_utils::{jeprof_memory_status, jeprof_purge_arena};
 
 use crate::status_server::profile::set_prof_active;
 
@@ -249,6 +249,11 @@ where
         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("");
@@ -805,6 +810,9 @@ where
                             (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
@@ -1,6 +1,6 @@
 // Copyright 2024 TiKV Project Authors. Licensed under Apache-2.0.
 
-use proxy_ffi::jemalloc_utils::{issue_mallctl, 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<()> {
@@ -126,3 +126,7 @@ pub fn jeprof_purge_arena() {
     }
     info!("jeprof_purge_arena purge {} arenas done", narenas);
 }
+
+pub fn jeprof_memory_status() -> String {
+    return get_malloc_stats();
+}