vector/codecs/encoding/
encoder.rs

1use bytes::BytesMut;
2use tokio_util::codec::Encoder as _;
3use vector_lib::codecs::{
4    CharacterDelimitedEncoder, NewlineDelimitedEncoder, TextSerializerConfig,
5    encoding::{Error, Framer, Serializer},
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, event::LogEvent};
193
194    use super::*;
195
196    #[derive(Debug, Clone)]
197    struct ParenEncoder;
198
199    impl ParenEncoder {
200        pub(super) const fn new() -> Self {
201            Self
202        }
203    }
204
205    impl tokio_util::codec::Encoder<()> for ParenEncoder {
206        type Error = BoxedFramingError;
207
208        fn encode(&mut self, _: (), dst: &mut BytesMut) -> Result<(), Self::Error> {
209            dst.reserve(2);
210            let inner = dst.split();
211            dst.put_u8(b'(');
212            dst.unsplit(inner);
213            dst.put_u8(b')');
214            Ok(())
215        }
216    }
217
218    #[derive(Debug, Clone)]
219    struct ErrorNthEncoder<T>(T, usize, usize)
220    where
221        T: tokio_util::codec::Encoder<(), Error = BoxedFramingError>;
222
223    impl<T> ErrorNthEncoder<T>
224    where
225        T: tokio_util::codec::Encoder<(), Error = BoxedFramingError>,
226    {
227        pub(super) const fn new(encoder: T, n: usize) -> Self {
228            Self(encoder, 0, n)
229        }
230    }
231
232    impl<T> tokio_util::codec::Encoder<()> for ErrorNthEncoder<T>
233    where
234        T: tokio_util::codec::Encoder<(), Error = BoxedFramingError>,
235    {
236        type Error = BoxedFramingError;
237
238        fn encode(&mut self, _: (), dst: &mut BytesMut) -> Result<(), Self::Error> {
239            self.0.encode((), dst)?;
240            let result = if self.1 == self.2 {
241                Err(Box::new(std::io::Error::other("error")) as _)
242            } else {
243                Ok(())
244            };
245            self.1 += 1;
246            result
247        }
248    }
249
250    #[tokio::test]
251    async fn test_encode_events_sink_empty() {
252        let encoder = Encoder::<Framer>::new(
253            Framer::Boxed(Box::new(ParenEncoder::new())),
254            TextSerializerConfig::default().build().into(),
255        );
256        let source = futures::stream::iter(vec![
257            Event::Log(LogEvent::from("foo")),
258            Event::Log(LogEvent::from("bar")),
259            Event::Log(LogEvent::from("baz")),
260        ])
261        .map(Ok);
262        let sink = Vec::new();
263        let mut framed = FramedWrite::new(sink, encoder);
264        source.forward(&mut framed).await.unwrap();
265        let sink = framed.into_inner();
266        assert_eq!(sink, b"(foo)(bar)(baz)");
267    }
268
269    #[tokio::test]
270    async fn test_encode_events_sink_non_empty() {
271        let encoder = Encoder::<Framer>::new(
272            Framer::Boxed(Box::new(ParenEncoder::new())),
273            TextSerializerConfig::default().build().into(),
274        );
275        let source = futures::stream::iter(vec![
276            Event::Log(LogEvent::from("bar")),
277            Event::Log(LogEvent::from("baz")),
278            Event::Log(LogEvent::from("bat")),
279        ])
280        .map(Ok);
281        let sink = Vec::from("(foo)");
282        let mut framed = FramedWrite::new(sink, encoder);
283        source.forward(&mut framed).await.unwrap();
284        let sink = framed.into_inner();
285        assert_eq!(sink, b"(foo)(bar)(baz)(bat)");
286    }
287
288    #[tokio::test]
289    async fn test_encode_events_sink_empty_handle_framing_error() {
290        let encoder = Encoder::<Framer>::new(
291            Framer::Boxed(Box::new(ErrorNthEncoder::new(ParenEncoder::new(), 1))),
292            TextSerializerConfig::default().build().into(),
293        );
294        let source = futures::stream::iter(vec![
295            Event::Log(LogEvent::from("foo")),
296            Event::Log(LogEvent::from("bar")),
297            Event::Log(LogEvent::from("baz")),
298        ])
299        .map(Ok);
300        let sink = Vec::new();
301        let mut framed = FramedWrite::new(sink, encoder);
302        assert!(source.forward(&mut framed).await.is_err());
303        framed.flush().await.unwrap();
304        let sink = framed.into_inner();
305        assert_eq!(sink, b"(foo)");
306    }
307
308    #[tokio::test]
309    async fn test_encode_events_sink_non_empty_handle_framing_error() {
310        let encoder = Encoder::<Framer>::new(
311            Framer::Boxed(Box::new(ErrorNthEncoder::new(ParenEncoder::new(), 1))),
312            TextSerializerConfig::default().build().into(),
313        );
314        let source = futures::stream::iter(vec![
315            Event::Log(LogEvent::from("bar")),
316            Event::Log(LogEvent::from("baz")),
317            Event::Log(LogEvent::from("bat")),
318        ])
319        .map(Ok);
320        let sink = Vec::from("(foo)");
321        let mut framed = FramedWrite::new(sink, encoder);
322        assert!(source.forward(&mut framed).await.is_err());
323        framed.flush().await.unwrap();
324        let sink = framed.into_inner();
325        assert_eq!(sink, b"(foo)(bar)");
326    }
327
328    #[tokio::test]
329    async fn test_encode_batch_newline() {
330        let encoder = Encoder::<Framer>::new(
331            Framer::NewlineDelimited(NewlineDelimitedEncoder::default()),
332            TextSerializerConfig::default().build().into(),
333        );
334        let source = futures::stream::iter(vec![
335            Event::Log(LogEvent::from("bar")),
336            Event::Log(LogEvent::from("baz")),
337            Event::Log(LogEvent::from("bat")),
338        ])
339        .map(Ok);
340        let sink: Vec<u8> = Vec::new();
341        let mut framed = FramedWrite::new(sink, encoder);
342        source.forward(&mut framed).await.unwrap();
343        let sink = framed.into_inner();
344        assert_eq!(sink, b"bar\nbaz\nbat\n");
345    }
346}
vector/codecs/encoding/encoder.rs

vector/codecs/encoding/
encoder.rs
