Generators obey throttle wait_for errors (#1317)

blt · web-flow · commit 31ade89ffe1b · 2025-04-10T14:42:14.000-07:00
### What does this PR do?

This commit fixes a long-standing mystery: why when throughputs are low does
lading _sometimes_ emit faster than it should. It turns out the answer is
we failed to handle the capacity error that the throttle kicks out in such
conditions systematically.

This also implies that unless the user explicitly sets a max block size then
we should configure it for them, so the default of a flat 1MiB does not serve
well. I'll fix that in a follow-on PR.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,11 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [Unreleased]
+## Fixed
+- Fixed throttle capacity validation to prevent requests larger than maximum
+  capacity from being emitted.
+
 ## [0.25.8]
 ## Fixed
 - Lading will now ignore child processes when polling /proc if the children are
diff --git a/lading/src/generator/file_gen/logrotate.rs b/lading/src/generator/file_gen/logrotate.rs
@@ -368,17 +368,24 @@ impl Child {
             let blk = rcv.peek().await.expect("block cache should never be empty");
 
             tokio::select! {
-                _ = self.throttle.wait_for(blk.total_bytes) => {
-                    let blk = rcv.next().await.expect("failed to advance through the blocks");
-                    write_bytes(&blk,
-                                &mut fp,
-                                &mut total_bytes_written,
-                                bytes_per_second,
-                                total_names,
-                                maximum_bytes_per_log,
-                                &self.names,
-                                last_name,
-                                &self.labels).await?;
+                result = self.throttle.wait_for(blk.total_bytes) => {
+                    match result {
+                        Ok(()) => {
+                            let blk = rcv.next().await.expect("failed to advance through the blocks");
+                            write_bytes(&blk,
+                                    &mut fp,
+                                    &mut total_bytes_written,
+                                    bytes_per_second,
+                                    total_names,
+                                    maximum_bytes_per_log,
+                                    &self.names,
+                                    last_name,
+                                    &self.labels).await?;
+                        }
+                        Err(err) => {
+                            error!("Throttle request of {} is larger than throttle capacity. Block will be discarded. Error: {}", blk.total_bytes, err);
+                        }
+                    }
                 }
                 () = &mut shutdown_wait => {
                     fp.flush().await.map_err(|err| Error::IoFlush { err })?;
@@ -387,7 +394,6 @@ impl Child {
                     info!("shutdown signal received");
                     return Ok(());
                 },
-
             }
         }
     }
diff --git a/lading/src/generator/file_gen/traditional.rs b/lading/src/generator/file_gen/traditional.rs
@@ -35,7 +35,7 @@ use tokio::{
     sync::mpsc,
     task::{JoinError, JoinHandle},
 };
-use tracing::info;
+use tracing::{error, info};
 
 use crate::common::PeekableReceiver;
 use lading_payload::{
@@ -263,38 +263,45 @@ impl Child {
             let total_bytes = blk.total_bytes;
 
             tokio::select! {
-                _ = self.throttle.wait_for(total_bytes) => {
-                    let blk = rcv.next().await.expect("failed to advance through blocks"); // actually advance through the blocks
-                    let total_bytes = u64::from(total_bytes.get());
+                result = self.throttle.wait_for(total_bytes) => {
+                    match result {
+                        Ok(()) => {
+                            let blk = rcv.next().await.expect("failed to advance through blocks"); // actually advance through the blocks
+                            let total_bytes = u64::from(total_bytes.get());
 
-                    {
-                        fp.write_all(&blk.bytes).await?;
-                        counter!("bytes_written").increment(total_bytes);
-                        total_bytes_written += total_bytes;
-                    }
+                            {
+                                fp.write_all(&blk.bytes).await?;
+                                counter!("bytes_written").increment(total_bytes);
+                                total_bytes_written += total_bytes;
+                            }
 
-                    if total_bytes_written > maximum_bytes_per_file {
-                        fp.flush().await?;
-                        if self.rotate {
-                            // Delete file, leaving any open file handlers intact. This
-                            // includes our own `fp` for the time being.
-                            fs::remove_file(&path).await?;
+                            if total_bytes_written > maximum_bytes_per_file {
+                                fp.flush().await?;
+                                if self.rotate {
+                                    // Delete file, leaving any open file handlers intact. This
+                                    // includes our own `fp` for the time being.
+                                    fs::remove_file(&path).await?;
+                                }
+                                // Update `path` to point to the next indexed file.
+                                file_index = self.file_index.fetch_add(1, Ordering::Relaxed);
+                                path = path_from_template(&self.path_template, file_index);
+                                // Open a new fp to `path`, replacing `fp`. Any holders of the
+                                // file pointer still have it but the file no longer has a name.
+                                fp = BufWriter::with_capacity(
+                                    bytes_per_second,
+                                    fs::OpenOptions::new()
+                                        .create(true)
+                                        .truncate(false)
+                                        .write(true)
+                                        .open(&path)
+                                        .await?,
+                                );
+                                total_bytes_written = 0;
+                            }
+                        }
+                        Err(err) => {
+                            error!("Throttle request of {total_bytes} is larger than throttle capacity. Block will be discarded. Error: {err}");
                         }
-                        // Update `path` to point to the next indexed file.
-                        file_index = self.file_index.fetch_add(1, Ordering::Relaxed);
-                        path = path_from_template(&self.path_template, file_index);
-                        // Open a new fp to `path`, replacing `fp`. Any holders of the
-                        // file pointer still have it but the file no longer has a name.
-                        fp = BufWriter::with_capacity(
-                            bytes_per_second,
-                            fs::OpenOptions::new()
-                                .create(true)
-                                .truncate(false)
-                                .write(true)
-                                .open(&path)
-                                .await?,
-                        );
-                        total_bytes_written = 0;
                     }
                 }
                 () = &mut shutdown_wait => {
diff --git a/lading/src/generator/http.rs b/lading/src/generator/http.rs
@@ -22,7 +22,7 @@ use once_cell::sync::OnceCell;
 use rand::{SeedableRng, prelude::StdRng};
 use serde::{Deserialize, Serialize};
 use tokio::sync::{Semaphore, mpsc};
-use tracing::info;
+use tracing::{error, info};
 
 use crate::common::PeekableReceiver;
 use lading_payload::block::{self, Block};
@@ -227,30 +227,38 @@ impl Http {
             }
 
             tokio::select! {
-                _ = self.throttle.wait_for(total_bytes) => {
-                    let client = client.clone();
-                    let labels = labels.clone();
+                result = self.throttle.wait_for(total_bytes) => {
+                    match result {
+                        Ok(()) => {
+                            let client = client.clone();
+                            let labels = labels.clone();
 
-                    let permit = CONNECTION_SEMAPHORE.get().expect("Connection Semaphore is being initialized or cell is empty").acquire().await.expect("Connection Semaphore has already closed");
-                    tokio::spawn(async move {
-                        counter!("requests_sent", &labels).increment(1);
-                        match client.request(request).await {
-                            Ok(response) => {
-                                counter!("bytes_written", &labels).increment(block_length as u64);
-                                let status = response.status();
-                                let mut status_labels = labels.clone();
-                                status_labels
-                                    .push(("status_code".to_string(), status.as_u16().to_string()));
-                                counter!("request_ok", &status_labels).increment(1);
-                            }
-                            Err(err) => {
-                                let mut error_labels = labels.clone();
-                                error_labels.push(("error".to_string(), err.to_string()));
-                                counter!("request_failure", &error_labels).increment(1);
-                            }
+                            let permit = CONNECTION_SEMAPHORE.get().expect("Connection Semaphore is being initialized or cell is empty").acquire().await.expect("Connection Semaphore has already closed");
+                            tokio::spawn(async move {
+                                counter!("requests_sent", &labels).increment(1);
+                                match client.request(request).await {
+                                    Ok(response) => {
+                                        counter!("bytes_written", &labels).increment(block_length as u64);
+                                        let status = response.status();
+                                        let mut status_labels = labels.clone();
+                                        status_labels
+                                            .push(("status_code".to_string(), status.as_u16().to_string()));
+                                        counter!("request_ok", &status_labels).increment(1);
+                                    }
+                                    Err(err) => {
+                                        let mut error_labels = labels.clone();
+                                        error_labels.push(("error".to_string(), err.to_string()));
+                                        counter!("request_failure", &error_labels).increment(1);
+                                    }
+                                }
+                                drop(permit);
+                            });
+
+                        }
+                        Err(err) => {
+                            error!("Throttle request of {total_bytes} is larger than throttle capacity. Block will be discarded. Error: {err}");
                         }
-                        drop(permit);
-                    });
+                    }
                 },
                 () = &mut shutdown_wait => {
                     info!("shutdown signal received");
diff --git a/lading/src/generator/splunk_hec.rs b/lading/src/generator/splunk_hec.rs
@@ -39,7 +39,7 @@ use tokio::{
     sync::{Semaphore, SemaphorePermit, mpsc},
     time::timeout,
 };
-use tracing::info;
+use tracing::{error, info};
 
 use crate::{common::PeekableReceiver, generator::splunk_hec::acknowledgements::Channel};
 use lading_payload::block::{self, Block};
@@ -282,30 +282,37 @@ impl SplunkHec {
             let total_bytes = blk.total_bytes;
 
             tokio::select! {
-                _ = self.throttle.wait_for(total_bytes) => {
-                    let client = client.clone();
-                    let labels = labels.clone();
-                    let uri = uri.clone();
+                result = self.throttle.wait_for(total_bytes) => {
+                    match result {
+                        Ok(()) => {
+                            let client = client.clone();
+                            let labels = labels.clone();
+                            let uri = uri.clone();
 
-                    let blk = rcv.next().await.expect("failed to advance through blocks"); // actually advance through the blocks
-                    let body = crate::full(blk.bytes.clone());
-                    let block_length = blk.bytes.len();
+                            let blk = rcv.next().await.expect("failed to advance through blocks"); // actually advance through the blocks
+                            let body = crate::full(blk.bytes.clone());
+                            let block_length = blk.bytes.len();
 
-                    let request = Request::builder()
-                        .method(Method::POST)
-                        .uri(uri)
-                        .header(AUTHORIZATION, format!("Splunk {}", self.token))
-                        .header(CONTENT_LENGTH, block_length)
-                        .header(SPLUNK_HEC_CHANNEL_HEADER, channel.id())
-                        .body(body)?;
+                            let request = Request::builder()
+                                .method(Method::POST)
+                                .uri(uri)
+                                .header(AUTHORIZATION, format!("Splunk {}", self.token))
+                                .header(CONTENT_LENGTH, block_length)
+                                .header(SPLUNK_HEC_CHANNEL_HEADER, channel.id())
+                                .body(body)?;
 
-                    // NOTE once JoinSet is in tokio stable we can make this
-                    // much, much tidier by spawning requests in the JoinSet. I
-                    // think we could also possibly have the send request return
-                    // the AckID, meaning we could just keep the channel logic
-                    // in this main loop here and avoid the AckService entirely.
-                    let permit = CONNECTION_SEMAPHORE.get().expect("Connecton Semaphore is empty or being initialized").acquire().await.expect("Semaphore has already been closed");
-                    tokio::spawn(send_hec_request(permit, block_length, labels, channel, client, request, request_shutdown.clone()));
+                            // NOTE once JoinSet is in tokio stable we can make this
+                            // much, much tidier by spawning requests in the JoinSet. I
+                            // think we could also possibly have the send request return
+                            // the AckID, meaning we could just keep the channel logic
+                            // in this main loop here and avoid the AckService entirely.
+                            let permit = CONNECTION_SEMAPHORE.get().expect("Connecton Semaphore is empty or being initialized").acquire().await.expect("Semaphore has already been closed");
+                            tokio::spawn(send_hec_request(permit, block_length, labels, channel, client, request, request_shutdown.clone()));
+                        }
+                        Err(err) => {
+                            error!("Throttle request of {total_bytes} is larger than throttle capacity. Block will be discarded. Error: {err}");
+                        }
+                    }
                 }
                 () = &mut shutdown_wait => {
                     info!("shutdown signal received");
diff --git a/lading/src/generator/tcp.rs b/lading/src/generator/tcp.rs
@@ -23,7 +23,7 @@ use metrics::{counter, gauge};
 use rand::{SeedableRng, rngs::StdRng};
 use serde::{Deserialize, Serialize};
 use tokio::{io::AsyncWriteExt, net::TcpStream, sync::mpsc};
-use tracing::{info, trace};
+use tracing::{error, info, trace};
 
 use crate::common::PeekableReceiver;
 use lading_payload::block::{self, Block};
@@ -177,20 +177,27 @@ impl Tcp {
             let total_bytes = blk.total_bytes;
 
             tokio::select! {
-                _ = self.throttle.wait_for(total_bytes) => {
-                    let blk = rcv.next().await.expect("failed to advance through the blocks"); // actually advance through the blocks
-                    match connection.write_all(&blk.bytes).await {
+                result = self.throttle.wait_for(total_bytes) => {
+                    match result {
                         Ok(()) => {
-                            counter!("bytes_written", &self.metric_labels).increment(u64::from(blk.total_bytes.get()));
-                            counter!("packets_sent", &self.metric_labels).increment(1);
+                            let blk = rcv.next().await.expect("failed to advance through the blocks"); // actually advance through the blocks
+                            match connection.write_all(&blk.bytes).await {
+                                Ok(()) => {
+                                    counter!("bytes_written", &self.metric_labels).increment(u64::from(blk.total_bytes.get()));
+                                    counter!("packets_sent", &self.metric_labels).increment(1);
+                                }
+                                Err(err) => {
+                                    trace!("write failed: {}", err);
+
+                                    let mut error_labels = self.metric_labels.clone();
+                                    error_labels.push(("error".to_string(), err.to_string()));
+                                    counter!("request_failure", &error_labels).increment(1);
+                                    current_connection = None;
+                                }
+                            }
                         }
                         Err(err) => {
-                            trace!("write failed: {}", err);
-
-                            let mut error_labels = self.metric_labels.clone();
-                            error_labels.push(("error".to_string(), err.to_string()));
-                            counter!("request_failure", &error_labels).increment(1);
-                            current_connection = None;
+                            error!("Throttle request of {total_bytes} is larger than throttle capacity. Block will be discarded. Error: {err}");
                         }
                     }
                 }
diff --git a/lading/src/generator/udp.rs b/lading/src/generator/udp.rs
diff --git a/lading/src/generator/unix_datagram.rs b/lading/src/generator/unix_datagram.rs
diff --git a/lading/src/generator/unix_stream.rs b/lading/src/generator/unix_stream.rs
