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
32#[cfg(feature = "flamegraph")]
33pub use util::FlamegraphOptions;
34use util::{parse_jeheap, ProfStartTime};
35
36/// Activate jemalloc profiling.
37pub async fn activate_jemalloc_profiling() {
38    let Some(ctl) = PROF_CTL.as_ref() else {
39        tracing::warn!("jemalloc profiling is disabled and cannot be activated");
40        return;
41    };
42
43    let mut ctl = ctl.lock().await;
44    if ctl.activated() {
45        return;
46    }
47
48    match ctl.activate() {
49        Ok(()) => tracing::info!("jemalloc profiling activated"),
50        Err(err) => tracing::warn!("could not activate jemalloc profiling: {err}"),
51    }
52}
53
54/// Deactivate jemalloc profiling.
55pub async fn deactivate_jemalloc_profiling() {
56    let Some(ctl) = PROF_CTL.as_ref() else {
57        return; // jemalloc not enabled
58    };
59
60    let mut ctl = ctl.lock().await;
61    if !ctl.activated() {
62        return;
63    }
64
65    match ctl.deactivate() {
66        Ok(()) => tracing::info!("jemalloc profiling deactivated"),
67        Err(err) => tracing::warn!("could not deactivate jemalloc profiling: {err}"),
68    }
69}
70
71/// Per-process singleton for controlling jemalloc profiling.
72pub static PROF_CTL: Lazy<Option<Arc<Mutex<JemallocProfCtl>>>> =
73    Lazy::new(|| JemallocProfCtl::get().map(|ctl| Arc::new(Mutex::new(ctl))));
74
75/// Metadata about a jemalloc heap profiler.
76#[derive(Copy, Clone, Debug)]
77pub struct JemallocProfMetadata {
78    pub start_time: Option<ProfStartTime>,
79}
80
81/// A handle to control jemalloc profiling.
82#[derive(Debug)]
83pub struct JemallocProfCtl {
84    md: JemallocProfMetadata,
85}
86
87impl JemallocProfCtl {
88    // Creates and returns the global singleton.
89    fn get() -> Option<Self> {
90        // SAFETY: "opt.prof" is documented as being readable and returning a bool:
91        // http://jemalloc.net/jemalloc.3.html#opt.prof
92        let prof_enabled: bool = unsafe { raw::read(b"opt.prof\0") }.unwrap();
93        if prof_enabled {
94            // SAFETY: "opt.prof_active" is documented as being readable and returning a bool:
95            // http://jemalloc.net/jemalloc.3.html#opt.prof_active
96            let prof_active: bool = unsafe { raw::read(b"opt.prof_active\0") }.unwrap();
97            let start_time = if prof_active {
98                Some(ProfStartTime::TimeImmemorial)
99            } else {
100                None
101            };
102            let md = JemallocProfMetadata { start_time };
103            Some(Self { md })
104        } else {
105            None
106        }
107    }
108
109    /// Returns the base 2 logarithm of the sample rate (average interval, in bytes, between allocation samples).
110    pub fn lg_sample(&self) -> size_t {
111        // SAFETY: "prof.lg_sample" is documented as being readable and returning size_t:
112        // https://jemalloc.net/jemalloc.3.html#opt.lg_prof_sample
113        unsafe { raw::read(b"prof.lg_sample\0") }.unwrap()
114    }
115
116    /// Returns the metadata of the profiler.
117    pub fn get_md(&self) -> JemallocProfMetadata {
118        self.md
119    }
120
121    /// Returns whether the profiler is active.
122    pub fn activated(&self) -> bool {
123        self.md.start_time.is_some()
124    }
125
126    /// Activate the profiler and if unset, set the start time to the current time.
127    pub fn activate(&mut self) -> Result<(), tikv_jemalloc_ctl::Error> {
128        // SAFETY: "prof.active" is documented as being writable and taking a bool:
129        // http://jemalloc.net/jemalloc.3.html#prof.active
130        unsafe { raw::write(b"prof.active\0", true) }?;
131        if self.md.start_time.is_none() {
132            self.md.start_time = Some(ProfStartTime::Instant(Instant::now()));
133        }
134        Ok(())
135    }
136
137    /// Deactivate the profiler.
138    pub fn deactivate(&mut self) -> Result<(), tikv_jemalloc_ctl::Error> {
139        // SAFETY: "prof.active" is documented as being writable and taking a bool:
140        // http://jemalloc.net/jemalloc.3.html#prof.active
141        unsafe { raw::write(b"prof.active\0", false) }?;
142        let rate = self.lg_sample();
143        // SAFETY: "prof.reset" is documented as being writable and taking a size_t:
144        // http://jemalloc.net/jemalloc.3.html#prof.reset
145        unsafe { raw::write(b"prof.reset\0", rate) }?;
146
147        self.md.start_time = None;
148        Ok(())
149    }
150
151    /// Dump a profile into a temporary file and return it.
152    pub fn dump(&mut self) -> anyhow::Result<std::fs::File> {
153        let f = NamedTempFile::new()?;
154        let path = CString::new(f.path().as_os_str().as_encoded_bytes()).unwrap();
155
156        // SAFETY: "prof.dump" is documented as being writable and taking a C string as input:
157        // http://jemalloc.net/jemalloc.3.html#prof.dump
158        unsafe { raw::write(b"prof.dump\0", path.as_ptr()) }?;
159        Ok(f.into_file())
160    }
161
162    /// Dump a profile in pprof format (gzipped protobuf) and
163    /// return a buffer with its contents.
164    pub fn dump_pprof(&mut self) -> anyhow::Result<Vec<u8>> {
165        let f = self.dump()?;
166        let dump_reader = BufReader::new(f);
167        let profile = parse_jeheap(dump_reader, MAPPINGS.as_deref())?;
168        let pprof = profile.to_pprof(("inuse_space", "bytes"), ("space", "bytes"), None);
169        Ok(pprof)
170    }
171
172    /// Dump a profile flamegraph in SVG format.
173    #[cfg(feature = "flamegraph")]
174    pub fn dump_flamegraph(&mut self) -> anyhow::Result<Vec<u8>> {
175        let mut opts = FlamegraphOptions::default();
176        opts.title = "inuse_space".to_string();
177        opts.count_name = "bytes".to_string();
178        self.dump_flamegraph_with_options(&mut opts)
179    }
180
181    /// Dump a profile flamegraph in SVG format with the given options.
182    #[cfg(feature = "flamegraph")]
183    pub fn dump_flamegraph_with_options(
184        &mut self,
185        opts: &mut FlamegraphOptions,
186    ) -> anyhow::Result<Vec<u8>> {
187        let f = self.dump()?;
188        let dump_reader = BufReader::new(f);
189        let profile = parse_jeheap(dump_reader, MAPPINGS.as_deref())?;
190        profile.to_flamegraph(opts)
191    }
192}
jemalloc_pprof/lib.rs

jemalloc_pprof/
lib.rs
