radiate_core/stats/
set.rs

1use crate::{
2    Metric, MetricUpdate,
3    stats::{Meta, Tag, TagType, fmt, metric_fields},
4};
5use radiate_error::RadiateError;
6use radiate_expr::{ProjectExpr, SelectOp};
7use radiate_utils::{AnyValue, SmallStr};
8#[cfg(feature = "serde")]
9use serde::{Deserialize, Serialize};
10use std::{
11    collections::HashMap,
12    fmt::{Debug, Display},
13};
14
15#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
16#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
17#[repr(transparent)]
18pub(crate) struct MetricIdx(u32);
19
20impl MetricIdx {
21    #[inline(always)]
22    pub(crate) const fn new(idx: u32) -> Self {
23        MetricIdx(idx)
24    }
25
26    #[inline(always)]
27    pub(crate) const fn as_usize(self) -> usize {
28        self.0 as usize
29    }
30}
31
32#[derive(PartialEq)]
33pub struct MetricSetSummary {
34    pub metrics: usize,
35    pub updates: f32,
36}
37
38#[derive(Clone, Default, PartialEq)]
39pub struct MetricSet {
40    metrics: Vec<Metric>,
41    name_lookup: HashMap<SmallStr, MetricIdx>,
42    meta: Meta,
43}
44
45impl MetricSet {
46    pub fn new() -> Self {
47        MetricSet {
48            metrics: Vec::new(),
49            name_lookup: HashMap::new(),
50            meta: Meta::default(),
51        }
52    }
53
54    pub fn bump(&mut self, generation: usize) {
55        self.meta.generation = generation;
56    }
57
58    pub fn generation(&self) -> usize {
59        self.meta.generation
60    }
61
62    #[inline]
63    pub(crate) fn upsert_at<'a>(&mut self, idx: MetricIdx, update: impl Into<MetricUpdate<'a>>) {
64        let generation = self.meta.generation;
65        let metric = &mut self.metrics[idx.as_usize()];
66
67        metric.set_generation(generation);
68        metric.apply_update(update.into());
69
70        self.meta.update_count += 1;
71    }
72
73    #[inline(always)]
74    pub fn upsert<'a>(&mut self, key: impl AsRef<str>, metric: impl Into<MetricUpdate<'a>>) {
75        let metric_update = metric.into();
76        let idx = self.resolve(&key);
77        self.upsert_at(idx, metric_update);
78    }
79
80    #[inline(always)]
81    pub fn upsert_tagged<'a>(
82        &mut self,
83        key: impl AsRef<str>,
84        metric: impl Into<MetricUpdate<'a>>,
85        tag: TagType,
86    ) {
87        let metric_update = metric.into();
88        let idx = self.resolve(&key);
89        if let Some(metric) = self.metrics.get_mut(idx.as_usize()) {
90            metric.add_tag(tag);
91            self.upsert_at(idx, metric_update);
92        }
93    }
94
95    #[inline(always)]
96    pub fn keys(&self) -> impl Iterator<Item = SmallStr> {
97        self.metrics.iter().map(|m| m.name().clone())
98    }
99
100    #[inline(always)]
101    pub fn replace(&mut self, metric: impl Into<Metric>) {
102        let metric = metric.into();
103        if let Some(&idx) = self.name_lookup.get(metric.name().as_str()) {
104            self.metrics[idx.as_usize()] = metric;
105        } else {
106            let idx = MetricIdx::new(self.metrics.len() as u32);
107            self.name_lookup.insert(metric.name().clone(), idx);
108            self.metrics.push(metric);
109        }
110    }
111
112    #[inline(always)]
113    pub fn iter_tagged(&self, tag: TagType) -> impl Iterator<Item = &Metric> {
114        self.metrics.iter().filter(move |m| m.tags().has(tag))
115    }
116
117    #[inline(always)]
118    pub fn tags(&self) -> impl Iterator<Item = TagType> {
119        self.metrics
120            .iter()
121            .fold(Tag::empty(), |acc, m| acc.union(m.tags()))
122            .into_iter()
123    }
124
125    #[inline(always)]
126    pub fn iter(&self) -> impl Iterator<Item = &Metric> {
127        self.metrics.iter()
128    }
129
130    #[inline(always)]
131    pub fn add(&mut self, metric: Metric) {
132        self.replace(metric);
133    }
134
135    #[inline(always)]
136    pub fn get(&self, name: impl AsRef<str>) -> Option<&Metric> {
137        self.name_lookup
138            .get(name.as_ref())
139            .and_then(|idx| self.metrics.get(idx.as_usize()))
140    }
141
142    #[inline(always)]
143    pub fn clear(&mut self) {
144        for m in &mut self.metrics {
145            m.clear_values();
146        }
147        self.meta.update_count = 0;
148    }
149
150    #[inline(always)]
151    pub fn contains_key(&self, name: impl AsRef<str>) -> bool {
152        self.name_lookup.contains_key(name.as_ref())
153    }
154
155    pub fn remove_samples(&mut self) {
156        for m in &mut self.metrics {
157            if m.tags().has(TagType::Distribution) {
158                m.clear_samples();
159            }
160        }
161    }
162
163    #[inline(always)]
164    pub fn len(&self) -> usize {
165        self.metrics.len()
166    }
167
168    pub fn is_empty(&self) -> bool {
169        self.metrics.is_empty()
170    }
171
172    pub fn summary(&self) -> MetricSetSummary {
173        MetricSetSummary {
174            metrics: self.metrics.len(),
175            updates: self.meta.update_count as f32,
176        }
177    }
178
179    pub fn dashboard(&self) -> String {
180        fmt::render_full(self).unwrap_or_default()
181    }
182
183    /// Resolve a name to a stable [`MetricIdx`], registering an empty metric if
184    /// the name has not been seen before. The returned handle is valid for the
185    /// lifetime of this `MetricSet`.
186    #[inline]
187    fn resolve(&mut self, name: impl AsRef<str>) -> MetricIdx {
188        if let Some(&idx) = self.name_lookup.get(name.as_ref()) {
189            return idx;
190        }
191
192        let idx = MetricIdx::new(self.metrics.len() as u32);
193        let name = SmallStr::from(name.as_ref());
194        self.name_lookup.insert(name.clone(), idx);
195        self.metrics.push(Metric::new(name));
196        idx
197    }
198}
199
200impl<'a> ProjectExpr<'a> for &MetricSet {
201    #[inline]
202    fn select(&'a self, sel: &SelectOp) -> Result<AnyValue<'a>, RadiateError> {
203        (*self).select(sel)
204    }
205}
206
207impl<'a> ProjectExpr<'a> for MetricSet {
208    #[inline]
209    fn select(&'a self, sel: &SelectOp) -> Result<AnyValue<'a>, RadiateError> {
210        match sel {
211            SelectOp::Field(name) => self
212                .get(name)
213                .map(|metric| metric.select(&SelectOp::Field(metric_fields::LAST_VALUE)))
214                .unwrap_or(Ok(AnyValue::Null)),
215            SelectOp::Nested { parent, child } => {
216                if let SelectOp::Field(name) = parent.as_ref()
217                    && let Some(metric) = self.get(name)
218                {
219                    return metric.select(child);
220                }
221
222                Ok(AnyValue::Null)
223            }
224            _ => Ok(AnyValue::Null),
225        }
226    }
227}
228
229impl From<Vec<Metric>> for MetricSet {
230    fn from(metrics: Vec<Metric>) -> Self {
231        let mut by_name = HashMap::with_capacity(metrics.len());
232        for (i, m) in metrics.iter().enumerate() {
233            by_name.insert(m.name().clone(), MetricIdx::new(i as u32));
234        }
235
236        MetricSet {
237            metrics,
238            name_lookup: by_name,
239            meta: Meta::default(),
240        }
241    }
242}
243
244impl From<&[Metric]> for MetricSet {
245    fn from(metrics: &[Metric]) -> Self {
246        Self::from(metrics.to_vec())
247    }
248}
249
250impl<'a, S, T> From<(S, Vec<T>)> for MetricSet
251where
252    S: AsRef<str>,
253    T: Into<MetricUpdate<'a>>,
254{
255    fn from(tuple: (S, Vec<T>)) -> Self {
256        let (name, updates) = tuple;
257        let mut set = MetricSet::new();
258        for update in updates {
259            set.upsert(name.as_ref(), update.into());
260        }
261
262        set
263    }
264}
265
266impl Display for MetricSet {
267    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
268        let summary = self.summary();
269        let out = format!(
270            "[{} metrics, {:.0} updates]",
271            summary.metrics, summary.updates
272        );
273        write!(f, "{out}\n{}", fmt::render_full(self).unwrap_or_default())?;
274        Ok(())
275    }
276}
277
278impl Debug for MetricSet {
279    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
280        writeln!(f, "MetricSet {{")?;
281        writeln!(f, "{}", fmt::render_dashboard(self).unwrap_or_default())?;
282        write!(f, "}}")
283    }
284}
285
286#[cfg(feature = "serde")]
287impl Serialize for MetricSet {
288    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
289    where
290        S: serde::Serializer,
291    {
292        self.metrics.serialize(serializer)
293    }
294}
295
296#[cfg(feature = "serde")]
297impl<'de> Deserialize<'de> for MetricSet {
298    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
299    where
300        D: serde::Deserializer<'de>,
301    {
302        let metrics = Vec::<Metric>::deserialize(deserializer)?;
303        let mut by_name = HashMap::with_capacity(metrics.len());
304        for (i, m) in metrics.iter().enumerate() {
305            by_name.insert(m.name().clone(), MetricIdx::new(i as u32));
306        }
307        Ok(MetricSet {
308            metrics,
309            name_lookup: by_name,
310            meta: Meta::default(),
311        })
312    }
313}
314
315#[cfg(test)]
316mod tests {
317    use super::*;
318
319    #[test]
320    fn resolve_returns_stable_handle() {
321        let mut set = MetricSet::new();
322        let name = SmallStr::from_static("test.metric");
323
324        let idx1 = set.resolve(&name);
325        let idx2 = set.resolve(&name);
326        assert_eq!(idx1, idx2);
327
328        set.upsert_at(idx1, 1.0);
329        set.upsert_at(idx1, 2.0);
330        set.upsert_at(idx1, 3.0);
331
332        let m = set.get(name.as_str()).unwrap();
333        assert_eq!(m.count(), 3);
334        assert_eq!(m.sum(), 6.0);
335    }
336
337    #[test]
338    fn resolve_assigns_sequential_indices() {
339        let mut set = MetricSet::new();
340        let a = set.resolve(&SmallStr::from_static("a"));
341        let b = set.resolve(&SmallStr::from_static("b"));
342        let c = set.resolve(&SmallStr::from_static("c"));
343        assert_eq!(a.as_usize(), 0);
344        assert_eq!(b.as_usize(), 1);
345        assert_eq!(c.as_usize(), 2);
346    }
347}
radiate_core/stats/set.rs

radiate_core/stats/
set.rs
