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

common_mem_prof/
jemalloc.rs
