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::{CString, c_char};
18use std::io::BufReader;
19use std::path::PathBuf;
20
21use error::{
22    ActivateProfSnafu, BuildTempPathSnafu, DeactivateProfSnafu, DumpProfileDataSnafu,
23    OpenTempFileSnafu, ProfilingNotEnabledSnafu, ReadOptProfSnafu, ReadProfActiveSnafu,
24};
25use jemalloc_pprof_mappings::MAPPINGS;
26use jemalloc_pprof_utils::{FlamegraphOptions, StackProfile, parse_jeheap};
27use snafu::{ResultExt, ensure};
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";
34const PROF_ACTIVE: &[u8] = b"prof.active\0";
35const PROF_GDUMP: &[u8] = b"prof.gdump\0";
36
37pub async fn dump_profile() -> Result<Vec<u8>> {
38    ensure!(is_prof_enabled()?, ProfilingNotEnabledSnafu);
39    let tmp_path = tempfile::tempdir().map_err(|_| {
40        BuildTempPathSnafu {
41            path: std::env::temp_dir(),
42        }
43        .build()
44    })?;
45
46    let mut path_buf = PathBuf::from(tmp_path.path());
47    path_buf.push("greptimedb.hprof");
48
49    let path = path_buf
50        .to_str()
51        .ok_or_else(|| BuildTempPathSnafu { path: &path_buf }.build())?
52        .to_string();
53
54    let mut bytes = CString::new(path.as_str())
55        .map_err(|_| BuildTempPathSnafu { path: &path_buf }.build())?
56        .into_bytes_with_nul();
57
58    {
59        // #safety: we always expect a valid temp file path to write profiling data to.
60        let ptr = bytes.as_mut_ptr() as *mut c_char;
61        unsafe {
62            tikv_jemalloc_ctl::raw::write(PROF_DUMP, ptr)
63                .context(DumpProfileDataSnafu { path: path_buf })?
64        }
65    }
66
67    let mut f = tokio::fs::File::open(path.as_str())
68        .await
69        .context(OpenTempFileSnafu { path: &path })?;
70    let mut buf = vec![];
71    let _ = f
72        .read_to_end(&mut buf)
73        .await
74        .context(OpenTempFileSnafu { path })?;
75    Ok(buf)
76}
77
78async fn dump_profile_to_stack_profile() -> Result<StackProfile> {
79    let profile = dump_profile().await?;
80    let profile = BufReader::new(profile.as_slice());
81    parse_jeheap(profile, MAPPINGS.as_deref()).context(ParseJeHeapSnafu)
82}
83
84pub async fn dump_pprof() -> Result<Vec<u8>> {
85    let profile = dump_profile_to_stack_profile().await?;
86    let pprof = profile.to_pprof(("inuse_space", "bytes"), ("space", "bytes"), None);
87    Ok(pprof)
88}
89
90pub async fn dump_flamegraph() -> Result<Vec<u8>> {
91    let profile = dump_profile_to_stack_profile().await?;
92    let mut opts = FlamegraphOptions::default();
93    opts.title = "inuse_space".to_string();
94    opts.count_name = "bytes".to_string();
95    let flamegraph = profile.to_flamegraph(&mut opts).context(FlamegraphSnafu)?;
96    Ok(flamegraph)
97}
98
99pub fn activate_heap_profile() -> Result<()> {
100    ensure!(is_prof_enabled()?, ProfilingNotEnabledSnafu);
101    unsafe {
102        tikv_jemalloc_ctl::raw::update(PROF_ACTIVE, true).context(ActivateProfSnafu)?;
103    }
104    Ok(())
105}
106
107pub fn deactivate_heap_profile() -> Result<()> {
108    ensure!(is_prof_enabled()?, ProfilingNotEnabledSnafu);
109    unsafe {
110        tikv_jemalloc_ctl::raw::update(PROF_ACTIVE, false).context(DeactivateProfSnafu)?;
111    }
112    Ok(())
113}
114
115pub fn is_heap_profile_active() -> Result<bool> {
116    unsafe { Ok(tikv_jemalloc_ctl::raw::read::<bool>(PROF_ACTIVE).context(ReadProfActiveSnafu)?) }
117}
118
119fn is_prof_enabled() -> Result<bool> {
120    // safety: OPT_PROF variable, if present, is always a boolean value.
121    Ok(unsafe { tikv_jemalloc_ctl::raw::read::<bool>(OPT_PROF).context(ReadOptProfSnafu)? })
122}
123
124pub fn set_gdump_active(active: bool) -> Result<()> {
125    ensure!(is_prof_enabled()?, ProfilingNotEnabledSnafu);
126    unsafe {
127        tikv_jemalloc_ctl::raw::update(PROF_GDUMP, active).context(error::UpdateGdumpSnafu)?;
128    }
129    Ok(())
130}
131
132pub fn is_gdump_active() -> Result<bool> {
133    // safety: PROF_GDUMP, if present, is a boolean value.
134    unsafe { Ok(tikv_jemalloc_ctl::raw::read::<bool>(PROF_GDUMP).context(error::ReadGdumpSnafu)?) }
135}
common_mem_prof/jemalloc.rs

common_mem_prof/
jemalloc.rs
