From 2f9f5cc5cf09f076601130749411e697c7158aa3 Mon Sep 17 00:00:00 2001
From: Nils Andreas Svee <me@lochnair.net>
Date: Mon, 10 Aug 2026 23:29:21 +0700
Subject: [PATCH] fix(cli): account for targets in stats memory warning

---
 crates/irtt-cli/src/cmd/client/run.rs | 53 +++++++++++++++++++++++++--
 1 file changed, 49 insertions(+), 4 deletions(-)

diff --git a/crates/irtt-cli/src/cmd/client/run.rs b/crates/irtt-cli/src/cmd/client/run.rs
index 6f144be..185178e 100644
--- a/crates/irtt-cli/src/cmd/client/run.rs
+++ b/crates/irtt-cli/src/cmd/client/run.rs
@@ -69,14 +69,14 @@ pub fn run_stream(
     )
     .map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err))?;
     let continuous = args.is_continuous();
+    let target_count = targets.len();
     #[cfg(feature = "stats")]
-    if let Some(warning) = finite_stats_memory_warning(&args) {
+    if let Some(warning) = finite_stats_memory_warning(&args, target_count) {
         eprintln!("{warning}");
     }
     if is_shutdown_requested(shutdown_requested) {
         return Ok(());
     }
-    let target_count = targets.len();
     let config = ManagedClientConfig {
         client: args.to_client_config(),
         pacing: args.pacing.into(),
@@ -409,12 +409,14 @@ fn stats_config(continuous: bool) -> StatsConfig {
     }
 }
 #[cfg(feature = "stats")]
-fn finite_stats_memory_warning(args: &ClientArgs) -> Option<String> {
+fn finite_stats_memory_warning(args: &ClientArgs, target_count: usize) -> Option<String> {
     if args.is_continuous() || args.duration.is_zero() {
         return None;
     }
+    let target_count = u64::try_from(target_count).unwrap_or(u64::MAX);
     let estimated_bytes = expected_probe_count(args.duration, args.interval)
-        .saturating_mul(FINITE_STATS_BYTES_PER_PROBE);
+        .saturating_mul(FINITE_STATS_BYTES_PER_PROBE)
+        .saturating_mul(target_count);
     if estimated_bytes < FINITE_STATS_MEMORY_WARNING_BYTES {
         return None;
     }
@@ -441,6 +443,49 @@ mod tests {
     };
     use std::sync::Arc;
 
+    #[cfg(feature = "stats")]
+    #[test]
+    fn finite_stats_memory_warning_accounts_for_target_count() {
+        use clap::Parser;
+
+        let args = ClientArgs::try_parse_from([
+            "irtt-cli",
+            "--duration",
+            "4200s",
+            "--interval",
+            "100ms",
+            "127.0.0.1:2112",
+        ])
+        .unwrap();
+
+        assert!(
+            finite_stats_memory_warning(&args, 1).is_none(),
+            "a single target's estimate should stay below the warning threshold"
+        );
+        assert!(
+            finite_stats_memory_warning(&args, 7).is_some(),
+            "seven targets' aggregate estimate should cross the warning threshold"
+        );
+    }
+
+    #[cfg(feature = "stats")]
+    #[test]
+    fn finite_stats_memory_warning_saturates_target_count_multiplication() {
+        use clap::Parser;
+
+        let args = ClientArgs::try_parse_from([
+            "irtt-cli",
+            "--duration",
+            "4200s",
+            "--interval",
+            "100ms",
+            "127.0.0.1:2112",
+        ])
+        .unwrap();
+
+        assert!(finite_stats_memory_warning(&args, usize::MAX).is_some());
+    }
+
     #[test]
     fn lagged_events_are_counted_saturatingly() {
         let mut dropped = 0_u64;