broxus_util/
alloc.rs

1use std::ffi::{c_void, CString};
2
3pub type Allocator = tikv_jemallocator::Jemalloc;
4
5pub const fn allocator() -> Allocator {
6    Allocator {}
7}
8
9pub fn set_jemalloc_param<T>(name: &str, mut value: T) {
10    let name_buffer = CString::new(name).unwrap();
11
12    let res = unsafe {
13        tikv_jemalloc_sys::mallctl(
14            name_buffer.as_ptr(),
15            std::ptr::null_mut(),
16            std::ptr::null_mut(),
17            &mut value as *mut T as *mut c_void,
18            std::mem::size_of::<T>(),
19        )
20    };
21
22    if res != 0 {
23        log::error!("Failed to set {name}: {}", errno::Errno(res));
24    }
25}
26
27#[cfg(all(feature = "alloc-profiling", unix))]
28pub mod profiling {
29    use std::ffi::CString;
30    use std::os::raw::c_char;
31    use std::os::unix::ffi::OsStrExt;
32    use std::path::Path;
33
34    pub use tikv_jemalloc_ctl::Error;
35    use tikv_jemalloc_ctl::{epoch, stats};
36
37    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
38    #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
39    pub struct JemallocStats {
40        pub allocated: u64,
41        pub active: u64,
42        pub metadata: u64,
43        pub resident: u64,
44        pub mapped: u64,
45        pub retained: u64,
46        pub dirty: u64,
47        pub fragmentation: u64,
48    }
49
50    #[cfg(feature = "metrics")]
51    macro_rules! set_metrics {
52    ($($metric_name:expr => $metric_value:expr),* $(,)?) => {
53        $(
54            metrics::gauge!($metric_name).set($metric_value as f64);
55        )*
56    };}
57
58    #[cfg(feature = "metrics")]
59    pub async fn allocator_metrics_loop() {
60        loop {
61            let Ok(s) = fetch_stats() else { break };
62            set_metrics!(
63                "jemalloc_allocated_bytes" => s.allocated,
64                "jemalloc_active_bytes" => s.active,
65                "jemalloc_metadata_bytes" => s.metadata,
66                "jemalloc_resident_bytes" => s.resident,
67                "jemalloc_mapped_bytes" => s.mapped,
68                "jemalloc_retained_bytes" => s.retained,
69                "jemalloc_dirty_bytes" => s.dirty,
70                "jemalloc_fragmentation_bytes" => s.fragmentation,
71            );
72            tokio::time::sleep(std::time::Duration::from_secs(5)).await;
73        }
74    }
75
76    pub fn fetch_stats() -> Result<JemallocStats, Error> {
77        // Stats are cached. Need to advance epoch to refresh.
78        epoch::advance()?;
79
80        Ok(JemallocStats {
81            allocated: stats::allocated::read()? as u64,
82            active: stats::active::read()? as u64,
83            metadata: stats::metadata::read()? as u64,
84            resident: stats::resident::read()? as u64,
85            mapped: stats::mapped::read()? as u64,
86            retained: stats::retained::read()? as u64,
87            dirty: (stats::resident::read()?
88                .saturating_sub(stats::active::read()?)
89                .saturating_sub(stats::metadata::read()?)) as u64,
90            fragmentation: (stats::active::read()?.saturating_sub(stats::allocated::read()?))
91                as u64,
92        })
93    }
94
95    const PROF_ACTIVE: &[u8] = b"prof.active\0";
96    const PROF_DUMP: &[u8] = b"prof.dump\0";
97
98    pub fn start() -> Result<(), Error> {
99        log::info!("starting profiler");
100        unsafe { tikv_jemalloc_ctl::raw::update(PROF_ACTIVE, true)? };
101        Ok(())
102    }
103
104    pub fn stop() -> Result<(), Error> {
105        log::info!("stopping profiler");
106        unsafe { tikv_jemalloc_ctl::raw::update(PROF_ACTIVE, false)? };
107        Ok(())
108    }
109
110    /// Dump the profile to the `path`.
111    pub fn dump<P>(path: P) -> Result<(), DumpError>
112    where
113        P: AsRef<Path>,
114    {
115        let mut bytes = CString::new(path.as_ref().as_os_str().as_bytes())
116            .map_err(|_| DumpError::InvalidPath)?
117            .into_bytes_with_nul();
118
119        let ptr = bytes.as_mut_ptr() as *mut c_char;
120        let res = unsafe { tikv_jemalloc_ctl::raw::write(PROF_DUMP, ptr) };
121        match res {
122            Ok(_) => {
123                log::info!("saved the profiling dump to {:?}", path.as_ref());
124                Ok(())
125            }
126            Err(e) => {
127                log::error!(
128                    "failed to dump the profiling info to {:?}: {e:?}",
129                    path.as_ref()
130                );
131                Err(DumpError::JemallocError(e.to_string()))
132            }
133        }
134    }
135
136    #[derive(thiserror::Error, Debug)]
137    pub enum DumpError {
138        #[error("invalid path to the dump")]
139        InvalidPath,
140        #[error("failed to dump the profiling info: {0}")]
141        JemallocError(String),
142    }
143}
144
145// run with `MALLOC_CONF="prof:true" cargo test --ignored`
146#[cfg(all(test, feature = "alloc-profiling"))]
147mod test {
148    const OPT_PROF: &[u8] = b"opt.prof\0";
149
150    fn is_profiling_on() -> bool {
151        match unsafe { tikv_jemalloc_ctl::raw::read(OPT_PROF) } {
152            Err(e) => {
153                // Shouldn't be possible since mem-profiling is set
154                panic!("is_profiling_on: {e:?}");
155            }
156            Ok(prof) => prof,
157        }
158    }
159
160    #[test]
161    #[ignore]
162    fn test_profile() {
163        use super::*;
164        use tempfile::TempDir;
165
166        let dir = TempDir::new().unwrap();
167        let path = dir.path().join("profile.txt");
168
169        profiling::start().unwrap();
170        assert!(is_profiling_on());
171
172        profiling::dump(path.as_path()).unwrap();
173        profiling::stop().unwrap();
174    }
175}
broxus_util/alloc.rs

broxus_util/
alloc.rs
