1 use chrono::{DateTime, Utc};
2 use derive_is_enum_variant::is_enum_variant;
3 use serde::{Deserialize, Serialize};
4 use std::collections::{BTreeMap, BTreeSet};
5 
6 #[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
7 pub struct Metric {
8     pub name: String,
9     pub timestamp: Option<DateTime<Utc>>,
10     pub tags: Option<BTreeMap<String, String>>,
11     pub kind: MetricKind,
12     #[serde(flatten)]
13     pub value: MetricValue,
14 }
15 
16 #[derive(Debug, Hash, Clone, PartialEq, Deserialize, Serialize, is_enum_variant)]
17 #[serde(rename_all = "snake_case")]
18 pub enum MetricKind {
19     Incremental,
20     Absolute,
21 }
22 
23 #[derive(Debug, Clone, PartialEq, Deserialize, Serialize, is_enum_variant)]
24 #[serde(rename_all = "snake_case")]
25 pub enum MetricValue {
26     Counter {
27         value: f64,
28     },
29     Gauge {
30         value: f64,
31     },
32     Set {
33         values: BTreeSet<String>,
34     },
35     Distribution {
36         values: Vec<f64>,
37         sample_rates: Vec<u32>,
38     },
39     AggregatedHistogram {
40         buckets: Vec<f64>,
41         counts: Vec<u32>,
42         count: u32,
43         sum: f64,
44     },
45     AggregatedSummary {
46         quantiles: Vec<f64>,
47         values: Vec<f64>,
48         count: u32,
49         sum: f64,
50     },
51 }
52 
53 impl Metric {
to_absolute(&self) -> Self54     pub fn to_absolute(&self) -> Self {
55         Self {
56             name: self.name.clone(),
57             timestamp: self.timestamp,
58             tags: self.tags.clone(),
59             kind: MetricKind::Absolute,
60             value: self.value.clone(),
61         }
62     }
63 
add(&mut self, other: &Self)64     pub fn add(&mut self, other: &Self) {
65         if other.kind.is_absolute() {
66             return;
67         }
68 
69         match (&mut self.value, &other.value) {
70             (MetricValue::Counter { ref mut value }, MetricValue::Counter { value: value2 }) => {
71                 *value += value2;
72             }
73             (MetricValue::Gauge { ref mut value }, MetricValue::Gauge { value: value2 }) => {
74                 *value += value2;
75             }
76             (MetricValue::Set { ref mut values }, MetricValue::Set { values: values2 }) => {
77                 for val in values2 {
78                     values.insert(val.to_string());
79                 }
80             }
81             (
82                 MetricValue::Distribution {
83                     ref mut values,
84                     ref mut sample_rates,
85                 },
86                 MetricValue::Distribution {
87                     values: values2,
88                     sample_rates: sample_rates2,
89                 },
90             ) => {
91                 values.extend_from_slice(&values2);
92                 sample_rates.extend_from_slice(&sample_rates2);
93             }
94             (
95                 MetricValue::AggregatedHistogram {
96                     ref buckets,
97                     ref mut counts,
98                     ref mut count,
99                     ref mut sum,
100                 },
101                 MetricValue::AggregatedHistogram {
102                     buckets: buckets2,
103                     counts: counts2,
104                     count: count2,
105                     sum: sum2,
106                 },
107             ) => {
108                 if buckets == buckets2 && counts.len() == counts2.len() {
109                     for (i, c) in counts2.iter().enumerate() {
110                         counts[i] += c;
111                     }
112                     *count += count2;
113                     *sum += sum2;
114                 }
115             }
116             _ => {}
117         }
118     }
119 
reset(&mut self)120     pub fn reset(&mut self) {
121         match &mut self.value {
122             MetricValue::Counter { ref mut value } => {
123                 *value = 0.0;
124             }
125             MetricValue::Gauge { ref mut value } => {
126                 *value = 0.0;
127             }
128             MetricValue::Set { ref mut values } => {
129                 values.clear();
130             }
131             MetricValue::Distribution {
132                 ref mut values,
133                 ref mut sample_rates,
134             } => {
135                 values.clear();
136                 sample_rates.clear();
137             }
138             MetricValue::AggregatedHistogram {
139                 ref mut counts,
140                 ref mut count,
141                 ref mut sum,
142                 ..
143             } => {
144                 for c in counts.iter_mut() {
145                     *c = 0;
146                 }
147                 *count = 0;
148                 *sum = 0.0;
149             }
150             MetricValue::AggregatedSummary {
151                 ref mut values,
152                 ref mut count,
153                 ref mut sum,
154                 ..
155             } => {
156                 for v in values.iter_mut() {
157                     *v = 0.0;
158                 }
159                 *count = 0;
160                 *sum = 0.0;
161             }
162         }
163     }
164 }
165 
166 #[cfg(test)]
167 mod test {
168     use super::*;
169     use chrono::{offset::TimeZone, DateTime, Utc};
170 
ts() -> DateTime<Utc>171     fn ts() -> DateTime<Utc> {
172         Utc.ymd(2018, 11, 14).and_hms_nano(8, 9, 10, 11)
173     }
174 
tags() -> BTreeMap<String, String>175     fn tags() -> BTreeMap<String, String> {
176         vec![
177             ("normal_tag".to_owned(), "value".to_owned()),
178             ("true_tag".to_owned(), "true".to_owned()),
179             ("empty_tag".to_owned(), "".to_owned()),
180         ]
181         .into_iter()
182         .collect()
183     }
184 
185     #[test]
merge_counters()186     fn merge_counters() {
187         let mut counter = Metric {
188             name: "counter".into(),
189             timestamp: None,
190             tags: None,
191             kind: MetricKind::Incremental,
192             value: MetricValue::Counter { value: 1.0 },
193         };
194 
195         let delta = Metric {
196             name: "counter".into(),
197             timestamp: Some(ts()),
198             tags: Some(tags()),
199             kind: MetricKind::Incremental,
200             value: MetricValue::Counter { value: 2.0 },
201         };
202 
203         counter.add(&delta);
204         assert_eq!(
205             counter,
206             Metric {
207                 name: "counter".into(),
208                 timestamp: None,
209                 tags: None,
210                 kind: MetricKind::Incremental,
211                 value: MetricValue::Counter { value: 3.0 },
212             }
213         )
214     }
215 
216     #[test]
merge_gauges()217     fn merge_gauges() {
218         let mut gauge = Metric {
219             name: "gauge".into(),
220             timestamp: None,
221             tags: None,
222             kind: MetricKind::Incremental,
223             value: MetricValue::Gauge { value: 1.0 },
224         };
225 
226         let delta = Metric {
227             name: "gauge".into(),
228             timestamp: Some(ts()),
229             tags: Some(tags()),
230             kind: MetricKind::Incremental,
231             value: MetricValue::Gauge { value: -2.0 },
232         };
233 
234         gauge.add(&delta);
235         assert_eq!(
236             gauge,
237             Metric {
238                 name: "gauge".into(),
239                 timestamp: None,
240                 tags: None,
241                 kind: MetricKind::Incremental,
242                 value: MetricValue::Gauge { value: -1.0 },
243             }
244         )
245     }
246 
247     #[test]
merge_sets()248     fn merge_sets() {
249         let mut set = Metric {
250             name: "set".into(),
251             timestamp: None,
252             tags: None,
253             kind: MetricKind::Incremental,
254             value: MetricValue::Set {
255                 values: vec!["old".into()].into_iter().collect(),
256             },
257         };
258 
259         let delta = Metric {
260             name: "set".into(),
261             timestamp: Some(ts()),
262             tags: Some(tags()),
263             kind: MetricKind::Incremental,
264             value: MetricValue::Set {
265                 values: vec!["new".into()].into_iter().collect(),
266             },
267         };
268 
269         set.add(&delta);
270         assert_eq!(
271             set,
272             Metric {
273                 name: "set".into(),
274                 timestamp: None,
275                 tags: None,
276                 kind: MetricKind::Incremental,
277                 value: MetricValue::Set {
278                     values: vec!["old".into(), "new".into()].into_iter().collect()
279                 },
280             }
281         )
282     }
283 
284     #[test]
merge_histograms()285     fn merge_histograms() {
286         let mut dist = Metric {
287             name: "hist".into(),
288             timestamp: None,
289             tags: None,
290             kind: MetricKind::Incremental,
291             value: MetricValue::Distribution {
292                 values: vec![1.0],
293                 sample_rates: vec![10],
294             },
295         };
296 
297         let delta = Metric {
298             name: "hist".into(),
299             timestamp: Some(ts()),
300             tags: Some(tags()),
301             kind: MetricKind::Incremental,
302             value: MetricValue::Distribution {
303                 values: vec![1.0],
304                 sample_rates: vec![20],
305             },
306         };
307 
308         dist.add(&delta);
309         assert_eq!(
310             dist,
311             Metric {
312                 name: "hist".into(),
313                 timestamp: None,
314                 tags: None,
315                 kind: MetricKind::Incremental,
316                 value: MetricValue::Distribution {
317                     values: vec![1.0, 1.0],
318                     sample_rates: vec![10, 20],
319                 },
320             }
321         )
322     }
323 }
324