parallel_connections for Unix Stream generator (#1421)

Author: blt

CHANGELOG.md

@@ -11,6 +11,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   'capacity' exists in the throttle struct itself, without breaking use-cases
   where bytes-per-second are specified directly. bytes-per-second implies a
   stable throttle.
+- `unix_stream` generator now supports `parallel_connections` in a manner similar to
+  `unix_datagram`.
 
 ## [0.26.0]
 ## Added

lading/src/generator.rs

@@ -197,7 +197,7 @@ impl Server {
 
                 Self::UnixStream(unix_stream::UnixStream::new(
                     config.general,
-                    conf,
+                    &conf,
                     shutdown,
                 )?)
             }

lading/src/generator/unix_stream.rs

@@ -18,11 +18,19 @@ use metrics::counter;
 use rand::{SeedableRng, rngs::StdRng};
 use serde::{Deserialize, Serialize};
 use std::{num::NonZeroU32, path::PathBuf, thread};
-use tokio::{net, sync::mpsc, task::JoinError};
+use tokio::{
+    net,
+    sync::{broadcast::Receiver, mpsc},
+    task::{JoinError, JoinSet},
+};
 use tracing::{debug, error, info, warn};
 
 use super::General;
 
+fn default_parallel_connections() -> u16 {
+    1
+}
+
 #[derive(Debug, Deserialize, Serialize, PartialEq)]
 #[serde(deny_unknown_fields)]
 /// Configuration of this generator.
@@ -33,7 +41,7 @@ pub struct Config {
     pub path: PathBuf,
     /// The payload variant
     pub variant: lading_payload::Config,
-    /// The bytes per second to send or receive from the target
+    /// The bytes per second to send or receive from the target, per connection.
     pub bytes_per_second: Option<byte_unit::Byte>,
     /// The maximum size in bytes of the cache of prebuilt messages
     pub maximum_prebuild_cache_size_bytes: byte_unit::Byte,
@@ -43,6 +51,10 @@ pub struct Config {
     /// Whether to use a fixed or streaming block cache
     #[serde(default = "lading_payload::block::default_cache_method")]
     pub block_cache_method: block::CacheMethod,
+    /// The total number of parallel connections to maintain, see documentation
+    /// on `bytes_per_second`.
+    #[serde(default = "default_parallel_connections")]
+    pub parallel_connections: u16,
     /// The load throttle configuration
     pub throttle: Option<crate::generator::common::BytesThrottleConfig>,
 }
@@ -59,6 +71,15 @@ pub enum Error {
     /// Subtask error
     #[error("Subtask failure: {0}")]
     Subtask(#[from] JoinError),
+    /// Child sub-task error.
+    #[error("Child join error: {0}")]
+    Child(JoinError),
+    /// Startup send error.
+    #[error("Startup send error: {0}")]
+    StartupSend(#[from] tokio::sync::broadcast::error::SendError<()>),
+    /// Child startup wait error.
+    #[error("Child startup wait error: {0}")]
+    StartupWait(#[from] tokio::sync::broadcast::error::RecvError),
     /// Byte error
     #[error("Bytes must not be negative: {0}")]
     Byte(#[from] byte_unit::ParseError),
@@ -82,11 +103,9 @@ pub enum Error {
 /// This generator is responsible for sending data to the target via UDS
 /// streams.
 pub struct UnixStream {
-    path: PathBuf,
-    throttle: Throttle,
-    block_cache: block::Cache,
-    metric_labels: Vec<(String, String)>,
+    handles: JoinSet<Result<(), Error>>,
     shutdown: lading_signal::Watcher,
+    startup: tokio::sync::broadcast::Sender<()>,
 }
 
 impl UnixStream {
@@ -103,7 +122,7 @@ impl UnixStream {
     #[allow(clippy::cast_possible_truncation)]
     pub fn new(
         general: General,
-        config: Config,
+        config: &Config,
         shutdown: lading_signal::Watcher,
     ) -> Result<Self, Error> {
         let mut rng = StdRng::from_seed(config.seed);
@@ -115,37 +134,51 @@ impl UnixStream {
             labels.push(("id".to_string(), id));
         }
 
-        let throttle_config = match (config.bytes_per_second, &config.throttle) {
-            (Some(bytes_per_second), None) => {
-                let bytes_per_second =
-                    NonZeroU32::new(bytes_per_second.as_u128() as u32).ok_or(Error::Zero)?;
-                lading_throttle::Config::Stable {
-                    maximum_capacity: bytes_per_second,
+        let (startup, _startup_rx) = tokio::sync::broadcast::channel(1);
+
+        let mut handles = JoinSet::new();
+        for _ in 0..config.parallel_connections {
+            let throttle_config = match (config.bytes_per_second, &config.throttle) {
+                (Some(bytes_per_second), None) => {
+                    let bytes_per_second =
+                        NonZeroU32::new(bytes_per_second.as_u128() as u32).ok_or(Error::Zero)?;
+                    lading_throttle::Config::Stable {
+                        maximum_capacity: bytes_per_second,
+                    }
                 }
-            }
-            (None, Some(throttle)) => (*throttle).try_into()?,
-            (Some(_), Some(_)) => return Err(Error::ConflictingThrottleConfig),
-            (None, None) => return Err(Error::NoThrottleConfig),
-        };
+                (None, Some(throttle)) => (*throttle).try_into()?,
+                (Some(_), Some(_)) => return Err(Error::ConflictingThrottleConfig),
+                (None, None) => return Err(Error::NoThrottleConfig),
+            };
+            let throttle = Throttle::new_with_config(throttle_config);
+
+            let total_bytes =
+                NonZeroU32::new(config.maximum_prebuild_cache_size_bytes.as_u128() as u32)
+                    .ok_or(Error::Zero)?;
+            let block_cache = match config.block_cache_method {
+                block::CacheMethod::Fixed => block::Cache::fixed(
+                    &mut rng,
+                    total_bytes,
+                    config.maximum_block_size.as_u128(),
+                    &config.variant,
+                )?,
+            };
+
+            let child = Child {
+                path: config.path.clone(),
+                block_cache,
+                throttle,
+                metric_labels: labels.clone(),
+                shutdown: shutdown.clone(),
+            };
 
-        let total_bytes =
-            NonZeroU32::new(config.maximum_prebuild_cache_size_bytes.as_u128() as u32)
-                .ok_or(Error::Zero)?;
-        let block_cache = match config.block_cache_method {
-            block::CacheMethod::Fixed => block::Cache::fixed(
-                &mut rng,
-                total_bytes,
-                config.maximum_block_size.as_u128(),
-                &config.variant,
-            )?,
-        };
+            handles.spawn(child.spin(startup.subscribe()));
+        }
 
         Ok(Self {
-            path: config.path,
-            block_cache,
-            throttle: Throttle::new_with_config(throttle_config),
-            metric_labels: labels,
+            handles,
             shutdown,
+            startup,
         })
     }
 
@@ -159,6 +192,32 @@ impl UnixStream {
     ///
     /// Function will panic if underlying byte capacity is not available.
     pub async fn spin(mut self) -> Result<(), Error> {
+        self.startup.send(())?;
+        self.shutdown.recv().await;
+        info!("shutdown signal received");
+        while let Some(res) = self.handles.join_next().await {
+            match res {
+                Ok(Ok(())) => continue,
+                Ok(Err(err)) => return Err(err),
+                Err(err) => return Err(Error::Child(err)),
+            }
+        }
+        Ok(())
+    }
+}
+
+#[derive(Debug)]
+struct Child {
+    path: PathBuf,
+    throttle: Throttle,
+    block_cache: block::Cache,
+    metric_labels: Vec<(String, String)>,
+    shutdown: lading_signal::Watcher,
+}
+
+impl Child {
+    async fn spin(mut self, mut startup_receiver: Receiver<()>) -> Result<(), Error> {
+        startup_receiver.recv().await?;
         debug!("UnixStream generator running");
 
         // Move the block_cache into an OS thread, exposing a channel between it