risingwave_common_heap_profiling/
profiler.rs

1// Copyright 2025 RisingWave Labs
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
15use std::ffi::CString;
16use std::fs;
17use std::path::Path;
18
19use parking_lot::Once;
20use risingwave_common::config::HeapProfilingConfig;
21use risingwave_common::util::resource_util;
22use tikv_jemalloc_ctl::{opt as jemalloc_opt, prof as jemalloc_prof};
23use tokio::time::{self, Duration};
24
25use super::AUTO_DUMP_SUFFIX;
26
27/// `HeapProfiler` automatically triggers heap profiling when memory usage is higher than the threshold.
28///
29/// To use it, both jemalloc's `opt.prof` and RisingWave's config `heap_profiling.enable_auto` must be set to true.
30pub struct HeapProfiler {
31    config: HeapProfilingConfig,
32    threshold_auto_dump_heap_profile: usize,
33    jemalloc_dump_mib: jemalloc_prof::dump_mib,
34    /// If jemalloc profiling is enabled
35    opt_prof: bool,
36}
37
38impl HeapProfiler {
39    /// # Arguments
40    ///
41    /// `total_memory` must be the total available memory for the process.
42    /// It will be compared with the process resident memory.
43    pub fn new(total_memory: usize, config: HeapProfilingConfig) -> Self {
44        let threshold_auto_dump_heap_profile =
45            (total_memory as f64 * config.threshold_auto as f64) as usize;
46        let jemalloc_dump_mib = jemalloc_prof::dump::mib().unwrap();
47        let opt_prof = jemalloc_opt::prof::read().unwrap();
48
49        Self {
50            config,
51            threshold_auto_dump_heap_profile,
52            jemalloc_dump_mib,
53            opt_prof,
54        }
55    }
56
57    fn auto_dump_heap_prof(&self) {
58        let time_prefix = chrono::Local::now().format("%Y-%m-%d-%H-%M-%S");
59        let file_name = format!("{}.{}", time_prefix, AUTO_DUMP_SUFFIX);
60
61        let file_path = Path::new(&self.config.dir)
62            .join(&file_name)
63            .to_str()
64            .expect("file path is not valid utf8")
65            .to_owned();
66        let file_path_c = CString::new(file_path).expect("0 byte in file path");
67
68        // FIXME(yuhao): `unsafe` here because `jemalloc_dump_mib.write` requires static lifetime
69        if let Err(e) = self
70            .jemalloc_dump_mib
71            .write(unsafe { &*(file_path_c.as_c_str() as *const _) })
72        {
73            tracing::warn!("Auto Jemalloc dump heap file failed! {:?}", e);
74        } else {
75            tracing::info!("Successfully dumped heap profile to {}", file_name);
76        }
77    }
78
79    /// Start the daemon task of auto heap profiling.
80    pub fn start(self) {
81        if !self.config.enable_auto || !self.opt_prof {
82            tracing::info!("Auto memory dump is disabled.");
83            return;
84        }
85
86        static START: Once = Once::new();
87        START.call_once(|| {
88            fs::create_dir_all(&self.config.dir).unwrap();
89            tokio::spawn(async move {
90                let mut interval = time::interval(Duration::from_millis(500));
91                let mut prev_used_memory_bytes = 0;
92                loop {
93                    interval.tick().await;
94                    let cur_used_memory_bytes = resource_util::memory::total_memory_used_bytes();
95
96                    // Dump heap profile when memory usage is crossing the threshold.
97                    if cur_used_memory_bytes > self.threshold_auto_dump_heap_profile
98                        && prev_used_memory_bytes <= self.threshold_auto_dump_heap_profile
99                    {
100                        self.auto_dump_heap_prof();
101                    }
102                    prev_used_memory_bytes = cur_used_memory_bytes;
103                }
104            });
105        })
106    }
107}
risingwave_common_heap_profiling/profiler.rs

risingwave_common_heap_profiling/
profiler.rs
