Auto-cleanup

Author: denisbsu

README.md

@@ -16,11 +16,6 @@ If you have dedicated storage for worker, please provide docker access to the ex
 docker run --rm --network host --volume [path to storage folder]:/app/data -it subsquid/worker-benchmark:latest
 ```
 
-In this case you would need to clean up storage folder after run.
-
-
-No storage folder - no cleanup!
-
 Expected output contains progress indicators, detailed performance breakdown and overall report.
 
 Please note that as bulk of required bandwidth used to download initial data from infrastructure, we are measuring exactly this download speed. It may show numbers less than expected and usually not a dealbreaker as long as stated network speed is more or equal 1 Gigabit.
@@ -40,6 +35,7 @@ If you want to mess with internals, here are more details. If simply build and r
 - --config-url: URL to fetch benchmark description.
 - --config-file: filename to save benchmark description.
 - --force-update: re-download benchmark data even it is already downloaded.
+- --no-cleanup: disables auto-cleanup.
 
 # 1000 samples
 

src/main.rs

@@ -86,6 +86,9 @@ pub struct Cli {
 
     #[arg(long, env, default_value = "false")]
     pub force_update: bool,
+
+    #[arg(long, env, default_value = "false")]
+    pub no_cleanup: bool,
 }
 
 #[tokio::main]
@@ -122,7 +125,7 @@ async fn main() {
     let mut download_speed: f64 = 0.;
     if should_download {
         let _ = download_file(&config_url, &config_file).await;
-        let reader = fs::File::open(config_file).unwrap();
+        let reader = fs::File::open(&config_file).unwrap();
         bench_set = serde_yaml::from_reader(reader).unwrap();
 
         let url = bench_set.source.url.clone();
@@ -139,7 +142,7 @@ async fn main() {
         archive.unpack(".").unwrap();
         println!("Successfully decompressed and extracted");
     } else {
-        let reader = fs::File::open(config_file).unwrap();
+        let reader = fs::File::open(&config_file).unwrap();
         bench_set = serde_yaml::from_reader(reader).unwrap();
     }
 
@@ -188,6 +191,15 @@ async fn main() {
     }
     pb.finish_with_message("Benchmarking done.");
 
+    // cleanup
+    if !args.no_cleanup {
+        for idx in 0..bench_set.len() {
+            let path = bench_set.get_path(idx);
+            let _ = std::fs::remove_dir_all(path);
+        }
+        let _ = std::fs::remove_file(config_file);
+    }
+
     let mut time_matrix: Vec<Vec<f64>> = Default::default();
     time_matrix.resize(width, Default::default());
     let chunks = time_collector.chunks(width);