use crate::telemetry::metrics::{Counter, Gauge, GaugeExt, Register, raw};
use commonware_allocator::mimalloc::{
self, Stats,
metrics::{Kind, METRICS},
};
use std::sync::{LazyLock, atomic::Ordering};
use tracing::warn;
static GLOBAL_IS_MIMALLOC: LazyLock<bool> = LazyLock::new(|| {
let global = mimalloc::is_global_allocator();
if !global {
warn!("mimalloc metrics enabled but mimalloc is not the global allocator");
}
global
});
type Read = fn(&Stats) -> u64;
pub struct Metrics {
gauges: Vec<(Gauge, Read)>,
counters: Vec<(Counter, Read)>,
}
impl Metrics {
pub fn init(registry: &mut impl Register) -> Option<Self> {
if !*GLOBAL_IS_MIMALLOC {
return None;
}
let mut registry = registry.sub_registry("allocator");
let mut gauges = Vec::new();
let mut counters = Vec::new();
for metric in &METRICS {
match metric.kind {
Kind::Gauge => {
let gauge = registry.register(metric.name, metric.help, raw::Gauge::default());
gauges.push((gauge, metric.read));
}
Kind::Counter => {
let counter =
registry.register(metric.name, metric.help, raw::Counter::default());
counters.push((counter, metric.read));
}
}
}
let mut metrics = Self { gauges, counters };
metrics.update();
Some(metrics)
}
pub fn update(&mut self) {
let Some(stats) = Stats::read() else {
return;
};
for (gauge, read) in &self.gauges {
let _ = gauge.try_set(read(&stats));
}
for (counter, read) in &self.counters {
counter.inner().fetch_max(read(&stats), Ordering::Relaxed);
}
}
}