vector/sinks/util/
metadata.rs

1use std::num::NonZeroUsize;
2
3use vector_lib::{
4    ByteSizeOf, EstimatedJsonEncodedSizeOf, config,
5    request_metadata::{GetEventCountTags, GroupedCountByteSize, RequestMetadata},
6};
7
8use super::request_builder::EncodeResult;
9
10#[derive(Clone, Default)]
11pub struct RequestMetadataBuilder {
12    event_count: usize,
13    events_byte_size: usize,
14    grouped_events_byte_size: GroupedCountByteSize,
15}
16
17impl RequestMetadataBuilder {
18    pub fn from_events<E>(events: &[E]) -> Self
19    where
20        E: ByteSizeOf + GetEventCountTags + EstimatedJsonEncodedSizeOf,
21    {
22        let mut size = config::telemetry().create_request_count_byte_size();
23
24        let mut events_byte_size = 0;
25
26        for event in events {
27            events_byte_size += event.size_of();
28            size.add_event(event, event.estimated_json_encoded_size_of());
29        }
30
31        Self {
32            event_count: events.len(),
33            events_byte_size,
34            grouped_events_byte_size: size,
35        }
36    }
37
38    pub fn from_event<E>(event: &E) -> Self
39    where
40        E: ByteSizeOf + GetEventCountTags + EstimatedJsonEncodedSizeOf,
41    {
42        let mut size = config::telemetry().create_request_count_byte_size();
43        size.add_event(event, event.estimated_json_encoded_size_of());
44
45        Self {
46            event_count: 1,
47            events_byte_size: event.size_of(),
48            grouped_events_byte_size: size,
49        }
50    }
51
52    pub const fn new(
53        event_count: usize,
54        events_byte_size: usize,
55        grouped_events_byte_size: GroupedCountByteSize,
56    ) -> Self {
57        Self {
58            event_count,
59            events_byte_size,
60            grouped_events_byte_size,
61        }
62    }
63
64    pub fn track_event<E>(&mut self, event: E)
65    where
66        E: ByteSizeOf + GetEventCountTags + EstimatedJsonEncodedSizeOf,
67    {
68        self.event_count += 1;
69        self.events_byte_size += event.size_of();
70        let json_size = event.estimated_json_encoded_size_of();
71        self.grouped_events_byte_size.add_event(&event, json_size);
72    }
73
74    /// Builds the [`RequestMetadata`] with the given size.
75    /// This is used when there is no encoder in the process to provide an `EncodeResult`
76    pub fn with_request_size(&self, size: NonZeroUsize) -> RequestMetadata {
77        let size = size.get();
78
79        RequestMetadata::new(
80            self.event_count,
81            self.events_byte_size,
82            size,
83            size,
84            self.grouped_events_byte_size.clone(),
85        )
86    }
87
88    /// Builds the [`RequestMetadata`] from the results of encoding.
89    /// `EncodeResult` provides us with the byte size before and after compression
90    /// and the json size of the events after transforming (dropping unwanted fields) but
91    /// before encoding.
92    pub fn build<T>(&self, result: &EncodeResult<T>) -> RequestMetadata {
93        RequestMetadata::new(
94            self.event_count,
95            self.events_byte_size,
96            result.uncompressed_byte_size,
97            result
98                .compressed_byte_size
99                .unwrap_or(result.uncompressed_byte_size),
100            // Building from an encoded result, we take the json size from the encoded since that has the size
101            // after transforming the event.
102            result.transformed_json_size.clone(),
103        )
104    }
105}
vector/sinks/util/metadata.rs

vector/sinks/util/
metadata.rs
