common_mem_prof/
jemalloc.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15mod error;
16
17use std::ffi::{c_char, CString};
18use std::io::BufReader;
19use std::path::PathBuf;
20
21use error::{
22    BuildTempPathSnafu, DumpProfileDataSnafu, OpenTempFileSnafu, ProfilingNotEnabledSnafu,
23    ReadOptProfSnafu,
24};
25use jemalloc_pprof_mappings::MAPPINGS;
26use jemalloc_pprof_utils::{parse_jeheap, FlamegraphOptions, StackProfile};
27use snafu::{ensure, ResultExt};
28use tokio::io::AsyncReadExt;
29
30use crate::error::{FlamegraphSnafu, ParseJeHeapSnafu, Result};
31
32const PROF_DUMP: &[u8] = b"prof.dump\0";
33const OPT_PROF: &[u8] = b"opt.prof\0";
34
35pub async fn dump_profile() -> Result<Vec<u8>> {
36    ensure!(is_prof_enabled()?, ProfilingNotEnabledSnafu);
37    let tmp_path = tempfile::tempdir().map_err(|_| {
38        BuildTempPathSnafu {
39            path: std::env::temp_dir(),
40        }
41        .build()
42    })?;
43
44    let mut path_buf = PathBuf::from(tmp_path.path());
45    path_buf.push("greptimedb.hprof");
46
47    let path = path_buf
48        .to_str()
49        .ok_or_else(|| BuildTempPathSnafu { path: &path_buf }.build())?
50        .to_string();
51
52    let mut bytes = CString::new(path.as_str())
53        .map_err(|_| BuildTempPathSnafu { path: &path_buf }.build())?
54        .into_bytes_with_nul();
55
56    {
57        // #safety: we always expect a valid temp file path to write profiling data to.
58        let ptr = bytes.as_mut_ptr() as *mut c_char;
59        unsafe {
60            tikv_jemalloc_ctl::raw::write(PROF_DUMP, ptr)
61                .context(DumpProfileDataSnafu { path: path_buf })?
62        }
63    }
64
65    let mut f = tokio::fs::File::open(path.as_str())
66        .await
67        .context(OpenTempFileSnafu { path: &path })?;
68    let mut buf = vec![];
69    let _ = f
70        .read_to_end(&mut buf)
71        .await
72        .context(OpenTempFileSnafu { path })?;
73    Ok(buf)
74}
75
76async fn dump_profile_to_stack_profile() -> Result<StackProfile> {
77    let profile = dump_profile().await?;
78    let profile = BufReader::new(profile.as_slice());
79    parse_jeheap(profile, MAPPINGS.as_deref()).context(ParseJeHeapSnafu)
80}
81
82pub async fn dump_pprof() -> Result<Vec<u8>> {
83    let profile = dump_profile_to_stack_profile().await?;
84    let pprof = profile.to_pprof(("inuse_space", "bytes"), ("space", "bytes"), None);
85    Ok(pprof)
86}
87
88pub async fn dump_flamegraph() -> Result<Vec<u8>> {
89    let profile = dump_profile_to_stack_profile().await?;
90    let mut opts = FlamegraphOptions::default();
91    opts.title = "inuse_space".to_string();
92    opts.count_name = "bytes".to_string();
93    let flamegraph = profile.to_flamegraph(&mut opts).context(FlamegraphSnafu)?;
94    Ok(flamegraph)
95}
96fn is_prof_enabled() -> Result<bool> {
97    // safety: OPT_PROF variable, if present, is always a boolean value.
98    Ok(unsafe { tikv_jemalloc_ctl::raw::read::<bool>(OPT_PROF).context(ReadOptProfSnafu)? })
99}
common_mem_prof/jemalloc.rs

common_mem_prof/
jemalloc.rs
