143 lines
3.7 KiB
Rust
143 lines
3.7 KiB
Rust
use bevy_ecs::prelude::*;
|
|
use std::{
|
|
collections::HashMap,
|
|
sync::{atomic, Mutex, RwLock},
|
|
};
|
|
|
|
use crate::{entity_registry::EntityRegistry, Stages};
|
|
|
|
#[derive(Default)]
|
|
pub struct Histogram<T>
|
|
where
|
|
T: Copy + Send + Sync,
|
|
{
|
|
inner: Mutex<Vec<T>>,
|
|
}
|
|
|
|
impl<T> Histogram<T>
|
|
where
|
|
T: Copy + Send + Sync,
|
|
{
|
|
#[inline(always)]
|
|
pub fn record(&self, val: T) {
|
|
self.inner.lock().unwrap().push(val);
|
|
}
|
|
}
|
|
|
|
#[derive(Default)]
|
|
pub struct Counter {
|
|
inner: atomic::AtomicU64,
|
|
}
|
|
|
|
impl Counter {
|
|
#[inline(always)]
|
|
pub fn increment(&self, value: u64) {
|
|
self.inner.fetch_add(value, atomic::Ordering::Relaxed);
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
|
pub struct MetricKey {
|
|
entity: Option<Entity>,
|
|
label: &'static str,
|
|
}
|
|
|
|
#[derive(Resource, Default)]
|
|
pub struct Metrics {
|
|
pub(crate) active: bool,
|
|
counters: RwLock<HashMap<MetricKey, Counter>>,
|
|
histograms: RwLock<HashMap<MetricKey, Histogram<u32>>>,
|
|
}
|
|
|
|
impl Metrics {
|
|
pub fn record_histogram(&self, entity: Option<Entity>, label: &'static str, value: u32) {
|
|
if self.active {
|
|
let key = MetricKey { entity, label };
|
|
let r_hist = self.histograms.read().unwrap();
|
|
if let Some(hist) = r_hist.get(&key) {
|
|
hist.record(value);
|
|
} else {
|
|
std::mem::drop(r_hist);
|
|
|
|
let mut histograms = self.histograms.write().unwrap();
|
|
histograms.insert(
|
|
key,
|
|
Histogram {
|
|
inner: vec![value].into(),
|
|
},
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn increment_counter(&self, entity: Option<Entity>, label: &'static str, value: u64) {
|
|
if self.active {
|
|
let key = MetricKey { entity, label };
|
|
let r_counters = self.counters.read().unwrap();
|
|
if let Some(counter) = r_counters.get(&key) {
|
|
counter.increment(value);
|
|
} else {
|
|
std::mem::drop(r_counters);
|
|
|
|
let mut counters = self.counters.write().unwrap();
|
|
counters.insert(
|
|
key,
|
|
Counter {
|
|
inner: value.into(),
|
|
},
|
|
);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(crate) fn consume_metrics(
|
|
world: &World,
|
|
) -> (
|
|
Vec<proxisim_models::dto::metrics::Counter>,
|
|
Vec<proxisim_models::dto::metrics::Histogram>,
|
|
) {
|
|
let metrics = world.resource::<Metrics>();
|
|
let entities = world.resource::<EntityRegistry>();
|
|
|
|
let counters = metrics
|
|
.counters
|
|
.try_write()
|
|
.unwrap()
|
|
.drain()
|
|
.map(|(key, value)| proxisim_models::dto::metrics::Counter {
|
|
entity: key
|
|
.entity
|
|
.as_ref()
|
|
.map(|e| entities.0.get(e).unwrap().clone())
|
|
.unwrap_or(proxisim_models::dto::metrics::EntityInfo::Global),
|
|
value: value.inner.load(atomic::Ordering::Relaxed),
|
|
label: key.label,
|
|
})
|
|
.collect();
|
|
|
|
let histograms = metrics
|
|
.histograms
|
|
.try_write()
|
|
.unwrap()
|
|
.drain()
|
|
.map(|(key, value)| proxisim_models::dto::metrics::Histogram {
|
|
entity: key
|
|
.entity
|
|
.as_ref()
|
|
.map(|e| entities.0.get(e).unwrap().clone())
|
|
.unwrap_or(proxisim_models::dto::metrics::EntityInfo::Global),
|
|
values: value.inner.into_inner().unwrap(),
|
|
label: key.label,
|
|
})
|
|
.collect();
|
|
|
|
(counters, histograms)
|
|
}
|
|
|
|
pub(crate) fn configure(stages: &mut Stages) {
|
|
stages.world.init_resource::<Metrics>();
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod test {}
|