profile/
memory_usage.rs

1//! Like [`std::time::Instant`], but for memory.
2//!
3//! Measures the total size of all currently allocated objects.
4use std::fmt;
5
6use cfg_if::cfg_if;
7
8#[derive(Copy, Clone)]
9pub struct MemoryUsage {
10    pub allocated: Bytes,
11}
12
13impl fmt::Display for MemoryUsage {
14    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
15        self.allocated.fmt(f)
16    }
17}
18
19impl std::ops::Sub for MemoryUsage {
20    type Output = MemoryUsage;
21    fn sub(self, rhs: MemoryUsage) -> MemoryUsage {
22        MemoryUsage { allocated: self.allocated - rhs.allocated }
23    }
24}
25
26impl MemoryUsage {
27    pub fn now() -> MemoryUsage {
28        cfg_if! {
29            if #[cfg(all(feature = "jemalloc", not(target_env = "msvc")))] {
30                jemalloc_ctl::epoch::advance().unwrap();
31                MemoryUsage {
32                    allocated: Bytes(jemalloc_ctl::stats::allocated::read().unwrap() as isize),
33                }
34            } else if #[cfg(all(target_os = "linux", target_env = "gnu"))] {
35                memusage_linux()
36            } else if #[cfg(windows)] {
37                // There doesn't seem to be an API for determining heap usage, so we try to
38                // approximate that by using the Commit Charge value.
39
40                use windows_sys::Win32::System::{Threading::*, ProcessStatus::*};
41                use std::mem::MaybeUninit;
42
43                let proc = unsafe { GetCurrentProcess() };
44                let mut mem_counters = MaybeUninit::uninit();
45                let cb = size_of::<PROCESS_MEMORY_COUNTERS>();
46                let ret = unsafe { GetProcessMemoryInfo(proc, mem_counters.as_mut_ptr(), cb as u32) };
47                assert!(ret != 0);
48
49                let usage = unsafe { mem_counters.assume_init().PagefileUsage };
50                MemoryUsage { allocated: Bytes(usage as isize) }
51            } else {
52                MemoryUsage { allocated: Bytes(0) }
53            }
54        }
55    }
56}
57
58#[cfg(all(target_os = "linux", target_env = "gnu", not(feature = "jemalloc")))]
59fn memusage_linux() -> MemoryUsage {
60    // Linux/glibc has 2 APIs for allocator introspection that we can use: mallinfo and mallinfo2.
61    // mallinfo uses `int` fields and cannot handle memory usage exceeding 2 GB.
62    // mallinfo2 is very recent, so its presence needs to be detected at runtime.
63    // Both are abysmally slow.
64
65    use std::sync::atomic::{AtomicUsize, Ordering};
66
67    static MALLINFO2: AtomicUsize = AtomicUsize::new(1);
68
69    let mut mallinfo2 = MALLINFO2.load(Ordering::Relaxed);
70    if mallinfo2 == 1 {
71        mallinfo2 = unsafe { libc::dlsym(libc::RTLD_DEFAULT, c"mallinfo2".as_ptr()) } as usize;
72        // NB: races don't matter here, since they'll always store the same value
73        MALLINFO2.store(mallinfo2, Ordering::Relaxed);
74    }
75
76    if mallinfo2 == 0 {
77        // mallinfo2 does not exist, use mallinfo.
78        let alloc = unsafe { libc::mallinfo() }.uordblks as isize;
79        MemoryUsage { allocated: Bytes(alloc) }
80    } else {
81        let mallinfo2: extern "C" fn() -> libc::mallinfo2 =
82            unsafe { std::mem::transmute(mallinfo2) };
83        let alloc = mallinfo2().uordblks as isize;
84        MemoryUsage { allocated: Bytes(alloc) }
85    }
86}
87
88#[derive(Default, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy)]
89pub struct Bytes(isize);
90
91impl Bytes {
92    pub fn new(bytes: isize) -> Bytes {
93        Bytes(bytes)
94    }
95}
96
97impl Bytes {
98    pub fn megabytes(self) -> isize {
99        self.0 / 1024 / 1024
100    }
101}
102
103impl fmt::Display for Bytes {
104    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
105        let bytes = self.0;
106        let mut value = bytes;
107        let mut suffix = "b";
108        if value.abs() > 4096 {
109            value /= 1024;
110            suffix = "kb";
111            if value.abs() > 4096 {
112                value /= 1024;
113                suffix = "mb";
114            }
115        }
116        f.pad(&format!("{value}{suffix}"))
117    }
118}
119
120impl std::ops::AddAssign<usize> for Bytes {
121    fn add_assign(&mut self, x: usize) {
122        self.0 += x as isize;
123    }
124}
125
126impl std::ops::Sub for Bytes {
127    type Output = Bytes;
128    fn sub(self, rhs: Bytes) -> Bytes {
129        Bytes(self.0 - rhs.0)
130    }
131}
profile/memory_usage.rs

profile/
memory_usage.rs
