Add option to enabled/disable client stats and fix #1771 (#2001)

* Add option to enabled/disable client stats and fix #1771

* more fix

* fix map_density

* even more fix

* remove need for vec in Aggregator::aggregate

* fix json weirdness - remove individual clients (is that all right? )

* Make pretty

libafl/src/monitors/disk.rs

@@ -72,7 +72,7 @@ executions = {}
 exec_sec = {}
 ",
                 format_duration_hms(&(cur_time - self.start_time())),
-                self.client_stats().len(),
+                self.client_stats_count(),
                 self.corpus_size(),
                 self.objective_size(),
                 self.total_execs(),
@@ -211,12 +211,11 @@ where
 
             let line = json!({
                 "run_time": current_time() - self.base.start_time(),
-                "clients": self.base.client_stats().len(),
+                "clients": self.client_stats_count(),
                 "corpus": self.base.corpus_size(),
                 "objectives": self.base.objective_size(),
                 "executions": self.base.total_execs(),
                 "exec_sec": self.base.execs_per_sec(),
-                "clients": &self.client_stats().get(1..)
             });
             writeln!(&file, "{line}").expect("Unable to write JSON to file");
         }

libafl/src/monitors/mod.rs

@@ -62,22 +62,20 @@ impl Aggregator {
 
     /// takes the key and the ref to clients stats then aggregate them all.
     fn aggregate(&mut self, name: &str, client_stats: &[ClientStats]) {
-        let mut gather = vec![];
+        let mut gather = client_stats
+            .iter()
+            .filter_map(|client| client.user_monitor.get(name));
 
-        for client in client_stats {
-            if let Some(x) = client.user_monitor.get(name) {
-                gather.push(x);
-            }
-        }
+        let gather_count = gather.clone().count();
 
-        let (mut init, op) = match gather.first() {
+        let (mut init, op) = match gather.next() {
             Some(x) => (x.value().clone(), x.aggregator_op().clone()),
             _ => {
                 return;
             }
         };
 
-        for item in gather.iter().skip(1) {
+        for item in gather {
             match op {
                 AggregatorOps::None => {
                     // Nothing
@@ -112,7 +110,7 @@ impl Aggregator {
 
         if let AggregatorOps::Avg = op {
             // if avg then divide last.
-            init = match init.stats_div(gather.len()) {
+            init = match init.stats_div(gather_count) {
                 Some(x) => x,
                 _ => {
                     return;
@@ -340,6 +338,8 @@ fn prettify_float(value: f64) -> String {
 /// A simple struct to keep track of client monitor
 #[derive(Debug, Clone, Default, Serialize, Deserialize)]
 pub struct ClientStats {
+    /// If this client is enabled. This is set to `true` the first time we see this client.
+    pub enabled: bool,
     // monitor (maybe we need a separated struct?)
     /// The corpus size for this client
     pub corpus_size: u64,
@@ -506,6 +506,14 @@ pub trait Monitor {
             .fold(0_u64, |acc, x| acc + x.corpus_size)
     }
 
+    /// Count the number of enabled client stats
+    fn client_stats_count(&self) -> usize {
+        self.client_stats()
+            .iter()
+            .filter(|client| client.enabled)
+            .count()
+    }
+
     /// Amount of elements in the objectives (combined for all children)
     fn objective_size(&self) -> u64 {
         self.client_stats()
@@ -538,14 +546,23 @@ pub trait Monitor {
 
     /// The client monitor for a specific id, creating new if it doesn't exist
     fn client_stats_insert(&mut self, client_id: ClientId) {
-        let client_stat_count = self.client_stats().len();
-        for _ in client_stat_count..(client_id.0 + 1) as usize {
+        let total_client_stat_count = self.client_stats().len();
+        for _ in total_client_stat_count..=(client_id.0) as usize {
             self.client_stats_mut().push(ClientStats {
-                last_window_time: current_time(),
-                start_time: current_time(),
+                enabled: false,
+                last_window_time: Duration::from_secs(0),
+                start_time: Duration::from_secs(0),
                 ..ClientStats::default()
             });
         }
+        let new_stat = self.client_stats_mut_for(client_id);
+        if !new_stat.enabled {
+            let timestamp = current_time();
+            // I have never seen this man in my life
+            new_stat.start_time = timestamp;
+            new_stat.last_window_time = timestamp;
+            new_stat.enabled = true;
+        }
     }
 
     /// Get mutable reference to client stats
@@ -673,7 +690,7 @@ impl Monitor for SimplePrintingMonitor {
             event_msg,
             sender_id.0,
             format_duration_hms(&(current_time() - self.start_time)),
-            self.client_stats().len(),
+            self.client_stats_count(),
             self.corpus_size(),
             self.objective_size(),
             self.total_execs(),
@@ -749,7 +766,7 @@ where
             event_msg,
             sender_id.0,
             format_duration_hms(&(current_time() - self.start_time)),
-            self.client_stats().len(),
+            self.client_stats_count(),
             self.corpus_size(),
             self.objective_size(),
             self.total_execs(),

libafl/src/monitors/multi.rs

@@ -75,7 +75,7 @@ where
             "[{}]  (GLOBAL) run time: {}, clients: {}, corpus: {}, objectives: {}, executions: {}, exec/sec: {}",
             head,
             format_duration_hms(&(current_time() - self.start_time)),
-            self.client_stats().len(),
+            self.client_stats_count(),
             self.corpus_size(),
             self.objective_size(),
             self.total_execs(),
@@ -106,7 +106,7 @@ where
         #[cfg(feature = "introspection")]
         {
             // Print the client performance monitor. Skip the Client 0 which is the broker
-            for (i, client) in self.client_stats.iter().skip(1).enumerate() {
+            for (i, client) in self.client_stats.iter().filter(|x| x.enabled).enumerate() {
                 let fmt = format!("Client {:03}:\n{}", i + 1, client.introspection_monitor);
                 (self.print_fn)(&fmt);
             }

libafl/src/monitors/prometheus.rs

@@ -142,7 +142,7 @@ where
                 stat: String::new(),
             })
             .set(run_time.try_into().unwrap()); // run time in seconds, which can be converted to a time format by Grafana or similar
-        let total_clients = self.client_stats().len().try_into().unwrap(); // convert usize to u64 (unlikely that # of clients will be > 2^64 -1...)
+        let total_clients = self.client_stats_count().try_into().unwrap(); // convert usize to u64 (unlikely that # of clients will be > 2^64 -1...)
         self.clients_count
             .get_or_create(&Labels {
                 client: sender_id.0,
@@ -156,7 +156,7 @@ where
             event_msg,
             sender_id.0,
             format_duration_hms(&(current_time() - self.start_time)),
-            self.client_stats().len(),
+            self.client_stats_count(),
             self.corpus_size(),
             self.objective_size(),
             self.total_execs(),

libafl/src/monitors/tui/mod.rs

@@ -1,6 +1,7 @@
 //! Monitor based on ratatui
 
 use alloc::{boxed::Box, string::ToString};
+use core::cmp;
 use std::{
     collections::VecDeque,
     fmt::Write as _,
@@ -336,12 +337,16 @@ pub struct TuiMonitor {
 }
 
 impl Monitor for TuiMonitor {
-    /// the client monitor, mutable
+    /// The client monitor, mutable
+    /// This also includes disabled "padding" clients.
+    /// Results should be filtered by `.enabled`.
     fn client_stats_mut(&mut self) -> &mut Vec<ClientStats> {
         &mut self.client_stats
     }
 
-    /// the client monitor
+    /// The client monitor
+    /// This also includes disabled "padding" clients.
+    /// Results should be filtered by `.enabled`.
     fn client_stats(&self) -> &[ClientStats] {
         &self.client_stats
     }
@@ -419,8 +424,8 @@ impl Monitor for TuiMonitor {
 
         #[cfg(feature = "introspection")]
         {
-            // Print the client performance monitor. Skip the Client 0 which is the broker
-            for (i, client) in self.client_stats.iter().skip(1).enumerate() {
+            // Print the client performance monitor. Skip the Client IDs that have never sent anything.
+            for (i, client) in self.client_stats.iter().filter(|x| x.enabled).enumerate() {
                 self.context
                     .write()
                     .unwrap()
@@ -484,25 +489,12 @@ impl TuiMonitor {
     }
 
     fn map_density(&self) -> String {
-        if self.client_stats.len() < 2 {
-            return "0%".to_string();
-        }
-        let mut max_map_density = self
-            .client_stats()
-            .get(1)
-            .unwrap()
-            .get_user_stats("edges")
-            .map_or("0%".to_string(), ToString::to_string);
-
-        for client in self.client_stats().iter().skip(2) {
-            let client_map_density = client
-                .get_user_stats("edges")
-                .map_or(String::new(), ToString::to_string);
-            if client_map_density > max_map_density {
-                max_map_density = client_map_density;
-            }
-        }
-        max_map_density
+        self.client_stats()
+            .iter()
+            .filter(|client| client.enabled)
+            .filter_map(|client| client.get_user_stats("edges"))
+            .map(ToString::to_string)
+            .fold("0%".to_string(), cmp::max)
     }
 
     fn item_geometry(&self) -> ItemGeometry {
@@ -512,7 +504,7 @@ impl TuiMonitor {
         }
         let mut ratio_a: u64 = 0;
         let mut ratio_b: u64 = 0;
-        for client in self.client_stats().iter().skip(1) {
+        for client in self.client_stats().iter().filter(|client| client.enabled) {
             let afl_stats = client
                 .get_user_stats("AflStats")
                 .map_or("None".to_string(), ToString::to_string);
@@ -555,7 +547,7 @@ impl TuiMonitor {
         if self.client_stats.len() > 1 {
             let mut new_path_time = Duration::default();
             let mut new_objectives_time = Duration::default();
-            for client in self.client_stats().iter().skip(1) {
+            for client in self.client_stats().iter().filter(|client| client.enabled) {
                 new_path_time = client.last_corpus_time.max(new_path_time);
                 new_objectives_time = client.last_objective_time.max(new_objectives_time);
             }