metrics_procession/
label_set.rs

1// this module uses a hashmap with the `Key` type which shouldn't really
2// change during its lifetime
3#![allow(clippy::mutable_key_type)]
4use std::collections::BTreeMap;
5
6use metrics::Key;
7use serde::{Deserialize, Serialize, de::Visitor, ser::SerializeSeq};
8
9/// A set of labels mapping from the original [`metrics::Key`] to a unique identifier
10/// and will be used to lookup what identifier to use when recording metrics events
11#[derive(Debug, PartialEq, Clone, Default)]
12pub struct LabelSet(pub BTreeMap<Key, u16>);
13
14impl LabelSet {
15    /// Get the identifier for the provided key
16    pub fn get(&self, key: &Key) -> Option<u16> {
17        self.0.get(key).copied()
18    }
19
20    /// ensure the [`metrics::Key`] is in the set, inserting a clone if not
21    /// already present, returning the correct identifier for the provided key
22    pub fn ensure_key(&mut self, key: &Key) -> u16 {
23        if let Some(v) = self.0.get(key) {
24            return *v;
25        }
26        let v = u16::try_from(self.0.len()).unwrap_or_else(|_| {
27            eprintln!("too many labels!!!");
28            u16::MAX
29        });
30        self.0.insert(key.clone(), v);
31        v
32    }
33}
34
35/// Helper struct for serializing the [`LabelSet`] set to avoid needing to re-allocate the
36/// strings owned by the [`metrics::Key`] type along with its value to make it possible
37/// to deserialize a serialized `LabelSet` with the correct key<->id mapping
38#[derive(Debug, Serialize)]
39struct SerKey<'a> {
40    key_name: &'a str,
41    labels: Vec<SerLabel<'a>>,
42    value: u16,
43}
44
45/// Helper struct for serializing just the key-value pair owned by a [`metrics::Label`]
46/// this will use a tuple of string references
47#[derive(Debug, Serialize, Deserialize)]
48struct SerLabel<'a>(&'a str, &'a str);
49impl<'a> From<SerLabel<'a>> for metrics::Label {
50    fn from(value: SerLabel<'a>) -> Self {
51        metrics::Label::new(value.0.to_string(), value.1.to_string())
52    }
53}
54
55/// Helper struct for serializing just the key-value pair owned by a [`metrics::Label`]
56/// this will use a tuple of string references
57struct SerLabels<'a>(Vec<SerLabel<'a>>);
58impl metrics::IntoLabels for SerLabels<'_> {
59    fn into_labels(self) -> Vec<metrics::Label> {
60        self.0.into_iter().map(metrics::Label::from).collect()
61    }
62}
63
64impl Serialize for LabelSet {
65    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
66    where
67        S: serde::Serializer,
68    {
69        let mut m = serializer.serialize_seq(Some(self.0.len()))?;
70        for (k, v) in self.0.iter() {
71            let ser_key = SerKey {
72                key_name: k.name(),
73                labels: k.labels().map(|l| SerLabel(l.key(), l.value())).collect(),
74                value: *v,
75            };
76            m.serialize_element(&ser_key)?;
77        }
78        m.end()
79    }
80}
81
82impl<'de> Deserialize<'de> for LabelSet {
83    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
84    where
85        D: serde::Deserializer<'de>,
86    {
87        struct LabelSetVisitor;
88
89        impl<'de> Visitor<'de> for LabelSetVisitor {
90            type Value = LabelSet;
91
92            fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
93                formatter.write_str("sequence of label set entries")
94            }
95            fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
96            where
97                A: serde::de::SeqAccess<'de>,
98            {
99                let mut ret = BTreeMap::new();
100                while let Some(element) = seq.next_element::<SerKey<'de>>()? {
101                    let SerKey {
102                        key_name,
103                        labels,
104                        value,
105                    } = element;
106                    let key = Key::from_parts(key_name.to_string(), SerLabels(labels));
107                    ret.insert(key, value);
108                }
109                Ok(LabelSet(ret))
110            }
111        }
112
113        deserializer.deserialize_seq(LabelSetVisitor)
114    }
115}
116
117impl<'de> Deserialize<'de> for SerKey<'de> {
118    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
119    where
120        D: serde::Deserializer<'de>,
121    {
122        struct EntryVisitor;
123        impl<'de> Visitor<'de> for EntryVisitor {
124            type Value = SerKey<'de>;
125
126            fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
127                formatter.write_str("map of label data")
128            }
129
130            fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
131            where
132                A: serde::de::MapAccess<'de>,
133            {
134                let mut key_name: Option<&str> = None;
135                let mut labels: Option<Vec<SerLabel<'_>>> = None;
136                let mut value: Option<u16> = None;
137                while let Some((k, v)) = map.next_entry()? {
138                    match (k, v) {
139                        ("key_name", SerKeyValue::KeyName(name)) => key_name = Some(name),
140                        ("value", SerKeyValue::Value(i)) => value = Some(i),
141                        ("labels", SerKeyValue::Labels(ls)) => labels = Some(ls),
142                        _ => {}
143                    }
144                    if key_name.is_some() && labels.is_some() && value.is_some() {
145                        break;
146                    }
147                }
148                Ok(SerKey {
149                    key_name: key_name.ok_or_else(|| {
150                        serde::de::Error::custom("key_name missing from label set entry")
151                    })?,
152                    labels: labels.ok_or_else(|| {
153                        serde::de::Error::custom("key_name missing from label set entry")
154                    })?,
155                    value: value.ok_or_else(|| {
156                        serde::de::Error::custom("key_name missing from label set entry")
157                    })?,
158                })
159            }
160        }
161
162        #[derive(Debug, Deserialize)]
163        #[serde(untagged)]
164        enum SerKeyValue<'a> {
165            KeyName(&'a str),
166            Labels(Vec<SerLabel<'a>>),
167            Value(u16),
168        }
169        deserializer.deserialize_map(EntryVisitor)
170    }
171}
metrics_procession/label_set.rs

metrics_procession/
label_set.rs
