feat: mem prof can gen flamegraph directly (#6073)

* feat: mem-prof

* fix: use enum&update how to

Cargo.lock

@@ -2285,8 +2285,11 @@ dependencies = [
 name = "common-mem-prof"
 version = "0.15.0"
 dependencies = [
+ "anyhow",
  "common-error",
  "common-macro",
+ "mappings",
+ "pprof_util",
  "snafu 0.8.5",
  "tempfile",
  "tikv-jemalloc-ctl",
@@ -3979,6 +3982,25 @@ dependencies = [
  "syn 2.0.100",
 ]
 
+[[package]]
+name = "env_filter"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "186e05a59d4c50738528153b83b0b0194d3a29507dfec16eccd4b342903397d0"
+dependencies = [
+ "log",
+]
+
+[[package]]
+name = "env_logger"
+version = "0.11.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "13c863f0904021b108aa8b2f55046443e6b1ebde8fd4a15c399893aae4fa069f"
+dependencies = [
+ "env_filter",
+ "log",
+]
+
 [[package]]
 name = "equator"
 version = "0.2.2"
@@ -5648,6 +5670,28 @@ dependencies = [
  "str_stack",
 ]
 
+[[package]]
+name = "inferno"
+version = "0.12.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2094aecddc672e902cd773bad7071542f63641e01e9187c3bba4b43005e837e9"
+dependencies = [
+ "ahash 0.8.11",
+ "clap 4.5.19",
+ "crossbeam-channel",
+ "crossbeam-utils",
+ "dashmap",
+ "env_logger",
+ "indexmap 2.9.0",
+ "itoa",
+ "log",
+ "num-format",
+ "once_cell",
+ "quick-xml 0.37.5",
+ "rgb",
+ "str_stack",
+]
+
 [[package]]
 name = "influxdb_line_protocol"
 version = "0.1.0"
@@ -6553,6 +6597,19 @@ dependencies = [
  "thiserror 1.0.64",
 ]
 
+[[package]]
+name = "mappings"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e434981a332777c2b3062652d16a55f8e74fa78e6b1882633f0d77399c84fc2a"
+dependencies = [
+ "anyhow",
+ "libc",
+ "once_cell",
+ "pprof_util",
+ "tracing",
+]
+
 [[package]]
 name = "match_cfg"
 version = "0.1.0"
@@ -8651,7 +8708,7 @@ dependencies = [
  "cfg-if",
  "criterion 0.5.1",
  "findshlibs",
- "inferno",
+ "inferno 0.11.21",
  "libc",
  "log",
  "nix 0.26.4",
@@ -8668,6 +8725,21 @@ dependencies = [
  "thiserror 1.0.64",
 ]
 
+[[package]]
+name = "pprof_util"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9fa015c78eed2130951e22c58d2095849391e73817ab2e74f71b0b9f63dd8416"
+dependencies = [
+ "anyhow",
+ "backtrace",
+ "flate2",
+ "inferno 0.12.2",
+ "num",
+ "paste",
+ "prost 0.13.5",
+]
+
 [[package]]
 name = "ppv-lite86"
 version = "0.2.20"
@@ -9252,6 +9324,15 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "quick-xml"
+version = "0.37.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "331e97a1af0bf59823e6eadffe373d7b27f485be8748f71471c662c1f269b7fb"
+dependencies = [
+ "memchr",
+]
+
 [[package]]
 name = "quinn"
 version = "0.11.5"

docs/how-to/how-to-profile-memory.md

@@ -44,6 +44,10 @@ Dump memory profiling data through HTTP API:
 
 ```bash
 curl -X POST localhost:4000/debug/prof/mem > greptime.hprof
+# or output flamegraph directly
+curl -X POST "localhost:4000/debug/prof/mem?output=flamegraph" > greptime.svg
+# or output pprof format
+curl -X POST "localhost:4000/debug/prof/mem?output=proto" > greptime.pprof
 ```
 
 You can periodically dump profiling data and compare them to find the delta memory usage.

src/common/mem-prof/Cargo.toml

@@ -16,6 +16,12 @@ tokio.workspace = true
 
 [target.'cfg(not(windows))'.dependencies]
 tikv-jemalloc-ctl = { version = "0.6", features = ["use_std", "stats"] }
+jemalloc-pprof-utils = { version = "0.7", package = "pprof_util", features = [
+    "flamegraph",
+    "symbolize",
+] } # for parsing jemalloc prof dump
+jemalloc-pprof-mappings = { version = "0.7", package = "mappings" } # for get the name of functions in the prof dump
+anyhow = "1"
 
 [target.'cfg(not(windows))'.dependencies.tikv-jemalloc-sys]
 features = ["stats", "profiling", "unprefixed_malloc_on_supported_platforms"]

src/common/mem-prof/src/error.rs

@@ -30,12 +30,25 @@ pub enum Error {
 
     #[snafu(display("Memory profiling is not supported"))]
     ProfilingNotSupported,
+
+    #[snafu(display("Failed to parse jeheap profile: {}", err))]
+    ParseJeHeap {
+        #[snafu(source)]
+        err: anyhow::Error,
+    },
+
+    #[snafu(display("Failed to dump profile data to flamegraph: {}", err))]
+    Flamegraph {
+        #[snafu(source)]
+        err: anyhow::Error,
+    },
 }
 
 impl ErrorExt for Error {
     fn status_code(&self) -> StatusCode {
         match self {
             Error::Internal { source } => source.status_code(),
+            Error::ParseJeHeap { .. } | Error::Flamegraph { .. } => StatusCode::Internal,
             Error::ProfilingNotSupported => StatusCode::Unsupported,
         }
     }

src/common/mem-prof/src/jemalloc.rs

@@ -15,16 +15,19 @@
 mod error;
 
 use std::ffi::{c_char, CString};
+use std::io::BufReader;
 use std::path::PathBuf;
 
 use error::{
     BuildTempPathSnafu, DumpProfileDataSnafu, OpenTempFileSnafu, ProfilingNotEnabledSnafu,
     ReadOptProfSnafu,
 };
+use jemalloc_pprof_mappings::MAPPINGS;
+use jemalloc_pprof_utils::{parse_jeheap, FlamegraphOptions, StackProfile};
 use snafu::{ensure, ResultExt};
 use tokio::io::AsyncReadExt;
 
-use crate::error::Result;
+use crate::error::{FlamegraphSnafu, ParseJeHeapSnafu, Result};
 
 const PROF_DUMP: &[u8] = b"prof.dump\0";
 const OPT_PROF: &[u8] = b"opt.prof\0";
@@ -70,6 +73,26 @@ pub async fn dump_profile() -> Result<Vec<u8>> {
     Ok(buf)
 }
 
+async fn dump_profile_to_stack_profile() -> Result<StackProfile> {
+    let profile = dump_profile().await?;
+    let profile = BufReader::new(profile.as_slice());
+    parse_jeheap(profile, MAPPINGS.as_deref()).context(ParseJeHeapSnafu)
+}
+
+pub async fn dump_pprof() -> Result<Vec<u8>> {
+    let profile = dump_profile_to_stack_profile().await?;
+    let pprof = profile.to_pprof(("inuse_space", "bytes"), ("space", "bytes"), None);
+    Ok(pprof)
+}
+
+pub async fn dump_flamegraph() -> Result<Vec<u8>> {
+    let profile = dump_profile_to_stack_profile().await?;
+    let mut opts = FlamegraphOptions::default();
+    opts.title = "inuse_space".to_string();
+    opts.count_name = "bytes".to_string();
+    let flamegraph = profile.to_flamegraph(&mut opts).context(FlamegraphSnafu)?;
+    Ok(flamegraph)
+}
 fn is_prof_enabled() -> Result<bool> {
     // safety: OPT_PROF variable, if present, is always a boolean value.
     Ok(unsafe { tikv_jemalloc_ctl::raw::read::<bool>(OPT_PROF).context(ReadOptProfSnafu)? })

src/common/mem-prof/src/lib.rs

@@ -17,7 +17,7 @@ pub mod error;
 #[cfg(not(windows))]
 mod jemalloc;
 #[cfg(not(windows))]
-pub use jemalloc::dump_profile;
+pub use jemalloc::{dump_flamegraph, dump_pprof, dump_profile};
 
 #[cfg(windows)]
 pub async fn dump_profile() -> error::Result<Vec<u8>> {

src/servers/src/http/mem_prof.rs

@@ -12,22 +12,47 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use axum::extract::Query;
 use axum::http::StatusCode;
 use axum::response::IntoResponse;
+use serde::{Deserialize, Serialize};
+
+/// Output format.
+#[derive(Debug, Default, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum Output {
+    /// google’s pprof format report in protobuf.
+    Proto,
+    /// jeheap text format. Define by jemalloc.
+    #[default]
+    Text,
+    /// svg flamegraph.
+    Flamegraph,
+}
+
+#[derive(Default, Serialize, Deserialize, Debug)]
+#[serde(default)]
+pub struct MemPprofQuery {
+    output: Output,
+}
 
 #[cfg(feature = "mem-prof")]
 #[axum_macros::debug_handler]
-pub async fn mem_prof_handler() -> crate::error::Result<impl IntoResponse> {
+pub async fn mem_prof_handler(
+    Query(req): Query<MemPprofQuery>,
+) -> crate::error::Result<impl IntoResponse> {
     use snafu::ResultExt;
 
     use crate::error::DumpProfileDataSnafu;
 
-    Ok((
-        StatusCode::OK,
-        common_mem_prof::dump_profile()
-            .await
-            .context(DumpProfileDataSnafu)?,
-    ))
+    let dump = match req.output {
+        Output::Proto => common_mem_prof::dump_pprof().await,
+        Output::Text => common_mem_prof::dump_profile().await,
+        Output::Flamegraph => common_mem_prof::dump_flamegraph().await,
+    }
+    .context(DumpProfileDataSnafu)?;
+
+    Ok((StatusCode::OK, dump))
 }
 
 #[cfg(not(feature = "mem-prof"))]