From 61ac4ea7be0d1bbed8d4670705fda1c270b883a2 Mon Sep 17 00:00:00 2001
From: clesmian <18008727+clesmian@users.noreply.github.com>
Date: Tue, 30 Apr 2024 15:05:33 +0200
Subject: [PATCH] Improve `OnDiskTOMLMonitor` (#2125)

* Allow for more frequent updates of TOML monitor

* Don't skip first client

* Reduce code duplication

* Immediately write first TOML file

* Rust fmt

* Use same client numbering as other monitors

* Fmt
---
 libafl/src/monitors/disk.rs | 23 +++++++++++++++--------
 1 file changed, 15 insertions(+), 8 deletions(-)

diff --git a/libafl/src/monitors/disk.rs b/libafl/src/monitors/disk.rs
index f977eb1da5..6df195bc63 100644
--- a/libafl/src/monitors/disk.rs
+++ b/libafl/src/monitors/disk.rs
@@ -22,6 +22,7 @@ where
     base: M,
     filename: PathBuf,
     last_update: Duration,
+    update_interval: Duration,
 }
 
 impl<M> Monitor for OnDiskTOMLMonitor<M>
@@ -55,7 +56,7 @@ where
     fn display(&mut self, event_msg: &str, sender_id: ClientId) {
         let cur_time = current_time();
 
-        if (cur_time - self.last_update).as_secs() >= 60 {
+        if cur_time - self.last_update >= self.update_interval {
             self.last_update = cur_time;
 
             let mut file = File::create(&self.filename).expect("Failed to open the TOML file");
@@ -80,7 +81,7 @@ exec_sec = {}
             )
             .expect("Failed to write to the TOML file");
 
-            for (i, client) in self.client_stats_mut().iter_mut().skip(1).enumerate() {
+            for (i, client) in self.client_stats_mut().iter_mut().enumerate() {
                 let exec_sec = client.execs_per_sec(cur_time);
 
                 write!(
@@ -92,11 +93,7 @@ objectives = {}
 executions = {}
 exec_sec = {}
 ",
-                    i + 1,
-                    client.corpus_size,
-                    client.objective_size,
-                    client.executions,
-                    exec_sec
+                    i, client.corpus_size, client.objective_size, client.executions, exec_sec
                 )
                 .expect("Failed to write to the TOML file");
 
@@ -125,13 +122,23 @@ where
     /// Create new [`OnDiskTOMLMonitor`]
     #[must_use]
     pub fn new<P>(filename: P, base: M) -> Self
+    where
+        P: Into<PathBuf>,
+    {
+        Self::with_update_interval(filename, base, Duration::from_secs(60))
+    }
+
+    /// Create new [`OnDiskTOMLMonitor`] with custom update interval
+    #[must_use]
+    pub fn with_update_interval<P>(filename: P, base: M, update_interval: Duration) -> Self
     where
         P: Into<PathBuf>,
     {
         Self {
             base,
             filename: filename.into(),
-            last_update: current_time(),
+            last_update: current_time() - update_interval,
+            update_interval,
         }
     }
 }