vector/transforms/aggregate/
transform.rs

1use super::{AggregateConfig, AggregationMode};
2
3use std::{
4    collections::{HashMap, hash_map::Entry},
5    pin::Pin,
6    time::Duration,
7};
8
9use async_stream::stream;
10use futures::{Stream, StreamExt};
11use vector_lib::event::{
12    MetricValue,
13    metric::{Metric, MetricData, MetricKind, MetricSeries},
14};
15
16use crate::{
17    event::{Event, EventMetadata},
18    internal_events::{AggregateEventRecorded, AggregateFlushed, AggregateUpdateFailed},
19    transforms::TaskTransform,
20};
21
22#[derive(Clone, Debug, Default, PartialEq)]
23enum InnerMode {
24    /// Default mode. Sums incremental metrics and uses the latest value for absolute metrics.
25    #[default]
26    Auto,
27
28    /// Sums incremental metrics; absolute metrics pass through unchanged.
29    Sum,
30
31    /// Returns the latest value for absolute metrics; incremental metrics pass through unchanged.
32    Latest,
33
34    /// Counts metrics for incremental and absolute metrics
35    Count,
36
37    /// Returns difference between latest value for absolute; incremental metrics pass through unchanged.
38    Diff {
39        prev_map: HashMap<MetricSeries, MetricEntry>,
40    },
41
42    /// Max value of absolute metric; incremental metrics pass through unchanged.
43    Max,
44
45    /// Min value of absolute metric; incremental metrics pass through unchanged.
46    Min,
47
48    /// Mean value of absolute metric; incremental metrics pass through unchanged.
49    Mean {
50        multi_map: HashMap<MetricSeries, Vec<MetricEntry>>,
51    },
52
53    /// Stdev value of absolute metric; incremental metrics pass through unchanged.
54    Stdev {
55        multi_map: HashMap<MetricSeries, Vec<MetricEntry>>,
56    },
57}
58
59impl From<AggregationMode> for InnerMode {
60    fn from(value: AggregationMode) -> Self {
61        match value {
62            AggregationMode::Auto => InnerMode::Auto,
63            AggregationMode::Sum => InnerMode::Sum,
64            AggregationMode::Latest => InnerMode::Latest,
65            AggregationMode::Count => InnerMode::Count,
66            AggregationMode::Diff => InnerMode::Diff {
67                prev_map: HashMap::default(),
68            },
69            AggregationMode::Max => InnerMode::Max,
70            AggregationMode::Min => InnerMode::Min,
71            AggregationMode::Mean => InnerMode::Mean {
72                multi_map: HashMap::default(),
73            },
74            AggregationMode::Stdev => InnerMode::Stdev {
75                multi_map: HashMap::default(),
76            },
77        }
78    }
79}
80
81type MetricEntry = (MetricData, EventMetadata);
82
83#[derive(Debug)]
84pub struct Aggregate {
85    interval: Duration,
86    map: HashMap<MetricSeries, MetricEntry>,
87    mode: InnerMode,
88}
89
90impl Aggregate {
91    pub fn new(config: &AggregateConfig) -> crate::Result<Self> {
92        Ok(Self {
93            interval: Duration::from_millis(config.interval_ms),
94            map: Default::default(),
95            mode: config.mode.into(),
96        })
97    }
98
99    pub fn record(&mut self, event: Event) -> Option<Event> {
100        let (series, data, metadata) = event.into_metric().into_parts();
101
102        match (&mut self.mode, data.kind) {
103            (InnerMode::Sum, MetricKind::Absolute)
104            | (InnerMode::Latest | InnerMode::Diff { .. }, MetricKind::Incremental)
105            | (InnerMode::Max | InnerMode::Min, MetricKind::Incremental)
106            | (InnerMode::Mean { .. } | InnerMode::Stdev { .. }, MetricKind::Incremental) => {
107                return Some(Event::Metric(Metric::from_parts(series, data, metadata)));
108            }
109            (InnerMode::Auto | InnerMode::Sum, MetricKind::Incremental) => {
110                self.record_sum(series, data, metadata);
111            }
112            (InnerMode::Auto, MetricKind::Absolute)
113            | (InnerMode::Latest | InnerMode::Diff { .. }, MetricKind::Absolute) => {
114                self.map.insert(series, (data, metadata));
115            }
116            (InnerMode::Count, _) => {
117                self.record_count(series, data, metadata);
118            }
119            (InnerMode::Max | InnerMode::Min, MetricKind::Absolute) => {
120                self.record_comparison(series, data, metadata);
121            }
122            (
123                InnerMode::Mean { multi_map } | InnerMode::Stdev { multi_map },
124                MetricKind::Absolute,
125            ) => {
126                if matches!(data.value, MetricValue::Gauge { value: _ }) {
127                    match multi_map.entry(series) {
128                        Entry::Occupied(mut entry) => entry.get_mut().push((data, metadata)),
129                        Entry::Vacant(entry) => {
130                            entry.insert(vec![(data, metadata)]);
131                        }
132                    }
133                }
134            }
135        }
136        emit!(AggregateEventRecorded);
137        None
138    }
139
140    fn record_count(
141        &mut self,
142        series: MetricSeries,
143        mut data: MetricData,
144        metadata: EventMetadata,
145    ) {
146        let mut count_data = data.clone();
147        let existing = self.map.entry(series).or_insert_with(|| {
148            *data.value_mut() = MetricValue::Counter { value: 0f64 };
149            (data.clone(), metadata.clone())
150        });
151        *count_data.value_mut() = MetricValue::Counter { value: 1f64 };
152        if existing.0.kind == data.kind && existing.0.update(&count_data) {
153            existing.1.merge(metadata);
154        } else {
155            emit!(AggregateUpdateFailed);
156        }
157    }
158
159    fn record_sum(&mut self, series: MetricSeries, data: MetricData, metadata: EventMetadata) {
160        match self.map.entry(series) {
161            Entry::Occupied(mut entry) => {
162                let existing = entry.get_mut();
163                // In order to update (add) the new and old kind's must match
164                if existing.0.kind == data.kind && existing.0.update(&data) {
165                    existing.1.merge(metadata);
166                } else {
167                    emit!(AggregateUpdateFailed);
168                    *existing = (data, metadata);
169                }
170            }
171            Entry::Vacant(entry) => {
172                entry.insert((data, metadata));
173            }
174        }
175    }
176
177    fn record_comparison(
178        &mut self,
179        series: MetricSeries,
180        data: MetricData,
181        metadata: EventMetadata,
182    ) {
183        match self.map.entry(series) {
184            Entry::Occupied(mut entry) => {
185                let existing = entry.get_mut();
186                // In order to update (add) the new and old kind's must match
187                if existing.0.kind == data.kind {
188                    if let MetricValue::Gauge {
189                        value: existing_value,
190                    } = existing.0.value()
191                        && let MetricValue::Gauge { value: new_value } = data.value()
192                    {
193                        let should_update = match self.mode {
194                            InnerMode::Max => new_value > existing_value,
195                            InnerMode::Min => new_value < existing_value,
196                            _ => false,
197                        };
198                        if should_update {
199                            *existing = (data, metadata);
200                        }
201                    }
202                } else {
203                    emit!(AggregateUpdateFailed);
204                    *existing = (data, metadata);
205                }
206            }
207            Entry::Vacant(entry) => {
208                entry.insert((data, metadata));
209            }
210        }
211    }
212
213    pub fn flush_into(&mut self, output: &mut Vec<Event>) {
214        let map = std::mem::take(&mut self.map);
215        for (series, entry) in map.clone().into_iter() {
216            let mut metric = Metric::from_parts(series, entry.0, entry.1);
217            if let InnerMode::Diff { prev_map } = &self.mode
218                && let Some(prev_entry) = prev_map.get(metric.series())
219                && metric.data().kind == prev_entry.0.kind
220                && !metric.subtract(&prev_entry.0)
221            {
222                emit!(AggregateUpdateFailed);
223            }
224            output.push(Event::Metric(metric));
225        }
226
227        let multi_map = match &mut self.mode {
228            InnerMode::Mean { multi_map } | InnerMode::Stdev { multi_map } => {
229                std::mem::take(multi_map)
230            }
231            _ => HashMap::default(),
232        };
233
234        'outer: for (series, entries) in multi_map.into_iter() {
235            if entries.is_empty() {
236                continue;
237            }
238
239            let (mut final_sum, mut final_metadata) = entries.first().unwrap().clone();
240            for (data, metadata) in entries.iter().skip(1) {
241                if !final_sum.update(data) {
242                    // Incompatible types, skip this metric
243                    emit!(AggregateUpdateFailed);
244                    continue 'outer;
245                }
246                final_metadata.merge(metadata.clone());
247            }
248
249            let final_mean_value = if let MetricValue::Gauge { value } = final_sum.value_mut() {
250                // Entries are not empty so this is safe.
251                *value /= entries.len() as f64;
252                *value
253            } else {
254                0.0
255            };
256
257            let final_mean = final_sum.clone();
258            match self.mode {
259                InnerMode::Mean { .. } => {
260                    let metric = Metric::from_parts(series, final_mean, final_metadata);
261                    output.push(Event::Metric(metric));
262                }
263                InnerMode::Stdev { .. } => {
264                    let variance = entries
265                        .iter()
266                        .filter_map(|(data, _)| {
267                            if let MetricValue::Gauge { value } = data.value() {
268                                let diff = final_mean_value - value;
269                                Some(diff * diff)
270                            } else {
271                                None
272                            }
273                        })
274                        .sum::<f64>()
275                        / entries.len() as f64;
276                    let mut final_stdev = final_mean;
277                    if let MetricValue::Gauge { value } = final_stdev.value_mut() {
278                        *value = variance.sqrt()
279                    }
280                    let metric = Metric::from_parts(series, final_stdev, final_metadata);
281                    output.push(Event::Metric(metric));
282                }
283                _ => (),
284            }
285        }
286
287        if let InnerMode::Diff { prev_map } = &mut self.mode {
288            *prev_map = map;
289        }
290        emit!(AggregateFlushed);
291    }
292}
293
294impl TaskTransform<Event> for Aggregate {
295    fn transform(
296        mut self: Box<Self>,
297        mut input_rx: Pin<Box<dyn Stream<Item = Event> + Send>>,
298    ) -> Pin<Box<dyn Stream<Item = Event> + Send>>
299    where
300        Self: 'static,
301    {
302        let mut flush_stream = tokio::time::interval(self.interval);
303
304        Box::pin(stream! {
305            let mut output = Vec::new();
306            let mut done = false;
307            while !done {
308                tokio::select! {
309                    _ = flush_stream.tick() => {
310                        self.flush_into(&mut output);
311                    },
312                    maybe_event = input_rx.next() => {
313                        match maybe_event {
314                            None => {
315                                self.flush_into(&mut output);
316                                done = true;
317                            }
318                            Some(event) => {
319                                if let Some(passthrough) = self.record(event) {
320                                    output.push(passthrough);
321                                }
322                            }
323                        }
324                    }
325                };
326                for event in output.drain(..) {
327                    yield event;
328                }
329            }
330        })
331    }
332}
vector/transforms/aggregate/transform.rs

vector/transforms/aggregate/
transform.rs
