vector/codecs/encoding/
encoder.rs

1use bytes::BytesMut;
2use tokio_util::codec::Encoder as _;
3use vector_lib::codecs::{
4    encoding::{Error, Framer, Serializer},
5    CharacterDelimitedEncoder, NewlineDelimitedEncoder, TextSerializerConfig,
6};
7
8use crate::{
9    event::Event,
10    internal_events::{EncoderFramingError, EncoderSerializeError},
11};
12
13#[derive(Debug, Clone)]
14/// An encoder that can encode structured events into byte frames.
15pub struct Encoder<Framer>
16where
17    Framer: Clone,
18{
19    framer: Framer,
20    serializer: Serializer,
21}
22
23impl Default for Encoder<Framer> {
24    fn default() -> Self {
25        Self {
26            framer: NewlineDelimitedEncoder::default().into(),
27            serializer: TextSerializerConfig::default().build().into(),
28        }
29    }
30}
31
32impl Default for Encoder<()> {
33    fn default() -> Self {
34        Self {
35            framer: (),
36            serializer: TextSerializerConfig::default().build().into(),
37        }
38    }
39}
40
41impl<Framer> Encoder<Framer>
42where
43    Framer: Clone,
44{
45    /// Serialize the event without applying framing.
46    pub fn serialize(&mut self, event: Event, buffer: &mut BytesMut) -> Result<(), Error> {
47        let len = buffer.len();
48        let mut payload = buffer.split_off(len);
49
50        self.serialize_at_start(event, &mut payload)?;
51
52        buffer.unsplit(payload);
53
54        Ok(())
55    }
56
57    /// Serialize the event without applying framing, at the start of the provided buffer.
58    fn serialize_at_start(&mut self, event: Event, buffer: &mut BytesMut) -> Result<(), Error> {
59        self.serializer.encode(event, buffer).map_err(|error| {
60            emit!(EncoderSerializeError { error: &error });
61            Error::SerializingError(error)
62        })
63    }
64}
65
66impl Encoder<Framer> {
67    /// Creates a new `Encoder` with the specified `Serializer` to produce bytes
68    /// from a structured event, and the `Framer` to wrap these into a byte
69    /// frame.
70    pub const fn new(framer: Framer, serializer: Serializer) -> Self {
71        Self { framer, serializer }
72    }
73
74    /// Get the framer.
75    pub const fn framer(&self) -> &Framer {
76        &self.framer
77    }
78
79    /// Get the serializer.
80    pub const fn serializer(&self) -> &Serializer {
81        &self.serializer
82    }
83
84    /// Get the prefix that encloses a batch of events.
85    pub const fn batch_prefix(&self) -> &[u8] {
86        match (&self.framer, &self.serializer) {
87            (
88                Framer::CharacterDelimited(CharacterDelimitedEncoder { delimiter: b',' }),
89                Serializer::Json(_) | Serializer::NativeJson(_),
90            ) => b"[",
91            _ => &[],
92        }
93    }
94
95    /// Get the suffix that encloses a batch of events.
96    pub const fn batch_suffix(&self, empty: bool) -> &[u8] {
97        match (&self.framer, &self.serializer, empty) {
98            (
99                Framer::CharacterDelimited(CharacterDelimitedEncoder { delimiter: b',' }),
100                Serializer::Json(_) | Serializer::NativeJson(_),
101                _,
102            ) => b"]",
103            (Framer::NewlineDelimited(_), _, false) => b"\n",
104            _ => &[],
105        }
106    }
107
108    /// Get the HTTP content type.
109    pub const fn content_type(&self) -> &'static str {
110        match (&self.serializer, &self.framer) {
111            (Serializer::Json(_) | Serializer::NativeJson(_), Framer::NewlineDelimited(_)) => {
112                "application/x-ndjson"
113            }
114            (
115                Serializer::Gelf(_) | Serializer::Json(_) | Serializer::NativeJson(_),
116                Framer::CharacterDelimited(CharacterDelimitedEncoder { delimiter: b',' }),
117            ) => "application/json",
118            (Serializer::Native(_), _) | (Serializer::Protobuf(_), _) => "application/octet-stream",
119            (
120                Serializer::Avro(_)
121                | Serializer::Cef(_)
122                | Serializer::Csv(_)
123                | Serializer::Gelf(_)
124                | Serializer::Json(_)
125                | Serializer::Logfmt(_)
126                | Serializer::NativeJson(_)
127                | Serializer::RawMessage(_)
128                | Serializer::Text(_),
129                _,
130            ) => "text/plain",
131        }
132    }
133}
134
135impl Encoder<()> {
136    /// Creates a new `Encoder` with the specified `Serializer` to produce bytes
137    /// from a structured event.
138    pub const fn new(serializer: Serializer) -> Self {
139        Self {
140            framer: (),
141            serializer,
142        }
143    }
144
145    /// Get the serializer.
146    pub const fn serializer(&self) -> &Serializer {
147        &self.serializer
148    }
149}
150
151impl tokio_util::codec::Encoder<Event> for Encoder<Framer> {
152    type Error = Error;
153
154    fn encode(&mut self, event: Event, buffer: &mut BytesMut) -> Result<(), Self::Error> {
155        let len = buffer.len();
156        let mut payload = buffer.split_off(len);
157
158        self.serialize_at_start(event, &mut payload)?;
159
160        // Frame the serialized event.
161        self.framer.encode((), &mut payload).map_err(|error| {
162            emit!(EncoderFramingError { error: &error });
163            Error::FramingError(error)
164        })?;
165
166        buffer.unsplit(payload);
167
168        Ok(())
169    }
170}
171
172impl tokio_util::codec::Encoder<Event> for Encoder<()> {
173    type Error = Error;
174
175    fn encode(&mut self, event: Event, buffer: &mut BytesMut) -> Result<(), Self::Error> {
176        let len = buffer.len();
177        let mut payload = buffer.split_off(len);
178
179        self.serialize_at_start(event, &mut payload)?;
180
181        buffer.unsplit(payload);
182
183        Ok(())
184    }
185}
186
187#[cfg(test)]
188mod tests {
189    use bytes::BufMut;
190    use futures_util::{SinkExt, StreamExt};
191    use tokio_util::codec::FramedWrite;
192    use vector_lib::codecs::encoding::BoxedFramingError;
193    use vector_lib::event::LogEvent;
194
195    use super::*;
196
197    #[derive(Debug, Clone)]
198    struct ParenEncoder;
199
200    impl ParenEncoder {
201        pub(super) const fn new() -> Self {
202            Self
203        }
204    }
205
206    impl tokio_util::codec::Encoder<()> for ParenEncoder {
207        type Error = BoxedFramingError;
208
209        fn encode(&mut self, _: (), dst: &mut BytesMut) -> Result<(), Self::Error> {
210            dst.reserve(2);
211            let inner = dst.split();
212            dst.put_u8(b'(');
213            dst.unsplit(inner);
214            dst.put_u8(b')');
215            Ok(())
216        }
217    }
218
219    #[derive(Debug, Clone)]
220    struct ErrorNthEncoder<T>(T, usize, usize)
221    where
222        T: tokio_util::codec::Encoder<(), Error = BoxedFramingError>;
223
224    impl<T> ErrorNthEncoder<T>
225    where
226        T: tokio_util::codec::Encoder<(), Error = BoxedFramingError>,
227    {
228        pub(super) const fn new(encoder: T, n: usize) -> Self {
229            Self(encoder, 0, n)
230        }
231    }
232
233    impl<T> tokio_util::codec::Encoder<()> for ErrorNthEncoder<T>
234    where
235        T: tokio_util::codec::Encoder<(), Error = BoxedFramingError>,
236    {
237        type Error = BoxedFramingError;
238
239        fn encode(&mut self, _: (), dst: &mut BytesMut) -> Result<(), Self::Error> {
240            self.0.encode((), dst)?;
241            let result = if self.1 == self.2 {
242                Err(Box::new(std::io::Error::other("error")) as _)
243            } else {
244                Ok(())
245            };
246            self.1 += 1;
247            result
248        }
249    }
250
251    #[tokio::test]
252    async fn test_encode_events_sink_empty() {
253        let encoder = Encoder::<Framer>::new(
254            Framer::Boxed(Box::new(ParenEncoder::new())),
255            TextSerializerConfig::default().build().into(),
256        );
257        let source = futures::stream::iter(vec![
258            Event::Log(LogEvent::from("foo")),
259            Event::Log(LogEvent::from("bar")),
260            Event::Log(LogEvent::from("baz")),
261        ])
262        .map(Ok);
263        let sink = Vec::new();
264        let mut framed = FramedWrite::new(sink, encoder);
265        source.forward(&mut framed).await.unwrap();
266        let sink = framed.into_inner();
267        assert_eq!(sink, b"(foo)(bar)(baz)");
268    }
269
270    #[tokio::test]
271    async fn test_encode_events_sink_non_empty() {
272        let encoder = Encoder::<Framer>::new(
273            Framer::Boxed(Box::new(ParenEncoder::new())),
274            TextSerializerConfig::default().build().into(),
275        );
276        let source = futures::stream::iter(vec![
277            Event::Log(LogEvent::from("bar")),
278            Event::Log(LogEvent::from("baz")),
279            Event::Log(LogEvent::from("bat")),
280        ])
281        .map(Ok);
282        let sink = Vec::from("(foo)");
283        let mut framed = FramedWrite::new(sink, encoder);
284        source.forward(&mut framed).await.unwrap();
285        let sink = framed.into_inner();
286        assert_eq!(sink, b"(foo)(bar)(baz)(bat)");
287    }
288
289    #[tokio::test]
290    async fn test_encode_events_sink_empty_handle_framing_error() {
291        let encoder = Encoder::<Framer>::new(
292            Framer::Boxed(Box::new(ErrorNthEncoder::new(ParenEncoder::new(), 1))),
293            TextSerializerConfig::default().build().into(),
294        );
295        let source = futures::stream::iter(vec![
296            Event::Log(LogEvent::from("foo")),
297            Event::Log(LogEvent::from("bar")),
298            Event::Log(LogEvent::from("baz")),
299        ])
300        .map(Ok);
301        let sink = Vec::new();
302        let mut framed = FramedWrite::new(sink, encoder);
303        assert!(source.forward(&mut framed).await.is_err());
304        framed.flush().await.unwrap();
305        let sink = framed.into_inner();
306        assert_eq!(sink, b"(foo)");
307    }
308
309    #[tokio::test]
310    async fn test_encode_events_sink_non_empty_handle_framing_error() {
311        let encoder = Encoder::<Framer>::new(
312            Framer::Boxed(Box::new(ErrorNthEncoder::new(ParenEncoder::new(), 1))),
313            TextSerializerConfig::default().build().into(),
314        );
315        let source = futures::stream::iter(vec![
316            Event::Log(LogEvent::from("bar")),
317            Event::Log(LogEvent::from("baz")),
318            Event::Log(LogEvent::from("bat")),
319        ])
320        .map(Ok);
321        let sink = Vec::from("(foo)");
322        let mut framed = FramedWrite::new(sink, encoder);
323        assert!(source.forward(&mut framed).await.is_err());
324        framed.flush().await.unwrap();
325        let sink = framed.into_inner();
326        assert_eq!(sink, b"(foo)(bar)");
327    }
328
329    #[tokio::test]
330    async fn test_encode_batch_newline() {
331        let encoder = Encoder::<Framer>::new(
332            Framer::NewlineDelimited(NewlineDelimitedEncoder::default()),
333            TextSerializerConfig::default().build().into(),
334        );
335        let source = futures::stream::iter(vec![
336            Event::Log(LogEvent::from("bar")),
337            Event::Log(LogEvent::from("baz")),
338            Event::Log(LogEvent::from("bat")),
339        ])
340        .map(Ok);
341        let sink: Vec<u8> = Vec::new();
342        let mut framed = FramedWrite::new(sink, encoder);
343        source.forward(&mut framed).await.unwrap();
344        let sink = framed.into_inner();
345        assert_eq!(sink, b"bar\nbaz\nbat\n");
346    }
347}
vector/codecs/encoding/encoder.rs

vector/codecs/encoding/
encoder.rs
