Skip to content

Commit a27da1b

Browse files
Evian-ZhangEvian-Zhang
authored
Change Monitor API for more flexibility (#2927)
* Change Monitor API for more flexibility * Make clippy happy * Fix broken doc link
1 parent 72986fc commit a27da1b

File tree

7 files changed

+138
-78
lines changed

7 files changed

+138
-78
lines changed

‎libafl/src/events/broker_hooks/mod.rs

Lines changed: 18 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -120,8 +120,9 @@ where
120120
};
121121

122122
monitor.client_stats_insert(id);
123-
let client = monitor.client_stats_mut_for(id);
124-
client.update_corpus_size(*corpus_size as u64);
123+
monitor.update_client_stats_for(id, |client_stat| {
124+
client_stat.update_corpus_size(*corpus_size as u64);
125+
});
125126
monitor.display(event.name(), id);
126127
Ok(BrokerEventResult::Forward)
127128
}
@@ -132,8 +133,9 @@ where
132133
} => {
133134
// TODO: The monitor buffer should be added on client add.
134135
monitor.client_stats_insert(client_id);
135-
let client = monitor.client_stats_mut_for(client_id);
136-
client.update_executions(*executions, *time);
136+
monitor.update_client_stats_for(client_id, |client_stat| {
137+
client_stat.update_executions(*executions, *time);
138+
});
137139
monitor.display(event.name(), client_id);
138140
Ok(BrokerEventResult::Handled)
139141
}
@@ -143,8 +145,9 @@ where
143145
phantom: _,
144146
} => {
145147
monitor.client_stats_insert(client_id);
146-
let client = monitor.client_stats_mut_for(client_id);
147-
client.update_user_stats(name.clone(), value.clone());
148+
monitor.update_client_stats_for(client_id, |client_stat| {
149+
client_stat.update_user_stats(name.clone(), value.clone());
150+
});
148151
monitor.aggregate(name);
149152
monitor.display(event.name(), client_id);
150153
Ok(BrokerEventResult::Handled)
@@ -160,13 +163,12 @@ where
160163

161164
// Get the client for the staterestorer ID
162165
monitor.client_stats_insert(client_id);
163-
let client = monitor.client_stats_mut_for(client_id);
164-
165-
// Update the normal monitor for this client
166-
client.update_executions(*executions, *time);
167-
168-
// Update the performance monitor for this client
169-
client.update_introspection_monitor((**introspection_monitor).clone());
166+
monitor.update_client_stats_for(client_id, |client_stat| {
167+
// Update the normal monitor for this client
168+
client_stat.update_executions(*executions, *time);
169+
// Update the performance monitor for this client
170+
client_stat.update_introspection_monitor((**introspection_monitor).clone());
171+
});
170172

171173
// Display the monitor via `.display` only on core #1
172174
monitor.display(event.name(), client_id);
@@ -176,8 +178,9 @@ where
176178
}
177179
Event::Objective { objective_size, .. } => {
178180
monitor.client_stats_insert(client_id);
179-
let client = monitor.client_stats_mut_for(client_id);
180-
client.update_objective_size(*objective_size as u64);
181+
monitor.update_client_stats_for(client_id, |client_stat| {
182+
client_stat.update_objective_size(*objective_size as u64);
183+
});
181184
monitor.display(event.name(), client_id);
182185
Ok(BrokerEventResult::Handled)
183186
}

‎libafl/src/events/simple.rs

Lines changed: 17 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -209,9 +209,9 @@ where
209209
match event {
210210
Event::NewTestcase { corpus_size, .. } => {
211211
monitor.client_stats_insert(ClientId(0));
212-
monitor
213-
.client_stats_mut_for(ClientId(0))
214-
.update_corpus_size(*corpus_size as u64);
212+
monitor.update_client_stats_for(ClientId(0), |client_stat| {
213+
client_stat.update_corpus_size(*corpus_size as u64);
214+
});
215215
monitor.display(event.name(), ClientId(0));
216216
Ok(BrokerEventResult::Handled)
217217
}
@@ -220,18 +220,18 @@ where
220220
} => {
221221
// TODO: The monitor buffer should be added on client add.
222222
monitor.client_stats_insert(ClientId(0));
223-
let client = monitor.client_stats_mut_for(ClientId(0));
224-
225-
client.update_executions(*executions, *time);
223+
monitor.update_client_stats_for(ClientId(0), |client_stat| {
224+
client_stat.update_executions(*executions, *time);
225+
});
226226

227227
monitor.display(event.name(), ClientId(0));
228228
Ok(BrokerEventResult::Handled)
229229
}
230230
Event::UpdateUserStats { name, value, .. } => {
231231
monitor.client_stats_insert(ClientId(0));
232-
monitor
233-
.client_stats_mut_for(ClientId(0))
234-
.update_user_stats(name.clone(), value.clone());
232+
monitor.update_client_stats_for(ClientId(0), |client_stat| {
233+
client_stat.update_user_stats(name.clone(), value.clone());
234+
});
235235
monitor.aggregate(name);
236236
monitor.display(event.name(), ClientId(0));
237237
Ok(BrokerEventResult::Handled)
@@ -245,17 +245,18 @@ where
245245
} => {
246246
// TODO: The monitor buffer should be added on client add.
247247
monitor.client_stats_insert(ClientId(0));
248-
let client = monitor.client_stats_mut_for(ClientId(0));
249-
client.update_executions(*executions, *time);
250-
client.update_introspection_monitor((**introspection_monitor).clone());
248+
monitor.update_client_stats_for(ClientId(0), |client_stat| {
249+
client_stat.update_executions(*executions, *time);
250+
client_stat.update_introspection_monitor((**introspection_monitor).clone());
251+
});
251252
monitor.display(event.name(), ClientId(0));
252253
Ok(BrokerEventResult::Handled)
253254
}
254255
Event::Objective { objective_size, .. } => {
255256
monitor.client_stats_insert(ClientId(0));
256-
monitor
257-
.client_stats_mut_for(ClientId(0))
258-
.update_objective_size(*objective_size as u64);
257+
monitor.update_client_stats_for(ClientId(0), |client_stat| {
258+
client_stat.update_objective_size(*objective_size as u64);
259+
});
259260
monitor.display(event.name(), ClientId(0));
260261
Ok(BrokerEventResult::Handled)
261262
}
@@ -542,7 +543,7 @@ where
542543

543544
// reload the state of the monitor to display the correct stats after restarts
544545
monitor.set_start_time(start_time);
545-
*monitor.client_stats_mut() = clients_stats;
546+
monitor.update_all_client_stats(clients_stats);
546547

547548
(
548549
Some(state),

‎libafl/src/events/tcp.rs

Lines changed: 18 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -327,8 +327,9 @@ where
327327
client_id
328328
};
329329
monitor.client_stats_insert(id);
330-
let client = monitor.client_stats_mut_for(id);
331-
client.update_corpus_size(*corpus_size as u64);
330+
monitor.update_client_stats_for(id, |client| {
331+
client.update_corpus_size(*corpus_size as u64);
332+
});
332333
monitor.display(event.name(), id);
333334
Ok(BrokerEventResult::Forward)
334335
}
@@ -339,8 +340,9 @@ where
339340
} => {
340341
// TODO: The monitor buffer should be added on client add.
341342
monitor.client_stats_insert(client_id);
342-
let client = monitor.client_stats_mut_for(client_id);
343-
client.update_executions(*executions, *time);
343+
monitor.update_client_stats_for(client_id, |client| {
344+
client.update_executions(*executions, *time);
345+
});
344346
monitor.display(event.name(), client_id);
345347
Ok(BrokerEventResult::Handled)
346348
}
@@ -350,8 +352,9 @@ where
350352
phantom: _,
351353
} => {
352354
monitor.client_stats_insert(client_id);
353-
let client = monitor.client_stats_mut_for(client_id);
354-
client.update_user_stats(name.clone(), value.clone());
355+
monitor.update_client_stats_for(client_id, |client| {
356+
client.update_user_stats(name.clone(), value.clone());
357+
});
355358
monitor.aggregate(name);
356359
monitor.display(event.name(), client_id);
357360
Ok(BrokerEventResult::Handled)
@@ -367,13 +370,12 @@ where
367370

368371
// Get the client for the staterestorer ID
369372
monitor.client_stats_insert(client_id);
370-
let client = monitor.client_stats_mut_for(client_id);
371-
372-
// Update the normal monitor for this client
373-
client.update_executions(*executions, *time);
374-
375-
// Update the performance monitor for this client
376-
client.update_introspection_monitor((**introspection_monitor).clone());
373+
monitor.update_client_stats_for(client_id, |client| {
374+
// Update the normal monitor for this client
375+
client.update_executions(*executions, *time);
376+
// Update the performance monitor for this client
377+
client.update_introspection_monitor((**introspection_monitor).clone());
378+
});
377379

378380
// Display the monitor via `.display` only on core #1
379381
monitor.display(event.name(), client_id);
@@ -383,8 +385,9 @@ where
383385
}
384386
Event::Objective { objective_size, .. } => {
385387
monitor.client_stats_insert(client_id);
386-
let client = monitor.client_stats_mut_for(client_id);
387-
client.update_objective_size(*objective_size as u64);
388+
monitor.update_client_stats_for(client_id, |client| {
389+
client.update_objective_size(*objective_size as u64);
390+
});
388391
monitor.display(event.name(), client_id);
389392
Ok(BrokerEventResult::Handled)
390393
}

‎libafl/src/monitors/disk.rs

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -81,8 +81,13 @@ exec_sec = {}
8181
)
8282
.expect("Failed to write to the Toml file");
8383

84-
for (i, client) in self.client_stats_mut().iter_mut().enumerate() {
85-
let exec_sec = client.execs_per_sec(cur_time);
84+
for i in 0..(self.client_stats().len()) {
85+
let client_id = ClientId(i as u32);
86+
let exec_sec = self.update_client_stats_for(client_id, |client_stat| {
87+
client_stat.execs_per_sec(cur_time)
88+
});
89+
90+
let client = self.client_stats_for(client_id);
8691

8792
write!(
8893
&mut file,

‎libafl/src/monitors/mod.rs

Lines changed: 72 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -502,6 +502,10 @@ impl ClientStats {
502502
/// The monitor trait keeps track of all the client's monitor, and offers methods to display them.
503503
pub trait Monitor {
504504
/// The client monitor (mutable)
505+
///
506+
/// This method is for internal usage only, you shall never call this method directly.
507+
/// If you want to update one client stats, use [`update_client_stats_for`][Self::update_client_stats_for]. If you
508+
/// want to update all client stats together, use [`update_all_client_stats`][Self::update_all_client_stats].
505509
fn client_stats_mut(&mut self) -> &mut Vec<ClientStats>;
506510

507511
/// The client monitor
@@ -571,19 +575,33 @@ pub trait Monitor {
571575
..ClientStats::default()
572576
});
573577
}
574-
let new_stat = self.client_stats_mut_for(client_id);
575-
if !new_stat.enabled {
576-
let timestamp = current_time();
577-
// I have never seen this man in my life
578-
new_stat.start_time = timestamp;
579-
new_stat.last_window_time = timestamp;
580-
new_stat.enabled = true;
581-
}
578+
self.update_client_stats_for(client_id, |new_stat| {
579+
if !new_stat.enabled {
580+
let timestamp = current_time();
581+
// I have never seen this man in my life
582+
new_stat.start_time = timestamp;
583+
new_stat.last_window_time = timestamp;
584+
new_stat.enabled = true;
585+
}
586+
});
587+
}
588+
589+
/// Update sepecific client stats.
590+
///
591+
/// The update function is restricted as `Fn` instead of `FnMut` or `FnOnce` since in some
592+
/// monitors, the `update` will be called multiple times, and is assumed as stateless.
593+
fn update_client_stats_for<T, F: Fn(&mut ClientStats) -> T>(
594+
&mut self,
595+
client_id: ClientId,
596+
update: F,
597+
) -> T {
598+
let client_stat = &mut self.client_stats_mut()[client_id.0 as usize];
599+
update(client_stat)
582600
}
583601

584-
/// Get mutable reference to client stats
585-
fn client_stats_mut_for(&mut self, client_id: ClientId) -> &mut ClientStats {
586-
&mut self.client_stats_mut()[client_id.0 as usize]
602+
/// Update all client stats. This will clear all previous client stats, and fill in the new client stats.
603+
fn update_all_client_stats(&mut self, new_client_stats: Vec<ClientStats>) {
604+
*self.client_stats_mut() = new_client_stats;
587605
}
588606

589607
/// Get immutable reference to client stats
@@ -791,7 +809,7 @@ where
791809

792810
if self.print_user_monitor {
793811
self.client_stats_insert(sender_id);
794-
let client = self.client_stats_mut_for(sender_id);
812+
let client = self.client_stats_for(sender_id);
795813
for (key, val) in &client.user_monitor {
796814
write!(fmt, ", {key}: {val}").unwrap();
797815
}
@@ -1303,18 +1321,17 @@ impl Default for ClientPerfMonitor {
13031321
}
13041322
}
13051323

1306-
/// A combined monitor consisting of multiple [`Monitor`]s
1324+
// The client stats of first and second monitor will always be maintained
1325+
// to be consistent
1326+
/// A combined monitor consisting of multiple [`Monitor`]s.
1327+
///
1328+
/// Note that the `CombinedMonitor` should never be the base monitor of other wrapper
1329+
/// monitors.
13071330
#[derive(Debug, Clone)]
13081331
pub struct CombinedMonitor<A, B> {
13091332
first: A,
13101333
second: B,
13111334
start_time: Duration,
1312-
/// Client stats. This will be maintained to be consistent with
1313-
/// client stats of first and second monitor.
1314-
///
1315-
/// Currently, the client stats will be synced to first and second
1316-
/// before each display call.
1317-
client_stats: Vec<ClientStats>,
13181335
}
13191336

13201337
impl<A: Monitor, B: Monitor> CombinedMonitor<A, B> {
@@ -1323,22 +1340,24 @@ impl<A: Monitor, B: Monitor> CombinedMonitor<A, B> {
13231340
let start_time = current_time();
13241341
first.set_start_time(start_time);
13251342
second.set_start_time(start_time);
1343+
first.update_all_client_stats(vec![]);
1344+
second.update_all_client_stats(vec![]);
13261345
Self {
13271346
first,
13281347
second,
13291348
start_time,
1330-
client_stats: vec![],
13311349
}
13321350
}
13331351
}
13341352

13351353
impl<A: Monitor, B: Monitor> Monitor for CombinedMonitor<A, B> {
1354+
/// Never call this method.
13361355
fn client_stats_mut(&mut self) -> &mut Vec<ClientStats> {
1337-
&mut self.client_stats
1356+
panic!("client_stats_mut of CombinedMonitor should never be called")
13381357
}
13391358

13401359
fn client_stats(&self) -> &[ClientStats] {
1341-
&self.client_stats
1360+
self.first.client_stats()
13421361
}
13431362

13441363
fn start_time(&self) -> Duration {
@@ -1351,15 +1370,42 @@ impl<A: Monitor, B: Monitor> Monitor for CombinedMonitor<A, B> {
13511370
self.second.set_start_time(time);
13521371
}
13531372

1373+
fn client_stats_insert(&mut self, client_id: ClientId) {
1374+
self.first.client_stats_insert(client_id);
1375+
self.second.client_stats_insert(client_id);
1376+
}
1377+
1378+
#[inline]
1379+
fn execs_per_sec(&mut self) -> f64 {
1380+
let execs_per_sec = self.first.execs_per_sec();
1381+
let _ = self.second.execs_per_sec();
1382+
execs_per_sec
1383+
}
1384+
13541385
fn display(&mut self, event_msg: &str, sender_id: ClientId) {
1355-
self.first.client_stats_mut().clone_from(&self.client_stats);
13561386
self.first.display(event_msg, sender_id);
1357-
self.second
1358-
.client_stats_mut()
1359-
.clone_from(&self.client_stats);
13601387
self.second.display(event_msg, sender_id);
13611388
}
13621389

1390+
/// The `update` will be called multiple times, and the result of first
1391+
/// invocation will be returned. Since the client stats are guaranteed
1392+
/// to be consistent across first and second monitor, the result should be
1393+
/// the same.
1394+
fn update_client_stats_for<T, F: Fn(&mut ClientStats) -> T>(
1395+
&mut self,
1396+
client_id: ClientId,
1397+
update: F,
1398+
) -> T {
1399+
let res = self.first.update_client_stats_for(client_id, &update);
1400+
let _ = self.second.update_client_stats_for(client_id, &update);
1401+
res
1402+
}
1403+
1404+
fn update_all_client_stats(&mut self, new_client_stats: Vec<ClientStats>) {
1405+
self.first.update_all_client_stats(new_client_stats.clone());
1406+
self.second.update_all_client_stats(new_client_stats);
1407+
}
1408+
13631409
fn aggregate(&mut self, name: &str) {
13641410
self.first.aggregate(name);
13651411
self.second.aggregate(name);

0 commit comments

Comments
 (0)