mz_prof/
jemalloc.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//! Utilities for getting jemalloc statistics, as well as exporting them as metrics.
17
18use std::time::Duration;
19
20use jemalloc_pprof::JemallocProfCtl;
21use mz_ore::cast::CastFrom;
22use mz_ore::metric;
23use mz_ore::metrics::{MetricsRegistry, UIntGauge};
24use pprof_util::ProfStartTime;
25use tikv_jemalloc_ctl::{epoch, raw, stats};
26
27#[allow(non_upper_case_globals)]
28#[unsafe(export_name = "malloc_conf")]
29pub static malloc_conf: &[u8] = b"prof:true,prof_active:true,lg_prof_sample:19\0";
30
31#[derive(Copy, Clone, Debug)]
32pub struct JemallocProfMetadata {
33    pub start_time: Option<ProfStartTime>,
34}
35
36// See stats.{allocated, active, ...} in http://jemalloc.net/jemalloc.3.html for details
37pub struct JemallocStats {
38    pub active: usize,
39    pub allocated: usize,
40    pub metadata: usize,
41    pub resident: usize,
42    pub retained: usize,
43}
44
45pub trait JemallocProfCtlExt {
46    fn dump_stats(&mut self, json_format: bool) -> anyhow::Result<String>;
47    fn stats(&self) -> anyhow::Result<JemallocStats>;
48
49    /// Pauses allocation sampling by clearing jemalloc's `prof.active` mallctl.
50    ///
51    /// Unlike [`JemallocProfCtl::deactivate`], this leaves `prof.reset`
52    /// untouched, so the profile accumulated so far and the profiling metadata
53    /// (start time, sample rate) survive. Pair with [`resume`](Self::resume) to
54    /// keep extending the same profile. Use this to briefly stop sampling, for
55    /// example while capturing a CPU profile, without discarding the heap
56    /// profile collected so far. A caller that intends to start a fresh profile
57    /// wants [`JemallocProfCtl::deactivate`] instead.
58    ///
59    /// Takes `&self` deliberately: pausing does not touch the tracked metadata,
60    /// only the global mallctl.
61    fn pause(&self) -> anyhow::Result<()>;
62
63    /// Resumes allocation sampling by setting jemalloc's `prof.active` mallctl.
64    ///
65    /// The counterpart to [`pause`](Self::pause). Sampling continues into the
66    /// profile that `pause` preserved.
67    fn resume(&self) -> anyhow::Result<()>;
68}
69
70impl JemallocProfCtlExt for JemallocProfCtl {
71    fn dump_stats(&mut self, json_format: bool) -> anyhow::Result<String> {
72        // Try to avoid allocations within `stats_print`
73        let mut buf = Vec::with_capacity(1 << 22);
74        let mut options = tikv_jemalloc_ctl::stats_print::Options::default();
75        options.json_format = json_format;
76        tikv_jemalloc_ctl::stats_print::stats_print(&mut buf, options)?;
77        Ok(String::from_utf8(buf)?)
78    }
79
80    fn stats(&self) -> anyhow::Result<JemallocStats> {
81        JemallocStats::get()
82    }
83
84    fn pause(&self) -> anyhow::Result<()> {
85        // SAFETY: "prof.active" is documented as writable and taking a bool:
86        // http://jemalloc.net/jemalloc.3.html#prof.active
87        unsafe { raw::write(b"prof.active\0", false) }?;
88        Ok(())
89    }
90
91    fn resume(&self) -> anyhow::Result<()> {
92        // SAFETY: "prof.active" is documented as writable and taking a bool:
93        // http://jemalloc.net/jemalloc.3.html#prof.active
94        unsafe { raw::write(b"prof.active\0", true) }?;
95        Ok(())
96    }
97}
98
99impl JemallocStats {
100    pub fn get() -> anyhow::Result<JemallocStats> {
101        epoch::advance()?;
102        Ok(JemallocStats {
103            active: stats::active::read()?,
104            allocated: stats::allocated::read()?,
105            metadata: stats::metadata::read()?,
106            resident: stats::resident::read()?,
107            retained: stats::retained::read()?,
108        })
109    }
110}
111
112/// Metrics for jemalloc.
113pub struct JemallocMetrics {
114    pub active: UIntGauge,
115    pub allocated: UIntGauge,
116    pub metadata: UIntGauge,
117    pub resident: UIntGauge,
118    pub retained: UIntGauge,
119}
120
121impl JemallocMetrics {
122    /// Registers the metrics into the provided metrics registry, and spawns
123    /// a task to keep the metrics up to date.
124    // `async` indicates that the Tokio runtime context is required.
125    #[allow(clippy::unused_async)]
126    pub async fn register_into(registry: &MetricsRegistry) {
127        let m = JemallocMetrics {
128            active: registry.register(metric!(
129                name: "jemalloc_active",
130                help: "Total number of bytes in active pages allocated by the application",
131            )),
132            allocated: registry.register(metric!(
133                name: "jemalloc_allocated",
134                help: "Total number of bytes allocated by the application",
135            )),
136            metadata: registry.register(metric!(
137                name: "jemalloc_metadata",
138                help: "Total number of bytes dedicated to metadata.",
139            )),
140            resident: registry.register(metric!(
141                name: "jemalloc_resident",
142                help: "Maximum number of bytes in physically resident data pages mapped",
143            )),
144            retained: registry.register(metric!(
145                name: "jemalloc_retained",
146                help: "Total number of bytes in virtual memory mappings",
147            )),
148        };
149
150        mz_ore::task::spawn(|| "jemalloc_stats_update", async move {
151            let mut interval = tokio::time::interval(Duration::from_secs(10));
152            interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
153            loop {
154                interval.tick().await;
155                if let Err(e) = m.update() {
156                    tracing::warn!("Error while updating jemalloc stats: {}", e);
157                }
158            }
159        });
160    }
161
162    fn update(&self) -> anyhow::Result<()> {
163        let s = JemallocStats::get()?;
164        self.active.set(u64::cast_from(s.active));
165        self.allocated.set(u64::cast_from(s.allocated));
166        self.metadata.set(u64::cast_from(s.metadata));
167        self.resident.set(u64::cast_from(s.resident));
168        self.retained.set(u64::cast_from(s.retained));
169        Ok(())
170    }
171}
mz_prof/jemalloc.rs

mz_prof/
jemalloc.rs
