jemalloc_pprof/
lib.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
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 in the LICENSE file at the
6// root of this repository, or online at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16#![doc = include_str!("../README.md")]
17
18use std::ffi::CString;
19
20use std::io::BufReader;
21use std::sync::Arc;
22use std::time::Instant;
23
24use libc::size_t;
25use once_cell::sync::Lazy;
26
27use mappings::MAPPINGS;
28use tempfile::NamedTempFile;
29use tikv_jemalloc_ctl::raw;
30use tokio::sync::Mutex;
31
32use util::parse_jeheap;
33#[cfg(feature = "flamegraph")]
34pub use util::FlamegraphOptions;
35pub use util::{BuildId, Mapping, ProfStartTime, StackProfile, StackProfileIter, WeightedStack};
36
37/// Activate jemalloc profiling.
38pub async fn activate_jemalloc_profiling() {
39    let Some(ctl) = PROF_CTL.as_ref() else {
40        tracing::warn!("jemalloc profiling is disabled and cannot be activated");
41        return;
42    };
43
44    let mut ctl = ctl.lock().await;
45    if ctl.activated() {
46        return;
47    }
48
49    match ctl.activate() {
50        Ok(()) => tracing::info!("jemalloc profiling activated"),
51        Err(err) => tracing::warn!("could not activate jemalloc profiling: {err}"),
52    }
53}
54
55/// Deactivate jemalloc profiling.
56pub async fn deactivate_jemalloc_profiling() {
57    let Some(ctl) = PROF_CTL.as_ref() else {
58        return; // jemalloc not enabled
59    };
60
61    let mut ctl = ctl.lock().await;
62    if !ctl.activated() {
63        return;
64    }
65
66    match ctl.deactivate() {
67        Ok(()) => tracing::info!("jemalloc profiling deactivated"),
68        Err(err) => tracing::warn!("could not deactivate jemalloc profiling: {err}"),
69    }
70}
71
72/// Per-process singleton for controlling jemalloc profiling.
73pub static PROF_CTL: Lazy<Option<Arc<Mutex<JemallocProfCtl>>>> =
74    Lazy::new(|| JemallocProfCtl::get().map(|ctl| Arc::new(Mutex::new(ctl))));
75
76/// Metadata about a jemalloc heap profiler.
77#[derive(Copy, Clone, Debug)]
78pub struct JemallocProfMetadata {
79    pub start_time: Option<ProfStartTime>,
80}
81
82/// A handle to control jemalloc profiling.
83#[derive(Debug)]
84pub struct JemallocProfCtl {
85    md: JemallocProfMetadata,
86}
87
88impl JemallocProfCtl {
89    // Creates and returns the global singleton.
90    fn get() -> Option<Self> {
91        // SAFETY: "opt.prof" is documented as being readable and returning a bool:
92        // http://jemalloc.net/jemalloc.3.html#opt.prof
93        let prof_enabled: bool = unsafe { raw::read(b"opt.prof\0") }.unwrap();
94        if prof_enabled {
95            // SAFETY: "opt.prof_active" is documented as being readable and returning a bool:
96            // http://jemalloc.net/jemalloc.3.html#opt.prof_active
97            let prof_active: bool = unsafe { raw::read(b"opt.prof_active\0") }.unwrap();
98            let start_time = if prof_active {
99                Some(ProfStartTime::TimeImmemorial)
100            } else {
101                None
102            };
103            let md = JemallocProfMetadata { start_time };
104            Some(Self { md })
105        } else {
106            None
107        }
108    }
109
110    /// Returns the base 2 logarithm of the sample rate (average interval, in bytes, between allocation samples).
111    pub fn lg_sample(&self) -> size_t {
112        // SAFETY: "prof.lg_sample" is documented as being readable and returning size_t:
113        // https://jemalloc.net/jemalloc.3.html#opt.lg_prof_sample
114        unsafe { raw::read(b"prof.lg_sample\0") }.unwrap()
115    }
116
117    /// Returns the metadata of the profiler.
118    pub fn get_md(&self) -> JemallocProfMetadata {
119        self.md
120    }
121
122    /// Returns whether the profiler is active.
123    pub fn activated(&self) -> bool {
124        self.md.start_time.is_some()
125    }
126
127    /// Activate the profiler and if unset, set the start time to the current time.
128    pub fn activate(&mut self) -> Result<(), tikv_jemalloc_ctl::Error> {
129        // SAFETY: "prof.active" is documented as being writable and taking a bool:
130        // http://jemalloc.net/jemalloc.3.html#prof.active
131        unsafe { raw::write(b"prof.active\0", true) }?;
132        if self.md.start_time.is_none() {
133            self.md.start_time = Some(ProfStartTime::Instant(Instant::now()));
134        }
135        Ok(())
136    }
137
138    /// Deactivate the profiler.
139    pub fn deactivate(&mut self) -> Result<(), tikv_jemalloc_ctl::Error> {
140        // SAFETY: "prof.active" is documented as being writable and taking a bool:
141        // http://jemalloc.net/jemalloc.3.html#prof.active
142        unsafe { raw::write(b"prof.active\0", false) }?;
143        let rate = self.lg_sample();
144        // SAFETY: "prof.reset" is documented as being writable and taking a size_t:
145        // http://jemalloc.net/jemalloc.3.html#prof.reset
146        unsafe { raw::write(b"prof.reset\0", rate) }?;
147
148        self.md.start_time = None;
149        Ok(())
150    }
151
152    /// Dump a profile into a temporary file and return it.
153    pub fn dump(&mut self) -> anyhow::Result<std::fs::File> {
154        let f = NamedTempFile::new()?;
155        let path = CString::new(f.path().as_os_str().as_encoded_bytes()).unwrap();
156
157        // SAFETY: "prof.dump" is documented as being writable and taking a C string as input:
158        // http://jemalloc.net/jemalloc.3.html#prof.dump
159        unsafe { raw::write(b"prof.dump\0", path.as_ptr()) }?;
160        Ok(f.into_file())
161    }
162
163    /// Dump a profile and return the parsed [`StackProfile`].
164    pub fn dump_profile(&mut self) -> anyhow::Result<StackProfile> {
165        let f = self.dump()?;
166        let dump_reader = BufReader::new(f);
167        parse_jeheap(dump_reader, MAPPINGS.as_deref())
168    }
169
170    /// Dump a profile in pprof format (gzipped protobuf) and
171    /// return a buffer with its contents.
172    pub fn dump_pprof(&mut self) -> anyhow::Result<Vec<u8>> {
173        let profile = self.dump_profile()?;
174        let pprof = profile.to_pprof(("inuse_space", "bytes"), ("space", "bytes"), None);
175        Ok(pprof)
176    }
177
178    /// Dump a profile flamegraph in SVG format.
179    #[cfg(feature = "flamegraph")]
180    pub fn dump_flamegraph(&mut self) -> anyhow::Result<Vec<u8>> {
181        let mut opts = FlamegraphOptions::default();
182        opts.title = "inuse_space".to_string();
183        opts.count_name = "bytes".to_string();
184        self.dump_flamegraph_with_options(&mut opts)
185    }
186
187    /// Dump a profile flamegraph in SVG format with the given options.
188    #[cfg(feature = "flamegraph")]
189    pub fn dump_flamegraph_with_options(
190        &mut self,
191        opts: &mut FlamegraphOptions,
192    ) -> anyhow::Result<Vec<u8>> {
193        let profile = self.dump_profile()?;
194        profile.to_flamegraph(opts)
195    }
196}
jemalloc_pprof/lib.rs

jemalloc_pprof/
lib.rs
