blocks

Author: TheSafo

lading/src/generator/common.rs

@@ -36,7 +36,8 @@ use lading_payload::block;
 
 use super::General;
 use crate::generator::common::{
-    BytesThrottleConfig, MetricsBuilder, ThrottleConversionError, create_throttle,
+    BlockThrottle, MetricsBuilder, ThrottleConfig, ThrottleConversionError, ThrottleMode,
+    create_throttle,
 };
 
 /// An enum to allow us to determine what operation caused an IO errror as the
@@ -137,6 +138,7 @@ pub struct Config {
     /// Sets the [`crate::payload::Config`] of this template.
     pub variant: lading_payload::Config,
     /// Defines the number of bytes that written in each log file.
+    #[deprecated(note = "Use load_profile bytes-per-second instead")]
     bytes_per_second: Option<Byte>,
     /// Defines the maximum internal cache of this log target. `file_gen` will
     /// pre-build its outputs up to the byte capacity specified here.
@@ -147,8 +149,9 @@ pub struct Config {
     /// Whether to use a fixed or streaming block cache
     #[serde(default = "lading_payload::block::default_cache_method")]
     block_cache_method: block::CacheMethod,
-    /// The load throttle configuration
-    pub throttle: Option<BytesThrottleConfig>,
+    /// Throughput profile controlling emission rate (bytes or blocks).
+    #[serde(default)]
+    pub load_profile: Option<ThrottleConfig>,
 }
 
 #[derive(Debug)]
@@ -214,8 +217,13 @@ impl Server {
         let mut handles = Vec::new();
 
         for idx in 0..config.concurrent_logs {
-            let throttle =
-                create_throttle(config.throttle.as_ref(), config.bytes_per_second.as_ref())?;
+            let legacy_bps = {
+                #[allow(deprecated)]
+                {
+                    config.bytes_per_second.as_ref()
+                }
+            };
+            let throughput_throttle = create_throttle(config.load_profile.as_ref(), legacy_bps)?;
 
             let mut dir_path = config.root.clone();
             let depth = rng.random_range(0..config.max_depth);
@@ -234,8 +242,9 @@ impl Server {
                 &basename,
                 config.total_rotations,
                 maximum_bytes_per_log,
+                maximum_block_size.get(),
                 Arc::clone(&block_cache),
-                throttle,
+                throughput_throttle,
                 shutdown.clone(),
                 child_labels,
             );
@@ -283,8 +292,9 @@ struct Child {
     names: Vec<PathBuf>,
     // The soft limit bytes per file that will trigger a rotation.
     maximum_bytes_per_log: NonZeroU32,
+    maximum_block_size: u32,
     block_cache: Arc<block::Cache>,
-    throttle: lading_throttle::Throttle,
+    throttle: BlockThrottle,
     shutdown: lading_signal::Watcher,
     labels: Vec<(String, String)>,
 }
@@ -295,8 +305,9 @@ impl Child {
         basename: &Path,
         total_rotations: u8,
         maximum_bytes_per_log: NonZeroU32,
+        maximum_block_size: u32,
         block_cache: Arc<block::Cache>,
-        throttle: lading_throttle::Throttle,
+        throttle: BlockThrottle,
         shutdown: lading_signal::Watcher,
         labels: Vec<(String, String)>,
     ) -> Self {
@@ -316,6 +327,7 @@ impl Child {
         Self {
             names,
             maximum_bytes_per_log,
+            maximum_block_size,
             block_cache,
             throttle,
             shutdown,
@@ -324,7 +336,11 @@ impl Child {
     }
 
     async fn spin(mut self) -> Result<(), Error> {
-        let buffer_capacity = self.throttle.maximum_capacity() as usize;
+        let mut handle = self.block_cache.handle();
+        let buffer_capacity = match self.throttle.mode {
+            ThrottleMode::Bytes => self.throttle.maximum_capacity() as usize,
+            ThrottleMode::Blocks => self.maximum_block_size as usize,
+        };
         let mut total_bytes_written: u64 = 0;
         let maximum_bytes_per_log: u64 = u64::from(self.maximum_bytes_per_log.get());
 
@@ -357,21 +373,16 @@ impl Child {
                 })?,
         );
 
-        let mut handle = self.block_cache.handle();
-
         let shutdown_wait = self.shutdown.recv();
         tokio::pin!(shutdown_wait);
         loop {
             // SAFETY: By construction the block cache will never be empty
             // except in the event of a catastrophic failure.
-            let total_bytes = self.block_cache.peek_next_size(&handle);
-
             tokio::select! {
-                result = self.throttle.wait_for(total_bytes) => {
+                result = self.throttle.wait_for_block(&self.block_cache, &handle) => {
                     match result {
                         Ok(()) => {
-                            let block = self.block_cache.advance(&mut handle);
-                            write_bytes(block,
+                            write_bytes(self.block_cache.advance(&mut handle),
                                     &mut fp,
                                     &mut total_bytes_written,
                                     buffer_capacity,