Attempt to recreate UDS stream if not available, rather than crash

This commit adjusts the UDS generator spin implementation to be more like the
UDP implementation in that we no longer crash a subtask if the UDS socket is not
available for opening yet. In instances where Datadog Agent is concerned the
socket can be available relatively late, harming reliable inspection.

Signed-off-by: Brian L. Troutwine <brian.troutwine@datadoghq.com>

Author: blt

src/generator/udp.rs

@@ -152,13 +152,14 @@ impl Udp {
             let mut blocks = self.block_cache.iter().cycle();
             let mut connection = Option::<UdpSocket>::None;
             loop {
-                let blk = blocks.next().unwrap();
-                let total_bytes = blk.total_bytes;
-                assert!(
-                    total_bytes.get() <= 65507,
-                    "UDP packet too large (over 65507 B)"
-                );
                 if let Some(sock) = &connection {
+                    let blk = blocks.next().unwrap();
+                    let total_bytes = blk.total_bytes;
+                    assert!(
+                        total_bytes.get() <= 65507,
+                        "UDP packet too large (over 65507 B)"
+                    );
+
                     tokio::task::yield_now().await;
                     self.rate_limiter.until_n_ready(total_bytes).await.unwrap();
 

src/generator/uds.rs

@@ -18,8 +18,8 @@ use std::{
     num::{NonZeroU32, NonZeroUsize},
     path::PathBuf,
 };
-use tokio::{io::Interest, net::UnixStream};
-use tracing::{debug, info};
+use tokio::{io::Interest, net::UnixStream, task::JoinError};
+use tracing::{debug, error, info, trace};
 
 #[derive(Debug, Deserialize, PartialEq, Eq)]
 /// Configuration of this generator.
@@ -51,6 +51,9 @@ pub enum Error {
     /// Generic IO error
     #[error("IO error: {0}")]
     Io(#[from] std::io::Error),
+    /// Subtask error
+    #[error("Subtask failure: {0}")]
+    Subtask(#[from] JoinError),
 }
 
 #[derive(Debug)]
@@ -142,62 +145,76 @@ impl Uds {
 
         let send_task = tokio::spawn(async move {
             let mut blocks = self.block_cache.iter().cycle();
+            let mut unix_stream = Option::<UnixStream>::None;
+
             loop {
-                // Connect to the stream interior to the spawn allows us to
-                // re-connect within the same task, does require that we unwrap
-                // failures, catastrophically failing lading.
-                let stream = UnixStream::connect(&self.path)
-                    .await
-                    .map_err(Error::Io)
-                    .unwrap();
-                let blk = blocks.next().unwrap();
-                let total_bytes = blk.total_bytes;
-
-                tokio::task::yield_now().await;
-                self.rate_limiter
-                    .until_n_ready(total_bytes)
-                    .await
-                    .map_err(Error::Governor)
-                    .unwrap();
-
-                // NOTE When we write into a unix stream it may be that only
-                // some of the written bytes make it through in which case we
-                // must cycle back around and try to write the remainder of the
-                // buffer.
-                let blk_max: usize = total_bytes.get() as usize;
-                let mut blk_offset = 0;
-                while blk_offset < blk_max {
-                    let ready = stream
-                        .ready(Interest::WRITABLE)
+                if let Some(stream) = &unix_stream {
+                    let blk = blocks.next().unwrap();
+                    let total_bytes = blk.total_bytes;
+
+                    tokio::task::yield_now().await;
+                    self.rate_limiter
+                        .until_n_ready(total_bytes)
                         .await
-                        .map_err(Error::Io)
-                        .unwrap(); // Cannot ? in a spawned task :<. Mimics UDP generator.
-                    if ready.is_writable() {
-                        // Try to write data, this may still fail with `WouldBlock`
-                        // if the readiness event is a false positive.
-                        match stream.try_write(&blk.bytes[blk_offset..]) {
-                            Ok(bytes) => {
-                                counter!("bytes_written", bytes as u64, &labels);
-                                blk_offset = bytes;
-                            }
-                            Err(ref e) if e.kind() == tokio::io::ErrorKind::WouldBlock => {
-                                // If the read side has hung up we will never
-                                // know and will keep attempting to write into
-                                // the stream. This yield means we won't hog the
-                                // whole CPU.
-                                tokio::task::yield_now().await;
-                                continue;
-                            }
-                            Err(err) => {
-                                debug!("write failed: {}", err);
+                        .map_err(Error::Governor)
+                        .unwrap();
+
+                    // NOTE When we write into a unix stream it may be that only
+                    // some of the written bytes make it through in which case we
+                    // must cycle back around and try to write the remainder of the
+                    // buffer.
+                    let blk_max: usize = total_bytes.get() as usize;
+                    let mut blk_offset = 0;
+                    while blk_offset < blk_max {
+                        let ready = stream
+                            .ready(Interest::WRITABLE)
+                            .await
+                            .map_err(Error::Io)
+                            .unwrap(); // Cannot ? in a spawned task :<. Mimics UDP generator.
+                        if ready.is_writable() {
+                            // Try to write data, this may still fail with `WouldBlock`
+                            // if the readiness event is a false positive.
+                            match stream.try_write(&blk.bytes[blk_offset..]) {
+                                Ok(bytes) => {
+                                    counter!("bytes_written", bytes as u64, &labels);
+                                    blk_offset = bytes;
+                                }
+                                Err(ref e) if e.kind() == tokio::io::ErrorKind::WouldBlock => {
+                                    // If the read side has hung up we will never
+                                    // know and will keep attempting to write into
+                                    // the stream. This yield means we won't hog the
+                                    // whole CPU.
+                                    tokio::task::yield_now().await;
+                                    continue;
+                                }
+                                Err(err) => {
+                                    debug!("write failed: {}", err);
 
-                                let mut error_labels = labels.clone();
-                                error_labels.push(("error".to_string(), err.to_string()));
-                                counter!("request_failure", 1, &error_labels);
-                                break; // while blk_offset
+                                    let mut error_labels = labels.clone();
+                                    error_labels.push(("error".to_string(), err.to_string()));
+                                    counter!("request_failure", 1, &error_labels);
+                                    break; // while blk_offset
+                                }
                             }
                         }
                     }
+                } else {
+                    // Connect to the stream interior to the spawn allows us to
+                    // re-connect within the same task, does require that we unwrap
+                    // failures, catastrophically failing lading.
+                    match UnixStream::connect(&self.path).await {
+                        Ok(sock) => {
+                            debug!("UDS socket opened for writing.");
+                            unix_stream = Some(sock);
+                        }
+                        Err(err) => {
+                            trace!("opening UDS path failed: {}", err);
+
+                            let mut error_labels = labels.clone();
+                            error_labels.push(("error".to_string(), err.to_string()));
+                            counter!("connection_failure", 1, &error_labels);
+                        }
+                    }
                 }
             }
         });