From d40f18f3772905d9c470bc3a207d99bd350561a1 Mon Sep 17 00:00:00 2001
From: "Dongjia \"toka\" Zhang" <tokazerkje@outlook.com>
Date: Fri, 21 Feb 2025 12:23:00 +0100
Subject: [PATCH] Fix finish_stage call (#3014)

* fixer

* clp

---------

Co-authored-by: toka <toka@tokas-MacBook-Air.local>
---
 libafl/src/stages/mod.rs        | 8 ++++++--
 libafl/src/stages/mutational.rs | 7 +------
 libafl/src/stages/sync.rs       | 5 -----
 libafl/src/stages/tmin.rs       | 3 ---
 libafl/src/stages/tuneable.rs   | 7 +------
 5 files changed, 8 insertions(+), 22 deletions(-)

diff --git a/libafl/src/stages/mod.rs b/libafl/src/stages/mod.rs
index 16a3c578dc..538e54c6f8 100644
--- a/libafl/src/stages/mod.rs
+++ b/libafl/src/stages/mod.rs
@@ -49,7 +49,7 @@ pub use verify_timeouts::{TimeoutsToVerify, VerifyTimeoutsStage};
 use crate::{
     corpus::{CorpusId, HasCurrentCorpusId},
     events::SendExiting,
-    state::{HasCurrentStageId, HasExecutions, Stoppable},
+    state::{HasCurrentStageId, HasExecutions, MaybeHasClientPerfMonitor, Stoppable},
     Error, HasNamedMetadata,
 };
 
@@ -151,7 +151,7 @@ impl<Head, Tail, E, EM, S, Z> StagesTuple<E, EM, S, Z> for (Head, Tail)
 where
     Head: Stage<E, EM, S, Z> + Restartable<S>,
     Tail: StagesTuple<E, EM, S, Z> + HasConstLen,
-    S: HasCurrentStageId + Stoppable,
+    S: HasCurrentStageId + Stoppable + MaybeHasClientPerfMonitor,
     EM: SendExiting,
 {
     /// Performs all stages in the tuple,
@@ -192,6 +192,10 @@ where
             }
         }
 
+        // Mark the elapsed time for the scheduler
+        #[cfg(feature = "introspection")]
+        state.introspection_stats_mut().finish_stage();
+
         if state.stop_requested() {
             state.discard_stop_request();
             manager.on_shutdown()?;
diff --git a/libafl/src/stages/mutational.rs b/libafl/src/stages/mutational.rs
index 5f44a916b4..b6900e98b3 100644
--- a/libafl/src/stages/mutational.rs
+++ b/libafl/src/stages/mutational.rs
@@ -163,12 +163,7 @@ where
         state: &mut S,
         manager: &mut EM,
     ) -> Result<(), Error> {
-        let ret = self.perform_mutational(fuzzer, executor, state, manager);
-
-        #[cfg(feature = "introspection")]
-        state.introspection_stats_mut().finish_stage();
-
-        ret
+        self.perform_mutational(fuzzer, executor, state, manager)
     }
 }
 
diff --git a/libafl/src/stages/sync.rs b/libafl/src/stages/sync.rs
index af6eaf1c50..b9a88b78a3 100644
--- a/libafl/src/stages/sync.rs
+++ b/libafl/src/stages/sync.rs
@@ -139,9 +139,6 @@ where
             fuzzer.evaluate_input(state, executor, manager, &input)?;
         }
 
-        #[cfg(feature = "introspection")]
-        state.introspection_stats_mut().finish_stage();
-
         Ok(())
     }
 }
@@ -307,8 +304,6 @@ where
         }
 
         self.client.process(fuzzer, state, executor, manager)?;
-        #[cfg(feature = "introspection")]
-        state.introspection_stats_mut().finish_stage();
         Ok(())
     }
 }
diff --git a/libafl/src/stages/tmin.rs b/libafl/src/stages/tmin.rs
index 20667f950a..01806b1391 100644
--- a/libafl/src/stages/tmin.rs
+++ b/libafl/src/stages/tmin.rs
@@ -91,9 +91,6 @@ where
     ) -> Result<(), Error> {
         self.perform_minification(fuzzer, executor, state, manager)?;
 
-        #[cfg(feature = "introspection")]
-        state.introspection_stats_mut().finish_stage();
-
         Ok(())
     }
 }
diff --git a/libafl/src/stages/tuneable.rs b/libafl/src/stages/tuneable.rs
index 92764d8352..4a5e86143a 100644
--- a/libafl/src/stages/tuneable.rs
+++ b/libafl/src/stages/tuneable.rs
@@ -210,12 +210,7 @@ where
         state: &mut S,
         manager: &mut EM,
     ) -> Result<(), Error> {
-        let ret = self.perform_mutational(fuzzer, executor, state, manager);
-
-        #[cfg(feature = "introspection")]
-        state.introspection_stats_mut().finish_stage();
-
-        ret
+        self.perform_mutational(fuzzer, executor, state, manager)
     }
 }