Skip to main content

metrics_procession/
iter.rs

1//! This module is responsible for creating iterators from a [`crate::Procession`]
2use std::sync::OnceLock;
3
4use metrics::Key;
5use serde::{
6    Deserialize, Serialize,
7    ser::{SerializeMap, SerializeSeq},
8};
9use time::{Duration, OffsetDateTime};
10
11/// Only used in cases of an emergency, when a [`metrics::Key`] can somehow be lost when
12/// attempting to create a [`Metric`]
13static EMPTY_KEY: OnceLock<Key> = OnceLock::new();
14
15use crate::{
16    chunk::Chunk,
17    event::{Entry, Event},
18    procession::Procession,
19};
20
21/// A single event cloned out of the [Procession], this representation will
22/// allocation the strings needed to represent the value w/o holding a reference
23/// the time [Procession] itself. This type can be serialized and deserialized
24/// and represents and "owned" version of the [MetricRef] type.
25#[derive(Debug, Clone, Serialize, Deserialize)]
26pub struct Metric {
27    pub when: OffsetDateTime,
28    pub event: Entry,
29    pub key: String,
30    pub labels: Vec<(String, String)>,
31}
32
33/// A single event borrowed from the [Procession], this representation
34/// will not cause any additional allocations and can be serialized, the
35/// timestamp is re-calculated as part of the construction but no other
36/// computation should occur.
37#[derive(Debug)]
38pub struct MetricRef<'a> {
39    /// The time this event occurred
40    pub when: OffsetDateTime,
41    /// The value emitted for the key
42    pub event: Entry,
43    /// The key and labels provided by the metrics crate
44    pub key: &'a Key,
45}
46
47impl Serialize for MetricRef<'_> {
48    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
49    where
50        S: serde::Serializer,
51    {
52        let mut m = serializer.serialize_map(Some(4))?;
53        m.serialize_entry("when", &self.when)?;
54        m.serialize_entry("event", &self.event)?;
55        m.serialize_entry("key", &self.key.name())?;
56        m.serialize_entry("labels", &LabelsSet(self.key))?;
57        m.end()
58    }
59}
60
61/// Helper for serializing/deserializing the key type
62struct LabelsSet<'a>(&'a Key);
63
64impl Serialize for LabelsSet<'_> {
65    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
66    where
67        S: serde::Serializer,
68    {
69        let labels = self.0.labels();
70        let mut s = serializer.serialize_seq(Some(labels.len()))?;
71        for label in labels {
72            s.serialize_element(&(label.key(), label.value()))?;
73        }
74        s.end()
75    }
76}
77
78/// An iterator that will clone values out of the source [`Procession`]
79///
80/// warning: this will re-allocate all of the [`String`]s from the [`metrics::Key`]
81/// type potentially many times.
82pub struct MetricsIterator<'a>(MetricsRefIterator<'a>);
83
84impl<'a> From<MetricsRefIterator<'a>> for MetricsIterator<'a> {
85    fn from(value: MetricsRefIterator<'a>) -> Self {
86        Self(value)
87    }
88}
89
90impl<'a> From<&'a Procession> for MetricsRefIterator<'a> {
91    fn from(value: &'a Procession) -> Self {
92        Self {
93            stream: value,
94            chunk_index: 0,
95            event_index: 0,
96        }
97    }
98}
99impl<'a> From<&'a Procession> for MetricsIterator<'a> {
100    fn from(value: &'a Procession) -> Self {
101        Self(MetricsRefIterator::from(value))
102    }
103}
104
105impl Iterator for MetricsIterator<'_> {
106    type Item = Metric;
107    fn next(&mut self) -> Option<Self::Item> {
108        let MetricRef { when, event, key } = self.0.next()?;
109        Some(Metric {
110            when,
111            event,
112            key: key.name().to_string(),
113            labels: key
114                .labels()
115                .map(|l| (l.key().to_string(), l.value().to_string()))
116                .collect(),
117        })
118    }
119}
120
121/// An iterator that will borrow values from the owning [`Procession`], unlike the [`MetricsIterator`]
122/// this will not perform any reallocations but can be safely `collect`ed, as long as the underlying
123/// [`Procession`] is not dropped, and serialized
124pub struct MetricsRefIterator<'a> {
125    stream: &'a Procession,
126    chunk_index: usize,
127    event_index: usize,
128}
129
130impl<'a> Iterator for MetricsRefIterator<'a> {
131    type Item = MetricRef<'a>;
132
133    fn next(&mut self) -> Option<Self::Item> {
134        let (event, chunk) = self.get_next_event()?;
135        let when = chunk.reference_time + Duration::milliseconds(event.ms as i64);
136        let Some(key) = self.stream.labels.0.iter().find_map(|(k, v)| {
137            if *v == event.label {
138                return Some(k);
139            }
140            None
141        }) else {
142            return Some(MetricRef {
143                when,
144                event: event.entry,
145                key: EMPTY_KEY.get_or_init(|| Key::from_name("")),
146            });
147        };
148        Some(MetricRef {
149            when,
150            event: event.entry,
151            key,
152        })
153    }
154}
155
156impl<'a> MetricsRefIterator<'a> {
157    /// This method will do the majority of the work needed by the `next` implementation
158    /// above. The returned [`Event`] represents the correct value that should come
159    /// next in the series but since it only contains the millisecond count since its owning
160    /// [`Chunk`]'s `reference_time` we will also return the correct [`Chunk`]
161    ///
162    /// This method will also handle the index management for the calculation of the next event
163    /// we should emit. If the current chunk is exhausted, it will reset the `event_index` and
164    /// increment the `chunk_index`, otherwise it will increment the `event_index` only
165    fn get_next_event<'s, 'r>(&'s mut self) -> Option<(&'r Event, &'r Chunk)>
166    where
167        'a: 'r,
168    {
169        let mut chunk = self.stream.chunks.get(self.chunk_index)?;
170        if let Some(event) = chunk.events.get(self.event_index) {
171            self.event_index += 1;
172            return Some((event, chunk));
173        }
174        self.chunk_index += 1;
175        self.event_index = 0;
176        chunk = self.stream.chunks.get(self.chunk_index)?;
177        let ret = chunk.events.get(self.event_index)?;
178        self.event_index += 1;
179        Some((ret, chunk))
180    }
181}
182
183impl PartialEq<MetricRef<'_>> for Metric {
184    fn eq(&self, other: &MetricRef) -> bool {
185        self.when.eq(&other.when)
186            && self.event.eq(&other.event)
187            && self.key.eq(other.key.name())
188            && self.labels.len() == other.key.labels().len()
189            && self
190                .labels
191                .iter()
192                .all(|(k, v)| other.key.labels().any(|l| k == l.key() && v == l.value()))
193    }
194}
195
196impl PartialEq<Metric> for MetricRef<'_> {
197    fn eq(&self, other: &Metric) -> bool {
198        other.eq(self)
199    }
200}
201
202#[cfg(test)]
203mod tests {
204    use metrics::{Key, Label};
205    use time::{Date, Time};
206
207    use crate::{event::Op, label_set::LabelSet};
208
209    use super::*;
210
211    #[test]
212    fn iter_works_as_expected() {
213        let time_stream = build_test_stream();
214        let iter = MetricsIterator::from(&time_stream);
215        let flattened: Vec<Metric> = iter.collect();
216        insta::assert_json_snapshot!(flattened);
217    }
218
219    #[test]
220    fn iter_ref_and_iter_match() {
221        let time_stream = build_test_stream();
222        let met_refs = MetricsRefIterator::from(&time_stream);
223        let mets = MetricsIterator::from(&time_stream);
224        for (l, r) in mets.zip(met_refs) {
225            assert_eq!(l, r);
226        }
227    }
228
229    fn build_test_stream() -> Procession {
230        let start = OffsetDateTime::new_utc(
231            Date::from_calendar_date(2025, time::Month::January, 1).unwrap(),
232            Time::from_hms(0, 0, 0).unwrap(),
233        );
234        let mut labels = LabelSet([].into_iter().collect());
235        let k1 = Key::from_name("no-labels");
236        let mut raw_labels = Vec::new();
237        raw_labels.push(labels.ensure_key(&k1));
238        let k2 = Key::from_parts("one-label", vec![Label::new("label", "value")]);
239        raw_labels.push(labels.ensure_key(&k2));
240        let k3 = Key::from_parts(
241            "two-labels",
242            vec![
243                Label::new("3label1", "value1"),
244                Label::new("3label2", "value2"),
245            ],
246        );
247        raw_labels.push(labels.ensure_key(&k3));
248        let k4 = Key::from_parts(
249            "three-labels",
250            vec![
251                Label::new("4label1", "value1"),
252                Label::new("4label2", "value2"),
253                Label::new("4label3", "value3"),
254            ],
255        );
256        raw_labels.push(labels.ensure_key(&k4));
257        let k5 = Key::from_parts(
258            "three-labels",
259            vec![
260                Label::new("5label1", "value1"),
261                Label::new("5label2", "value2"),
262                Label::new("5label3", "value3"),
263                Label::new("5label4", "value4"),
264            ],
265        );
266        raw_labels.push(labels.ensure_key(&k5));
267
268        let streams = (0..128)
269            .map(|v| {
270                let reference_time = start + Duration::minutes(v);
271                let events = (0..128)
272                    .map(|v| Event {
273                        entry: Entry::Counter {
274                            value: 1,
275                            op: Op::Add,
276                        },
277                        ms: v as u16,
278                        label: raw_labels[v % 5],
279                    })
280                    .collect();
281                Chunk {
282                    reference_time,
283                    events,
284                }
285            })
286            .collect();
287        Procession {
288            labels,
289            chunks: streams,
290        }
291    }
292}