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, 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
50impl JemallocProfCtlExt for JemallocProfCtl {
51    fn dump_stats(&mut self, json_format: bool) -> anyhow::Result<String> {
52        // Try to avoid allocations within `stats_print`
53        let mut buf = Vec::with_capacity(1 << 22);
54        let mut options = tikv_jemalloc_ctl::stats_print::Options::default();
55        options.json_format = json_format;
56        tikv_jemalloc_ctl::stats_print::stats_print(&mut buf, options)?;
57        Ok(String::from_utf8(buf)?)
58    }
59
60    fn stats(&self) -> anyhow::Result<JemallocStats> {
61        JemallocStats::get()
62    }
63}
64
65impl JemallocStats {
66    pub fn get() -> anyhow::Result<JemallocStats> {
67        epoch::advance()?;
68        Ok(JemallocStats {
69            active: stats::active::read()?,
70            allocated: stats::allocated::read()?,
71            metadata: stats::metadata::read()?,
72            resident: stats::resident::read()?,
73            retained: stats::retained::read()?,
74        })
75    }
76}
77
78/// Metrics for jemalloc.
79pub struct JemallocMetrics {
80    pub active: UIntGauge,
81    pub allocated: UIntGauge,
82    pub metadata: UIntGauge,
83    pub resident: UIntGauge,
84    pub retained: UIntGauge,
85}
86
87impl JemallocMetrics {
88    /// Registers the metrics into the provided metrics registry, and spawns
89    /// a task to keep the metrics up to date.
90    // `async` indicates that the Tokio runtime context is required.
91    #[allow(clippy::unused_async)]
92    pub async fn register_into(registry: &MetricsRegistry) {
93        let m = JemallocMetrics {
94            active: registry.register(metric!(
95                name: "jemalloc_active",
96                help: "Total number of bytes in active pages allocated by the application",
97            )),
98            allocated: registry.register(metric!(
99                name: "jemalloc_allocated",
100                help: "Total number of bytes allocated by the application",
101            )),
102            metadata: registry.register(metric!(
103                name: "jemalloc_metadata",
104                help: "Total number of bytes dedicated to metadata.",
105            )),
106            resident: registry.register(metric!(
107                name: "jemalloc_resident",
108                help: "Maximum number of bytes in physically resident data pages mapped",
109            )),
110            retained: registry.register(metric!(
111                name: "jemalloc_retained",
112                help: "Total number of bytes in virtual memory mappings",
113            )),
114        };
115
116        mz_ore::task::spawn(|| "jemalloc_stats_update", async move {
117            let mut interval = tokio::time::interval(Duration::from_secs(10));
118            interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
119            loop {
120                interval.tick().await;
121                if let Err(e) = m.update() {
122                    tracing::warn!("Error while updating jemalloc stats: {}", e);
123                }
124            }
125        });
126    }
127
128    fn update(&self) -> anyhow::Result<()> {
129        let s = JemallocStats::get()?;
130        self.active.set(u64::cast_from(s.active));
131        self.allocated.set(u64::cast_from(s.allocated));
132        self.metadata.set(u64::cast_from(s.metadata));
133        self.resident.set(u64::cast_from(s.resident));
134        self.retained.set(u64::cast_from(s.retained));
135        Ok(())
136    }
137}