From 5af43cfc7bb30b3d7f1b12f87f67b0aead7e9aa9 Mon Sep 17 00:00:00 2001
From: Colin Wahl <cwahl@awakesecurity.com>
Date: Mon, 7 Jan 2019 23:12:30 -0800
Subject: [PATCH] add cleanup command which is executed at the end of each
 benchmark

---
 src/hyperfine/app.rs       | 11 ++++++++
 src/hyperfine/benchmark.rs | 54 ++++++++++++++++++++++++++++++++++----
 src/hyperfine/types.rs     |  4 +++
 src/main.rs                |  2 ++
 4 files changed, 66 insertions(+), 5 deletions(-)

diff --git a/src/hyperfine/app.rs b/src/hyperfine/app.rs
index 62a22e8..f8ee004 100644
--- a/src/hyperfine/app.rs
+++ b/src/hyperfine/app.rs
@@ -82,6 +82,17 @@ fn build_app() -> App<'static, 'static> {
                      clearing disk caches, for example.",
                 ),
         )
+        .arg(
+            Arg::with_name("cleanup")
+                .long("cleanup")
+                .short("c")
+                .takes_value(true)
+                .value_name("CMD")
+                .help(
+                    "Execute CMD after each benchmark run. This is useful for \
+                     cleaning up resources allocated during prepare.",
+                ),
+        )
         .arg(
             Arg::with_name("parameter-scan")
                 .long("parameter-scan")
diff --git a/src/hyperfine/benchmark.rs b/src/hyperfine/benchmark.rs
index 54f4ae7..5291e35 100644
--- a/src/hyperfine/benchmark.rs
+++ b/src/hyperfine/benchmark.rs
@@ -144,19 +144,18 @@ pub fn mean_shell_spawning_time(
     })
 }
 
-/// Run the command specified by `--prepare`.
-fn run_preparation_command(
+fn run_intermediate_command(
     shell: &str,
     command: &Option<Command>,
     show_output: bool,
+    error_output: &'static str,
 ) -> io::Result<TimingResult> {
     if let &Some(ref cmd) = command {
         let res = time_shell_command(shell, cmd, show_output, CmdFailureAction::RaiseError, None);
         if res.is_err() {
             return Err(io::Error::new(
                 io::ErrorKind::Other,
-                "The preparation command terminated with a non-zero exit code. \
-                 Append ' || true' to the command if you are sure that this can be ignored.",
+                error_output
             ));
         }
         return res.map(|r| r.0);
@@ -164,6 +163,41 @@ fn run_preparation_command(
     Ok(TimingResult {
         ..Default::default()
     })
+
+}
+
+/// Run the command specified by `--prepare`.
+fn run_preparation_command(
+    shell: &str,
+    command: &Option<Command>,
+    show_output: bool,
+) -> io::Result<TimingResult> {
+    let error_output = "The preparation command terminated with a non-zero exit code. \
+                       Append ' || true' to the command if you are sure that this can be ignored.";
+
+    run_intermediate_command(
+        shell,
+        command,
+        show_output,
+        error_output
+    )
+}
+
+/// Run the command specified by `--cleanup`.
+fn run_cleanup_command(
+    shell: &str,
+    command: &Option<Command>,
+    show_output: bool,
+) -> io::Result<TimingResult> {
+    let error_output = "The cleanup command terminated with a non-zero exit code. \
+                       Append ' || true' to the command if you are sure that this can be ignored.";
+
+    run_intermediate_command(
+        shell,
+        command,
+        show_output,
+        error_output
+    )
 }
 
 /// Run the benchmark for a single shell command
@@ -213,7 +247,7 @@ pub fn run_benchmark(
         &options.output_style,
     );
 
-    // Run init / cleanup command
+    // Run init command
     let prepare_cmd = options
         .preparation_command
         .as_ref()
@@ -359,6 +393,16 @@ pub fn run_benchmark(
 
     println!(" ");
 
+    // Run cleanup command
+    let cleanup_cmd = options
+        .cleanup_command
+        .as_ref()
+        .map(|cleanup_command| match cmd.get_parameter() {
+            Some((param, value)) => Command::new_parametrized(cleanup_command, param, value),
+            None => Command::new(cleanup_command),
+        });
+    run_cleanup_command(&options.shell, &cleanup_cmd, options.show_output)?;
+
     Ok(BenchmarkResult::new(
         cmd.get_shell_command(),
         t_mean,
diff --git a/src/hyperfine/types.rs b/src/hyperfine/types.rs
index 1568a5e..d07fd91 100644
--- a/src/hyperfine/types.rs
+++ b/src/hyperfine/types.rs
@@ -113,6 +113,9 @@ pub struct HyperfineOptions {
     /// Command to run before each timing run
     pub preparation_command: Option<String>,
 
+    /// Command to run after each benchmark
+    pub cleanup_command: Option<String>,
+
     /// What color mode to use for output
     pub output_style: OutputStyleOption,
 
@@ -134,6 +137,7 @@ impl Default for HyperfineOptions {
             min_time_sec: 3.0,
             failure_action: CmdFailureAction::RaiseError,
             preparation_command: None,
+            cleanup_command: None,
             output_style: OutputStyleOption::Full,
             shell: DEFAULT_SHELL.to_string(),
             show_output: false,
diff --git a/src/main.rs b/src/main.rs
index 0d13f64..3dd6bfa 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -164,6 +164,8 @@ fn build_hyperfine_options(matches: &ArgMatches) -> Result<HyperfineOptions, Opt
 
     options.preparation_command = matches.value_of("prepare").map(String::from);
 
+    options.cleanup_command = matches.value_of("cleanup").map(String::from);
+
     options.show_output = matches.is_present("show-output");
 
     options.output_style = match matches.value_of("style") {