Change Monitor API for more flexibility (#2927)

* Change Monitor API for more flexibility

* Make clippy happy

* Fix broken doc link

libafl/src/events/broker_hooks/mod.rs

@@ -120,8 +120,9 @@ where
                 };
 
                 monitor.client_stats_insert(id);
-                let client = monitor.client_stats_mut_for(id);
+                monitor.update_client_stats_for(id, |client_stat| {
-                client.update_corpus_size(*corpus_size as u64);
+                    client_stat.update_corpus_size(*corpus_size as u64);
+                });
                 monitor.display(event.name(), id);
                 Ok(BrokerEventResult::Forward)
             }
@@ -132,8 +133,9 @@ where
             } => {
                 // TODO: The monitor buffer should be added on client add.
                 monitor.client_stats_insert(client_id);
-                let client = monitor.client_stats_mut_for(client_id);
+                monitor.update_client_stats_for(client_id, |client_stat| {
-                client.update_executions(*executions, *time);
+                    client_stat.update_executions(*executions, *time);
+                });
                 monitor.display(event.name(), client_id);
                 Ok(BrokerEventResult::Handled)
             }
@@ -143,8 +145,9 @@ where
                 phantom: _,
             } => {
                 monitor.client_stats_insert(client_id);
-                let client = monitor.client_stats_mut_for(client_id);
+                monitor.update_client_stats_for(client_id, |client_stat| {
-                client.update_user_stats(name.clone(), value.clone());
+                    client_stat.update_user_stats(name.clone(), value.clone());
+                });
                 monitor.aggregate(name);
                 monitor.display(event.name(), client_id);
                 Ok(BrokerEventResult::Handled)
@@ -160,13 +163,12 @@ where
 
                 // Get the client for the staterestorer ID
                 monitor.client_stats_insert(client_id);
-                let client = monitor.client_stats_mut_for(client_id);
+                monitor.update_client_stats_for(client_id, |client_stat| {
-
                     // Update the normal monitor for this client
-                client.update_executions(*executions, *time);
+                    client_stat.update_executions(*executions, *time);
-
                     // Update the performance monitor for this client
-                client.update_introspection_monitor((**introspection_monitor).clone());
+                    client_stat.update_introspection_monitor((**introspection_monitor).clone());
+                });
 
                 // Display the monitor via `.display` only on core #1
                 monitor.display(event.name(), client_id);
@@ -176,8 +178,9 @@ where
             }
             Event::Objective { objective_size, .. } => {
                 monitor.client_stats_insert(client_id);
-                let client = monitor.client_stats_mut_for(client_id);
+                monitor.update_client_stats_for(client_id, |client_stat| {
-                client.update_objective_size(*objective_size as u64);
+                    client_stat.update_objective_size(*objective_size as u64);
+                });
                 monitor.display(event.name(), client_id);
                 Ok(BrokerEventResult::Handled)
             }

libafl/src/events/simple.rs

@@ -209,9 +209,9 @@ where
         match event {
             Event::NewTestcase { corpus_size, .. } => {
                 monitor.client_stats_insert(ClientId(0));
-                monitor
+                monitor.update_client_stats_for(ClientId(0), |client_stat| {
-                    .client_stats_mut_for(ClientId(0))
+                    client_stat.update_corpus_size(*corpus_size as u64);
-                    .update_corpus_size(*corpus_size as u64);
+                });
                 monitor.display(event.name(), ClientId(0));
                 Ok(BrokerEventResult::Handled)
             }
@@ -220,18 +220,18 @@ where
             } => {
                 // TODO: The monitor buffer should be added on client add.
                 monitor.client_stats_insert(ClientId(0));
-                let client = monitor.client_stats_mut_for(ClientId(0));
+                monitor.update_client_stats_for(ClientId(0), |client_stat| {
-
+                    client_stat.update_executions(*executions, *time);
-                client.update_executions(*executions, *time);
+                });
 
                 monitor.display(event.name(), ClientId(0));
                 Ok(BrokerEventResult::Handled)
             }
             Event::UpdateUserStats { name, value, .. } => {
                 monitor.client_stats_insert(ClientId(0));
-                monitor
+                monitor.update_client_stats_for(ClientId(0), |client_stat| {
-                    .client_stats_mut_for(ClientId(0))
+                    client_stat.update_user_stats(name.clone(), value.clone());
-                    .update_user_stats(name.clone(), value.clone());
+                });
                 monitor.aggregate(name);
                 monitor.display(event.name(), ClientId(0));
                 Ok(BrokerEventResult::Handled)
@@ -245,17 +245,18 @@ where
             } => {
                 // TODO: The monitor buffer should be added on client add.
                 monitor.client_stats_insert(ClientId(0));
-                let client = monitor.client_stats_mut_for(ClientId(0));
+                monitor.update_client_stats_for(ClientId(0), |client_stat| {
-                client.update_executions(*executions, *time);
+                    client_stat.update_executions(*executions, *time);
-                client.update_introspection_monitor((**introspection_monitor).clone());
+                    client_stat.update_introspection_monitor((**introspection_monitor).clone());
+                });
                 monitor.display(event.name(), ClientId(0));
                 Ok(BrokerEventResult::Handled)
             }
             Event::Objective { objective_size, .. } => {
                 monitor.client_stats_insert(ClientId(0));
-                monitor
+                monitor.update_client_stats_for(ClientId(0), |client_stat| {
-                    .client_stats_mut_for(ClientId(0))
+                    client_stat.update_objective_size(*objective_size as u64);
-                    .update_objective_size(*objective_size as u64);
+                });
                 monitor.display(event.name(), ClientId(0));
                 Ok(BrokerEventResult::Handled)
             }
@@ -542,7 +543,7 @@ where
 
                 // reload the state of the monitor to display the correct stats after restarts
                 monitor.set_start_time(start_time);
-                *monitor.client_stats_mut() = clients_stats;
+                monitor.update_all_client_stats(clients_stats);
 
                 (
                     Some(state),

libafl/src/events/tcp.rs

@@ -327,8 +327,9 @@ where
                     client_id
                 };
                 monitor.client_stats_insert(id);
-                let client = monitor.client_stats_mut_for(id);
+                monitor.update_client_stats_for(id, |client| {
                     client.update_corpus_size(*corpus_size as u64);
+                });
                 monitor.display(event.name(), id);
                 Ok(BrokerEventResult::Forward)
             }
@@ -339,8 +340,9 @@ where
             } => {
                 // TODO: The monitor buffer should be added on client add.
                 monitor.client_stats_insert(client_id);
-                let client = monitor.client_stats_mut_for(client_id);
+                monitor.update_client_stats_for(client_id, |client| {
                     client.update_executions(*executions, *time);
+                });
                 monitor.display(event.name(), client_id);
                 Ok(BrokerEventResult::Handled)
             }
@@ -350,8 +352,9 @@ where
                 phantom: _,
             } => {
                 monitor.client_stats_insert(client_id);
-                let client = monitor.client_stats_mut_for(client_id);
+                monitor.update_client_stats_for(client_id, |client| {
                     client.update_user_stats(name.clone(), value.clone());
+                });
                 monitor.aggregate(name);
                 monitor.display(event.name(), client_id);
                 Ok(BrokerEventResult::Handled)
@@ -367,13 +370,12 @@ where
 
                 // Get the client for the staterestorer ID
                 monitor.client_stats_insert(client_id);
-                let client = monitor.client_stats_mut_for(client_id);
+                monitor.update_client_stats_for(client_id, |client| {
-
                     // Update the normal monitor for this client
                     client.update_executions(*executions, *time);
-
                     // Update the performance monitor for this client
                     client.update_introspection_monitor((**introspection_monitor).clone());
+                });
 
                 // Display the monitor via `.display` only on core #1
                 monitor.display(event.name(), client_id);
@@ -383,8 +385,9 @@ where
             }
             Event::Objective { objective_size, .. } => {
                 monitor.client_stats_insert(client_id);
-                let client = monitor.client_stats_mut_for(client_id);
+                monitor.update_client_stats_for(client_id, |client| {
                     client.update_objective_size(*objective_size as u64);
+                });
                 monitor.display(event.name(), client_id);
                 Ok(BrokerEventResult::Handled)
             }

libafl/src/monitors/disk.rs

@@ -81,8 +81,13 @@ exec_sec = {}
             )
             .expect("Failed to write to the Toml file");
 
-            for (i, client) in self.client_stats_mut().iter_mut().enumerate() {
+            for i in 0..(self.client_stats().len()) {
-                let exec_sec = client.execs_per_sec(cur_time);
+                let client_id = ClientId(i as u32);
+                let exec_sec = self.update_client_stats_for(client_id, |client_stat| {
+                    client_stat.execs_per_sec(cur_time)
+                });
+
+                let client = self.client_stats_for(client_id);
 
                 write!(
                     &mut file,

libafl/src/monitors/mod.rs

@@ -502,6 +502,10 @@ impl ClientStats {
 /// The monitor trait keeps track of all the client's monitor, and offers methods to display them.
 pub trait Monitor {
     /// The client monitor (mutable)
+    ///
+    /// This method is for internal usage only, you shall never call this method directly.
+    /// If you want to update one client stats, use [`update_client_stats_for`][Self::update_client_stats_for]. If you
+    /// want to update all client stats together, use [`update_all_client_stats`][Self::update_all_client_stats].
     fn client_stats_mut(&mut self) -> &mut Vec<ClientStats>;
 
     /// The client monitor
@@ -571,7 +575,7 @@ pub trait Monitor {
                 ..ClientStats::default()
             });
         }
-        let new_stat = self.client_stats_mut_for(client_id);
+        self.update_client_stats_for(client_id, |new_stat| {
             if !new_stat.enabled {
                 let timestamp = current_time();
                 // I have never seen this man in my life
@@ -579,11 +583,25 @@ pub trait Monitor {
                 new_stat.last_window_time = timestamp;
                 new_stat.enabled = true;
             }
+        });
     }
 
-    /// Get mutable reference to client stats
+    /// Update sepecific client stats.
-    fn client_stats_mut_for(&mut self, client_id: ClientId) -> &mut ClientStats {
+    ///
-        &mut self.client_stats_mut()[client_id.0 as usize]
+    /// The update function is restricted as `Fn` instead of `FnMut` or `FnOnce` since in some
+    /// monitors, the `update` will be called multiple times, and is assumed as stateless.
+    fn update_client_stats_for<T, F: Fn(&mut ClientStats) -> T>(
+        &mut self,
+        client_id: ClientId,
+        update: F,
+    ) -> T {
+        let client_stat = &mut self.client_stats_mut()[client_id.0 as usize];
+        update(client_stat)
+    }
+
+    /// Update all client stats. This will clear all previous client stats, and fill in the new client stats.
+    fn update_all_client_stats(&mut self, new_client_stats: Vec<ClientStats>) {
+        *self.client_stats_mut() = new_client_stats;
     }
 
     /// Get immutable reference to client stats
@@ -791,7 +809,7 @@ where
 
         if self.print_user_monitor {
             self.client_stats_insert(sender_id);
-            let client = self.client_stats_mut_for(sender_id);
+            let client = self.client_stats_for(sender_id);
             for (key, val) in &client.user_monitor {
                 write!(fmt, ", {key}: {val}").unwrap();
             }
@@ -1303,18 +1321,17 @@ impl Default for ClientPerfMonitor {
     }
 }
 
-/// A combined monitor consisting of multiple [`Monitor`]s
+// The client stats of first and second monitor will always be maintained
+// to be consistent
+/// A combined monitor consisting of multiple [`Monitor`]s.
+///
+/// Note that the `CombinedMonitor` should never be the base monitor of other wrapper
+/// monitors.
 #[derive(Debug, Clone)]
 pub struct CombinedMonitor<A, B> {
     first: A,
     second: B,
     start_time: Duration,
-    /// Client stats. This will be maintained to be consistent with
-    /// client stats of first and second monitor.
-    ///
-    /// Currently, the client stats will be synced to first and second
-    /// before each display call.
-    client_stats: Vec<ClientStats>,
 }
 
 impl<A: Monitor, B: Monitor> CombinedMonitor<A, B> {
@@ -1323,22 +1340,24 @@ impl<A: Monitor, B: Monitor> CombinedMonitor<A, B> {
         let start_time = current_time();
         first.set_start_time(start_time);
         second.set_start_time(start_time);
+        first.update_all_client_stats(vec![]);
+        second.update_all_client_stats(vec![]);
         Self {
             first,
             second,
             start_time,
-            client_stats: vec![],
         }
     }
 }
 
 impl<A: Monitor, B: Monitor> Monitor for CombinedMonitor<A, B> {
+    /// Never call this method.
     fn client_stats_mut(&mut self) -> &mut Vec<ClientStats> {
-        &mut self.client_stats
+        panic!("client_stats_mut of CombinedMonitor should never be called")
     }
 
     fn client_stats(&self) -> &[ClientStats] {
-        &self.client_stats
+        self.first.client_stats()
     }
 
     fn start_time(&self) -> Duration {
@@ -1351,15 +1370,42 @@ impl<A: Monitor, B: Monitor> Monitor for CombinedMonitor<A, B> {
         self.second.set_start_time(time);
     }
 
+    fn client_stats_insert(&mut self, client_id: ClientId) {
+        self.first.client_stats_insert(client_id);
+        self.second.client_stats_insert(client_id);
+    }
+
+    #[inline]
+    fn execs_per_sec(&mut self) -> f64 {
+        let execs_per_sec = self.first.execs_per_sec();
+        let _ = self.second.execs_per_sec();
+        execs_per_sec
+    }
+
     fn display(&mut self, event_msg: &str, sender_id: ClientId) {
-        self.first.client_stats_mut().clone_from(&self.client_stats);
         self.first.display(event_msg, sender_id);
-        self.second
-            .client_stats_mut()
-            .clone_from(&self.client_stats);
         self.second.display(event_msg, sender_id);
     }
 
+    /// The `update` will be called multiple times, and the result of first
+    /// invocation will be returned. Since the client stats are guaranteed
+    /// to be consistent across first and second monitor, the result should be
+    /// the same.
+    fn update_client_stats_for<T, F: Fn(&mut ClientStats) -> T>(
+        &mut self,
+        client_id: ClientId,
+        update: F,
+    ) -> T {
+        let res = self.first.update_client_stats_for(client_id, &update);
+        let _ = self.second.update_client_stats_for(client_id, &update);
+        res
+    }
+
+    fn update_all_client_stats(&mut self, new_client_stats: Vec<ClientStats>) {
+        self.first.update_all_client_stats(new_client_stats.clone());
+        self.second.update_all_client_stats(new_client_stats);
+    }
+
     fn aggregate(&mut self, name: &str) {
         self.first.aggregate(name);
         self.second.aggregate(name);

libafl/src/monitors/multi.rs

@@ -88,9 +88,10 @@ where
         (self.print_fn)(&global_fmt);
 
         self.client_stats_insert(sender_id);
-        let client = self.client_stats_mut_for(sender_id);
         let cur_time = current_time();
-        let exec_sec = client.execs_per_sec_pretty(cur_time);
+        let exec_sec =
+            self.update_client_stats_for(sender_id, |client| client.execs_per_sec_pretty(cur_time));
+        let client = self.client_stats_for(sender_id);
 
         let pad = " ".repeat(head.len());
         let mut fmt = format!(

libafl/src/monitors/tui/mod.rs

@@ -454,8 +454,9 @@ impl Monitor for TuiMonitor {
         }
 
         self.client_stats_insert(sender_id);
-        let client = self.client_stats_mut_for(sender_id);
+        let exec_sec =
-        let exec_sec = client.execs_per_sec_pretty(cur_time);
+            self.update_client_stats_for(sender_id, |client| client.execs_per_sec_pretty(cur_time));
+        let client = self.client_stats_for(sender_id);
 
         let sender = format!("#{}", sender_id.0);
         let pad = if event_msg.len() + sender.len() < 13 {