lot o cleanup

tobz · tobz · commit 6037898f289d · 2023-02-17T10:18:08.000-05:00
diff --git a/src/net.rs b/src/net.rs
@@ -1,3 +1,5 @@
+//! Networking-related helper functions.
+
 use std::{io, os::fd::AsRawFd, time::Duration};
 
 use socket2::{SockRef, TcpKeepalive};
diff --git a/src/sinks/statsd/tests.rs b/src/sinks/statsd/tests.rs
@@ -5,9 +5,7 @@ use tokio::{net::UdpSocket, sync::mpsc};
 use tokio_stream::wrappers::ReceiverStream;
 use tokio_util::{codec::BytesCodec, udp::UdpFramed};
 use vector_core::{
-    event::{
-        metric::TagValue, Event, Metric, MetricKind, MetricTags, MetricValue, StatisticKind,
-    },
+    event::{metric::TagValue, Event, Metric, MetricKind, MetricTags, MetricValue, StatisticKind},
     metric_tags,
 };
 
diff --git a/src/sinks/util/service/net/mod.rs b/src/sinks/util/service/net/mod.rs
@@ -6,7 +6,9 @@ pub use self::tcp::{TcpConnector, TcpConnectorConfig};
 pub use self::udp::{UdpConnector, UdpConnectorConfig};
 pub use self::unix::{UnixConnector, UnixConnectorConfig, UnixMode};
 
+use futures_util::future::BoxFuture;
 use snafu::Snafu;
+use tokio::sync::oneshot;
 use vector_config::configurable_component;
 
 /// Hostname and port tuple.
@@ -71,3 +73,21 @@ pub enum NetError {
     #[snafu(display("Failed to get socket back after send as channel closed unexpectedly."))]
     ServiceSocketChannelClosed,
 }
+
+pub enum ServiceState<C> {
+    /// The service is currently disconnected.
+    Disconnected,
+
+    /// The service is currently attempting to connect to the endpoint.
+    Connecting(BoxFuture<'static, C>),
+
+    /// The service is connected and idle.
+    Connected(C),
+
+    /// The service has an in-flight send to the socket.
+    ///
+    /// If the socket experiences an unrecoverable error during the send, `None` will be returned
+    /// over the channel to signal the need to establish a new connection rather than reusing the
+    /// existing connection.
+    Sending(oneshot::Receiver<Option<C>>),
+}
diff --git a/src/sinks/util/service/net/tcp.rs b/src/sinks/util/service/net/tcp.rs
@@ -4,7 +4,8 @@ use std::{
     time::Duration,
 };
 
-use futures::{future::BoxFuture, FutureExt};
+use futures::future::BoxFuture;
+use futures_util::FutureExt;
 use snafu::ResultExt;
 use tokio::{
     io::AsyncWriteExt,
@@ -24,7 +25,7 @@ use crate::{
     sinks::{util::retries::ExponentialBackoff, Healthcheck},
 };
 
-use super::{HostAndPort, NetError, net_error::*};
+use super::{net_error::*, HostAndPort, NetError};
 
 /// `TcpConnector` configuration.
 #[configurable_component]
diff --git a/src/sinks/util/service/net/udp.rs b/src/sinks/util/service/net/udp.rs
@@ -4,7 +4,8 @@ use std::{
     time::Duration,
 };
 
-use futures::{future::BoxFuture, FutureExt};
+use futures::future::BoxFuture;
+use futures_util::FutureExt;
 use snafu::ResultExt;
 use tokio::{net::UdpSocket, sync::oneshot, time::sleep};
 use tower::Service;
@@ -18,7 +19,7 @@ use crate::{
     sinks::{util::retries::ExponentialBackoff, Healthcheck},
 };
 
-use super::{HostAndPort, NetError, net_error::*};
+use super::{net_error::*, HostAndPort, NetError, ServiceState};
 
 /// `UdpConnector` configuration.
 #[configurable_component]
@@ -73,9 +74,7 @@ impl UdpConnector {
         let addr = SocketAddr::new(ip, self.address.port);
         let bind_address = find_bind_address(&addr);
 
-        let socket = UdpSocket::bind(bind_address)
-            .await
-            .context(FailedToBind)?;
+        let socket = UdpSocket::bind(bind_address).await.context(FailedToBind)?;
 
         if let Some(send_buffer_size) = self.send_buffer_size {
             if let Err(error) = net::set_send_buffer_size(&socket, send_buffer_size) {
@@ -117,34 +116,16 @@ impl UdpConnector {
     }
 }
 
-enum UdpServiceState {
-    /// The service is currently disconnected.
-    Disconnected,
-
-    /// The service is currently attempting to connect to the endpoint.
-    Connecting(BoxFuture<'static, UdpSocket>),
-
-    /// The service is connected and idle.
-    Connected(UdpSocket),
-
-    /// The service has an in-flight send to the socket.
-    ///
-    /// If the socket experiences an unrecoverable error during the send, `None` will be returned
-    /// over the channel to signal the need to establish a new connection rather than reusing the
-    /// existing connection.
-    Sending(oneshot::Receiver<Option<UdpSocket>>),
-}
-
 pub struct UdpService {
     connector: UdpConnector,
-    state: UdpServiceState,
+    state: ServiceState<UdpSocket>,
 }
 
 impl UdpService {
     const fn new(connector: UdpConnector) -> Self {
         Self {
             connector,
-            state: UdpServiceState::Disconnected,
+            state: ServiceState::Disconnected,
         }
     }
 }
@@ -157,24 +138,24 @@ impl Service<Vec<u8>> for UdpService {
     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         loop {
             self.state = match &mut self.state {
-                UdpServiceState::Disconnected => {
+                ServiceState::Disconnected => {
                     let connector = self.connector.clone();
-                    UdpServiceState::Connecting(Box::pin(async move {
-                        connector.connect_backoff().await
-                    }))
+                    ServiceState::Connecting(Box::pin(
+                        async move { connector.connect_backoff().await },
+                    ))
                 }
-                UdpServiceState::Connecting(fut) => {
+                ServiceState::Connecting(fut) => {
                     let socket = ready!(fut.poll_unpin(cx));
-                    UdpServiceState::Connected(socket)
+                    ServiceState::Connected(socket)
                 }
-                UdpServiceState::Connected(_) => break,
-                UdpServiceState::Sending(fut) => {
+                ServiceState::Connected(_) => break,
+                ServiceState::Sending(fut) => {
                     match ready!(fut.poll_unpin(cx)) {
                         // When a send concludes, and there's an error, the request future sends
                         // back `None`. Otherwise, it'll send back `Some(...)` with the socket.
                         Ok(maybe_socket) => match maybe_socket {
-                            Some(socket) => UdpServiceState::Connected(socket),
-                            None => UdpServiceState::Disconnected,
+                            Some(socket) => ServiceState::Connected(socket),
+                            None => ServiceState::Disconnected,
                         },
                         Err(_) => return Poll::Ready(Err(NetError::ServiceSocketChannelClosed)),
                     }
@@ -187,8 +168,8 @@ impl Service<Vec<u8>> for UdpService {
     fn call(&mut self, buf: Vec<u8>) -> Self::Future {
         let (tx, rx) = oneshot::channel();
 
-        let socket = match std::mem::replace(&mut self.state, UdpServiceState::Sending(rx)) {
-            UdpServiceState::Connected(socket) => socket,
+        let socket = match std::mem::replace(&mut self.state, ServiceState::Sending(rx)) {
+            ServiceState::Connected(socket) => socket,
             _ => panic!("poll_ready must be called first"),
         };
 
diff --git a/src/sinks/util/service/net/unix.rs b/src/sinks/util/service/net/unix.rs
@@ -6,7 +6,8 @@ use std::{
     time::Duration,
 };
 
-use futures::{future::BoxFuture, FutureExt};
+use futures::future::BoxFuture;
+use futures_util::FutureExt;
 use snafu::ResultExt;
 use tokio::{
     io::AsyncWriteExt,
@@ -26,7 +27,7 @@ use crate::{
     sinks::{util::retries::ExponentialBackoff, Healthcheck},
 };
 
-use super::{NetError, net_error::*};
+use super::{net_error::*, NetError, ServiceState};
 
 /// Unix socket modes.
 #[configurable_component]
@@ -120,14 +121,16 @@ pub struct UnixConnector {
 impl UnixConnector {
     async fn connect(&self) -> Result<(PathBuf, UnixEither), NetError> {
         let either_socket = match self.mode {
-            UnixMode::Datagram => UnixDatagram::unbound()
-                .context(FailedToBind)
-                .and_then(|datagram| {
-                    datagram
-                        .connect(&self.path)
-                        .context(FailedToConnect)
-                        .map(|_| UnixEither::Datagram(datagram))
-                })?,
+            UnixMode::Datagram => {
+                UnixDatagram::unbound()
+                    .context(FailedToBind)
+                    .and_then(|datagram| {
+                        datagram
+                            .connect(&self.path)
+                            .context(FailedToConnect)
+                            .map(|_| UnixEither::Datagram(datagram))
+                    })?
+            }
             UnixMode::Stream => UnixStream::connect(&self.path)
                 .await
                 .context(FailedToConnect)
@@ -175,34 +178,16 @@ impl UnixConnector {
     }
 }
 
-enum UnixServiceState {
-    /// The service is currently disconnected.
-    Disconnected,
-
-    /// The service is currently attempting to connect to the endpoint.
-    Connecting(BoxFuture<'static, UnixEither>),
-
-    /// The service is connected and idle.
-    Connected(UnixEither),
-
-    /// The service has an in-flight send to the socket.
-    ///
-    /// If the socket experiences an unrecoverable error during the send, `None` will be returned
-    /// over the channel to signal the need to establish a new connection rather than reusing the
-    /// existing connection.
-    Sending(oneshot::Receiver<Option<UnixEither>>),
-}
-
 pub struct UnixService {
     connector: UnixConnector,
-    state: UnixServiceState,
+    state: ServiceState<UnixEither>,
 }
 
 impl UnixService {
     const fn new(connector: UnixConnector) -> Self {
         Self {
             connector,
-            state: UnixServiceState::Disconnected,
+            state: ServiceState::Disconnected,
         }
     }
 }
@@ -215,24 +200,24 @@ impl Service<Vec<u8>> for UnixService {
     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         loop {
             self.state = match &mut self.state {
-                UnixServiceState::Disconnected => {
+                ServiceState::Disconnected => {
                     let connector = self.connector.clone();
-                    UnixServiceState::Connecting(Box::pin(async move {
-                        connector.connect_backoff().await
-                    }))
+                    ServiceState::Connecting(Box::pin(
+                        async move { connector.connect_backoff().await },
+                    ))
                 }
-                UnixServiceState::Connecting(fut) => {
+                ServiceState::Connecting(fut) => {
                     let socket = ready!(fut.poll_unpin(cx));
-                    UnixServiceState::Connected(socket)
+                    ServiceState::Connected(socket)
                 }
-                UnixServiceState::Connected(_) => break,
-                UnixServiceState::Sending(fut) => {
+                ServiceState::Connected(_) => break,
+                ServiceState::Sending(fut) => {
                     match ready!(fut.poll_unpin(cx)) {
                         // When a send concludes, and there's an error, the request future sends
                         // back `None`. Otherwise, it'll send back `Some(...)` with the socket.
                         Ok(maybe_socket) => match maybe_socket {
-                            Some(socket) => UnixServiceState::Connected(socket),
-                            None => UnixServiceState::Disconnected,
+                            Some(socket) => ServiceState::Connected(socket),
+                            None => ServiceState::Disconnected,
                         },
                         Err(_) => return Poll::Ready(Err(NetError::ServiceSocketChannelClosed)),
                     }
@@ -245,8 +230,8 @@ impl Service<Vec<u8>> for UnixService {
     fn call(&mut self, buf: Vec<u8>) -> Self::Future {
         let (tx, rx) = oneshot::channel();
 
-        let mut socket = match std::mem::replace(&mut self.state, UnixServiceState::Sending(rx)) {
-            UnixServiceState::Connected(socket) => socket,
+        let mut socket = match std::mem::replace(&mut self.state, ServiceState::Sending(rx)) {
+            ServiceState::Connected(socket) => socket,
             _ => panic!("poll_ready must be called first"),
         };
 

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,5 @@`
	`1`	`+//! Networking-related helper functions.`
	`2`	`+`
`1`	`3`	`use std::{io, os::fd::AsRawFd, time::Duration};`
`2`	`4`
`3`	`5`	`use socket2::{SockRef, TcpKeepalive};`