Add OnDiskJsonAggregateMonitor

libafl/src/monitors/disk_aggregate.rs

@@ -0,0 +1,125 @@
+//! Monitors that log aggregated stats to disk.
+
+use alloc::vec::Vec;
+use core::{
+    fmt::{Debug, Formatter},
+    time::Duration,
+};
+use std::{fs::OpenOptions, io::Write, path::PathBuf};
+
+use libafl_bolts::{current_time, ClientId};
+use serde_json::json;
+
+use crate::monitors::{Aggregator, ClientStats, Monitor};
+
+/// A monitor that wraps another monitor and logs aggregated stats to a JSON file.
+#[derive(Clone)]
+pub struct OnDiskJsonAggregateMonitor<M> {
+    base: M,
+    aggregator: Aggregator,
+    json_path: PathBuf,
+    last_update: Duration,
+    update_interval: Duration,
+}
+
+impl<M> Debug for OnDiskJsonAggregateMonitor<M>
+where
+    M: Debug,
+{
+    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
+        f.debug_struct("OnDiskJsonAggregateMonitor")
+            .field("base", &self.base)
+            .field("last_update", &self.last_update)
+            .field("update_interval", &self.update_interval)
+            .field("json_path", &self.json_path)
+            .finish_non_exhaustive()
+    }
+}
+
+impl<M> Monitor for OnDiskJsonAggregateMonitor<M>
+where
+    M: Monitor,
+{
+    fn client_stats_mut(&mut self) -> &mut Vec<ClientStats> {
+        self.base.client_stats_mut()
+    }
+
+    fn client_stats(&self) -> &[ClientStats] {
+        self.base.client_stats()
+    }
+
+    fn set_start_time(&mut self, time: Duration) {
+        self.base.set_start_time(time);
+    }
+
+    fn start_time(&self) -> Duration {
+        self.base.start_time()
+    }
+
+    fn aggregate(&mut self, name: &str) {
+        self.aggregator.aggregate(name, self.base.client_stats());
+        self.base.aggregate(name);
+    }
+
+    fn display(&mut self, event_msg: &str, sender_id: ClientId) {
+        // First let the base monitor handle its display
+        self.base.display(event_msg, sender_id);
+
+        // Write JSON stats if update interval has elapsed
+        let cur_time = current_time();
+        if cur_time - self.last_update >= self.update_interval {
+            self.last_update = cur_time;
+
+            let file = OpenOptions::new()
+                .append(true)
+                .create(true)
+                .open(&self.json_path)
+                .expect("Failed to open JSON logging file");
+
+            let mut json_value = json!({
+                "run_time": (cur_time - self.start_time()).as_secs(),
+                "clients": self.client_stats_count(),
+                "corpus": self.corpus_size(),
+                "objectives": self.objective_size(),
+                "executions": self.total_execs(),
+                "exec_sec": self.execs_per_sec(),
+            });
+
+            // Add all aggregated values directly to the root
+            if let Some(obj) = json_value.as_object_mut() {
+                obj.extend(
+                    self.aggregator
+                        .aggregated
+                        .iter()
+                        .map(|(k, v)| (k.clone(), json!(v))),
+                );
+            }
+
+            writeln!(&file, "{json_value}").expect("Unable to write JSON to file");
+        }
+    }
+}
+
+impl<M> OnDiskJsonAggregateMonitor<M> {
+    /// Creates a new [`OnDiskJsonAggregateMonitor`]
+    pub fn new<P>(base: M, json_path: P) -> Self
+    where
+        P: Into<PathBuf>,
+    {
+        Self::with_interval(json_path, base, Duration::from_secs(10))
+    }
+
+    /// Creates a new [`OnDiskJsonAggregateMonitor`] with custom update interval
+    pub fn with_interval<P>(json_path: P, base: M, update_interval: Duration) -> Self
+    where
+        P: Into<PathBuf>,
+    {
+        Self {
+            base,
+            aggregator: Aggregator::new(),
+            json_path: json_path.into(),
+            last_update: current_time() - update_interval,
+            update_interval,
+        }
+    }
+}

libafl/src/monitors/mod.rs

@@ -3,25 +3,39 @@
 pub mod multi;
 pub use multi::MultiMonitor;
 
+#[cfg(feature = "std")]
+pub mod disk;
+
+#[cfg(feature = "std")]
+pub mod disk_aggregate;
+
 #[cfg(all(feature = "tui_monitor", feature = "std"))]
 pub mod tui;
 
 #[cfg(all(feature = "prometheus_monitor", feature = "std"))]
 pub mod prometheus;
-use alloc::string::ToString;
 
-#[cfg(all(feature = "prometheus_monitor", feature = "std"))]
-pub use prometheus::PrometheusMonitor;
-#[cfg(feature = "std")]
-pub mod disk;
-use alloc::{borrow::Cow, fmt::Debug, string::String, vec::Vec};
+use alloc::{
+    borrow::Cow,
+    fmt::Debug,
+    string::{String, ToString},
+    vec::Vec,
+};
 use core::{fmt, fmt::Write, time::Duration};
 
-#[cfg(feature = "std")]
-pub use disk::{OnDiskJsonMonitor, OnDiskTomlMonitor};
 use hashbrown::HashMap;
 use libafl_bolts::{current_time, format_duration_hms, ClientId};
+#[cfg(all(feature = "prometheus_monitor", feature = "std"))]
+pub use prometheus::PrometheusMonitor;
+#[cfg(feature = "std")]
 use serde::{Deserialize, Serialize};
+#[cfg(all(feature = "tui_monitor", feature = "std"))]
+pub use tui::TuiMonitor;
+#[cfg(feature = "std")]
+pub use {
+    disk::{OnDiskJsonMonitor, OnDiskTomlMonitor},
+    disk_aggregate::OnDiskJsonAggregateMonitor,
+};
 
 #[cfg(feature = "afl_exec_sec")]
 const CLIENT_STATS_TIME_WINDOW_SECS: u64 = 5; // 5 seconds
@@ -313,15 +327,15 @@ impl fmt::Display for UserStatsValue {
 /// Prettifies float values for human-readable output
 fn prettify_float(value: f64) -> String {
     let (value, suffix) = match value {
-        value if value >= 1_000_000.0 => (value / 1_000_000.0, "M"),
-        value if value >= 1_000.0 => (value / 1_000.0, "k"),
+        value if value >= 1000000.0 => (value / 1000000.0, "M"),
+        value if value >= 1000.0 => (value / 1000.0, "k"),
         value => (value, ""),
     };
     match value {
-        value if value >= 1_000_000.0 => {
+        value if value >= 1000000.0 => {
             format!("{value:.2}{suffix}")
         }
-        value if value >= 1_000.0 => {
+        value if value >= 1000.0 => {
             format!("{value:.1}{suffix}")
         }
         value if value >= 100.0 => {