pprof-alloc 0.2.0

Allocation profiling and Linux memory telemetry for Rust services.
Documentation
use crate::stats::malloc_info;
use crate::stats::malloc_info::Error;
use crate::stats::malloc_info::info::{Malloc, SystemType, TotalType};
use prometheus_client::collector::Collector;
use prometheus_client::encoding::DescriptorEncoder;
use prometheus_client::metrics::gauge::ConstGauge;

pub fn info() -> Result<MallocInfo, Error> {
	malloc_info::malloc_info().map(MallocInfo)
}

#[cfg(target_os = "linux")]
#[cfg(target_env = "gnu")]
pub fn malloc_trim() {
	unsafe {
		let _ = libc::malloc_trim(0usize);
	}
}

#[derive(Debug)]
pub struct MallocInfo(pub Malloc);

impl MallocInfo {
	fn total_by_type(&self, r#type: TotalType) -> u64 {
		self
			.0
			.total
			.iter()
			.filter(|t| t.r#type == r#type)
			.map(|t| t.size as u64)
			.sum()
	}

	pub fn system_max(&self) -> u64 {
		self
			.0
			.system
			.iter()
			.find(|s| s.r#type == SystemType::Max)
			.map(|s| s.size as u64)
			.unwrap_or_default()
	}
	pub fn system_current(&self) -> u64 {
		self
			.0
			.system
			.iter()
			.find(|s| s.r#type == SystemType::Current)
			.map(|s| s.size as u64)
			.unwrap_or_default()
	}
	pub fn total(&self) -> u64 {
		self.0.total.iter().map(|t| t.size as u64).sum()
	}

	pub fn free_bytes(&self) -> u64 {
		self.total_by_type(TotalType::Fast) + self.total_by_type(TotalType::Rest)
	}

	pub fn mmap_bytes(&self) -> u64 {
		self.total_by_type(TotalType::Mmap)
	}

	pub fn in_use_bytes(&self) -> u64 {
		self
			.system_current()
			.saturating_sub(self.free_bytes())
			.saturating_add(self.mmap_bytes())
	}

	pub fn heaps(&self) -> u64 {
		self.0.heaps.len() as u64
	}
}

#[derive(Debug, Clone)]
pub struct PrometheusCollector {}

impl PrometheusCollector {
	pub fn register(registry: &mut prometheus_client::registry::Registry) {
		registry.register_collector(Box::new(Self {}))
	}
}

impl Collector for PrometheusCollector {
	fn encode(&self, mut encoder: DescriptorEncoder) -> Result<(), std::fmt::Error> {
		use prometheus_client::encoding::EncodeMetric;
		let Ok(s) = info() else {
			return Ok(());
		};
		let mut encode = |v: u64, n: &'static str, d: &str| {
			let metric = ConstGauge::new(v);
			let metric_encoder = encoder.encode_descriptor(n, d, None, metric.metric_type())?;
			metric.encode(metric_encoder)?;
			Ok(())
		};
		encode(s.system_max(), "malloc_max", "total peak memory")?;
		encode(s.system_current(), "malloc_current", "total current memory")?;
		encode(s.heaps(), "malloc_heaps", "current heaps used")?;
		Ok(())
	}
}