Skip to main content

pprof_alloc/stats/
malloc.rs

1use crate::stats::malloc_info;
2use crate::stats::malloc_info::Error;
3use crate::stats::malloc_info::info::{Malloc, SystemType, TotalType};
4use prometheus_client::collector::Collector;
5use prometheus_client::encoding::DescriptorEncoder;
6use prometheus_client::metrics::gauge::ConstGauge;
7
8pub fn info() -> Result<MallocInfo, Error> {
9	malloc_info::malloc_info().map(MallocInfo)
10}
11
12#[cfg(target_os = "linux")]
13#[cfg(target_env = "gnu")]
14pub fn malloc_trim() {
15	unsafe {
16		let _ = libc::malloc_trim(0usize);
17	}
18}
19
20#[derive(Debug)]
21pub struct MallocInfo(pub Malloc);
22
23impl MallocInfo {
24	fn total_by_type(&self, r#type: TotalType) -> u64 {
25		self
26			.0
27			.total
28			.iter()
29			.filter(|t| t.r#type == r#type)
30			.map(|t| t.size as u64)
31			.sum()
32	}
33
34	pub fn system_max(&self) -> u64 {
35		self
36			.0
37			.system
38			.iter()
39			.find(|s| s.r#type == SystemType::Max)
40			.map(|s| s.size as u64)
41			.unwrap_or_default()
42	}
43	pub fn system_current(&self) -> u64 {
44		self
45			.0
46			.system
47			.iter()
48			.find(|s| s.r#type == SystemType::Current)
49			.map(|s| s.size as u64)
50			.unwrap_or_default()
51	}
52	pub fn total(&self) -> u64 {
53		self.0.total.iter().map(|t| t.size as u64).sum()
54	}
55
56	pub fn free_bytes(&self) -> u64 {
57		self.total_by_type(TotalType::Fast) + self.total_by_type(TotalType::Rest)
58	}
59
60	pub fn mmap_bytes(&self) -> u64 {
61		self.total_by_type(TotalType::Mmap)
62	}
63
64	pub fn in_use_bytes(&self) -> u64 {
65		self
66			.system_current()
67			.saturating_sub(self.free_bytes())
68			.saturating_add(self.mmap_bytes())
69	}
70
71	pub fn heaps(&self) -> u64 {
72		self.0.heaps.len() as u64
73	}
74}
75
76#[derive(Debug, Clone)]
77pub struct PrometheusCollector {}
78
79impl PrometheusCollector {
80	pub fn register(registry: &mut prometheus_client::registry::Registry) {
81		registry.register_collector(Box::new(Self {}))
82	}
83}
84
85impl Collector for PrometheusCollector {
86	fn encode(&self, mut encoder: DescriptorEncoder) -> Result<(), std::fmt::Error> {
87		use prometheus_client::encoding::EncodeMetric;
88		let Ok(s) = info() else {
89			return Ok(());
90		};
91		let mut encode = |v: u64, n: &'static str, d: &str| {
92			let metric = ConstGauge::new(v);
93			let metric_encoder = encoder.encode_descriptor(n, d, None, metric.metric_type())?;
94			metric.encode(metric_encoder)?;
95			Ok(())
96		};
97		encode(s.system_max(), "malloc_max", "total peak memory")?;
98		encode(s.system_current(), "malloc_current", "total current memory")?;
99		encode(s.heaps(), "malloc_heaps", "current heaps used")?;
100		Ok(())
101	}
102}