Make generic encode_to_datagrams

Author: thomasqueirozb

src/sinks/util/datagram.rs

@@ -1,13 +1,14 @@
 use std::pin::Pin;
 
-use bytes::Bytes;
+use bytes::{Bytes, BytesMut};
 use futures::{StreamExt, stream::BoxStream};
 use futures_util::stream::Peekable;
 #[cfg(unix)]
 use std::path::PathBuf;
 use tokio::net::UdpSocket;
 #[cfg(unix)]
 use tokio::net::UnixDatagram;
+use tokio_util::codec::Encoder;
 use vector_lib::{
     codecs::encoding::{Chunker, Chunking},
     internal_event::{ByteSize, BytesSent, InternalEventHandle, RegisterInternalEvent},
@@ -16,7 +17,8 @@ use vector_lib::{
 #[cfg(unix)]
 use crate::internal_events::{UnixSendIncompleteError, UnixSocketSendError};
 use crate::{
-    event::{EventFinalizers, EventStatus},
+    codecs::Transformer,
+    event::{Event, EventFinalizers, EventStatus, Finalizable},
     internal_events::{
         SocketEventsSent, SocketMode, SocketSendError, UdpChunkingError, UdpSendIncompleteError,
     },
@@ -65,6 +67,32 @@ fn is_recoverable_socket_error(error: &std::io::Error) -> bool {
     )
 }
 
+/// Transforms and encodes a raw event stream into a stream of [`EncodedDatagram`]s
+/// ready to be passed to [`send_datagrams`].
+pub fn encode_to_datagrams<'a, E>(
+    input: BoxStream<'a, Event>,
+    transformer: Transformer,
+    mut encoder: E,
+) -> Peekable<BoxStream<'a, EncodedDatagram>>
+where
+    E: Encoder<Event, Error = vector_lib::codecs::encoding::Error> + Send + 'a,
+{
+    input
+        .map(move |mut event| {
+            transformer.transform(&mut event);
+            let finalizers = event.take_finalizers();
+            let mut bytes = BytesMut::new();
+            let bytes = if encoder.encode(event, &mut bytes).is_ok() {
+                Some(bytes.freeze())
+            } else {
+                None
+            };
+            EncodedDatagram { bytes, finalizers }
+        })
+        .boxed()
+        .peekable()
+}
+
 pub async fn send_datagrams(
     input: &mut Peekable<BoxStream<'_, EncodedDatagram>>,
     mut socket: DatagramSocket,

src/sinks/util/udp.rs

@@ -4,8 +4,7 @@ use std::{
 };
 
 use async_trait::async_trait;
-use bytes::BytesMut;
-use futures::{FutureExt, StreamExt, stream::BoxStream};
+use futures::{FutureExt, stream::BoxStream};
 use snafu::{ResultExt, Snafu};
 use tokio::{net::UdpSocket, time::sleep};
 use tokio_util::codec::Encoder;
@@ -17,13 +16,13 @@ use vector_lib::{
 
 use super::{
     SinkBuildError,
-    datagram::{DatagramSocket, EncodedDatagram, send_datagrams},
+    datagram::{DatagramSocket, encode_to_datagrams, send_datagrams},
 };
 use crate::{
     codecs::Transformer,
     common::backoff::ExponentialBackoff,
     dns,
-    event::{Event, Finalizable},
+    event::Event,
     internal_events::{UdpSocketConnectionEstablished, UdpSocketOutgoingConnectionError},
     net,
     sinks::{Healthcheck, VectorSink, util::StreamSink},
@@ -199,22 +198,7 @@ where
     E: Encoder<Event, Error = vector_lib::codecs::encoding::Error> + Clone + Send + Sync,
 {
     async fn run(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
-        let mut encoder = self.encoder.clone();
-        let transformer = self.transformer.clone();
-        let mut input = input
-            .map(move |mut event| {
-                transformer.transform(&mut event);
-                let finalizers = event.take_finalizers();
-                let mut bytes = BytesMut::new();
-                let bytes = if encoder.encode(event, &mut bytes).is_ok() {
-                    Some(bytes.freeze())
-                } else {
-                    None
-                };
-                EncodedDatagram { bytes, finalizers }
-            })
-            .boxed()
-            .peekable();
+        let mut input = encode_to_datagrams(input, self.transformer.clone(), self.encoder.clone());
 
         let chunker = self.chunker.clone();
         while Pin::new(&mut input).peek().await.is_some() {

src/sinks/util/unix.rs

@@ -22,7 +22,7 @@ use vector_lib::{
     json_size::JsonSize,
 };
 
-use super::datagram::{DatagramSocket, EncodedDatagram, send_datagrams};
+use super::datagram::{DatagramSocket, encode_to_datagrams, send_datagrams};
 use crate::{
     codecs::Transformer,
     common::backoff::ExponentialBackoff,
@@ -259,22 +259,7 @@ where
 
     async fn run_datagram(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
         let bytes_sent = register!(BytesSent::from(Protocol::UNIX));
-        let mut encoder = self.encoder.clone();
-        let transformer = self.transformer.clone();
-        let mut input = input
-            .map(move |mut event| {
-                transformer.transform(&mut event);
-                let finalizers = event.take_finalizers();
-                let mut bytes = BytesMut::new();
-                let bytes = if encoder.encode(event, &mut bytes).is_ok() {
-                    Some(bytes.freeze())
-                } else {
-                    None
-                };
-                EncodedDatagram { bytes, finalizers }
-            })
-            .boxed()
-            .peekable();
+        let mut input = encode_to_datagrams(input, self.transformer.clone(), self.encoder.clone());
 
         while Pin::new(&mut input).peek().await.is_some() {
             let socket = match self.connector.connect_backoff().await {