cleanup local job_id directory upon clean exit (#738)

src/agent/onefuzz-agent/Cargo.toml

@@ -35,6 +35,8 @@ reqwest-retry = { path = "../reqwest-retry" }
 onefuzz-telemetry = { path = "../onefuzz-telemetry" }
 stacktrace-parser = { path = "../stacktrace-parser" }
 path-absolutize = "3.0.6"
+atexit = { path = "../atexit" }
+remove_dir_all = "0.7"
 
 [dev-dependencies]
 tempfile = "3.2"

src/agent/onefuzz-agent/src/local/common.rs

@@ -1,10 +1,14 @@
 use crate::tasks::config::CommonConfig;
 use crate::tasks::utils::parse_key_value;
 use anyhow::Result;
+use backoff::{future::retry, Error as BackoffError, ExponentialBackoff};
 use clap::{App, Arg, ArgMatches};
 use onefuzz::jitter::delay_with_jitter;
 use onefuzz::{blob::BlobContainerUrl, monitor::DirectoryMonitor, syncdir::SyncedDir};
+use path_absolutize::Absolutize;
+use remove_dir_all::remove_dir_all;
 use reqwest::Url;
+use std::task::Poll;
 use std::{
     collections::HashMap,
     path::{Path, PathBuf},
@@ -12,10 +16,6 @@ use std::{
 };
 use uuid::Uuid;
 
-use backoff::{future::retry, Error as BackoffError, ExponentialBackoff};
-use path_absolutize::Absolutize;
-use std::task::Poll;
-
 pub const SETUP_DIR: &str = "setup_dir";
 pub const INPUTS_DIR: &str = "inputs_dir";
 pub const CRASHES_DIR: &str = "crashes_dir";
@@ -128,6 +128,12 @@ pub fn add_common_config(app: App<'static, 'static>) -> App<'static, 'static> {
             .takes_value(true)
             .required(false),
     )
+    .arg(
+        Arg::with_name("keep_job_dir")
+            .long("keep_job_dir")
+            .required(false)
+            .help("keep the local directory created for running the task"),
+    )
 }
 
 fn get_uuid(name: &str, args: &ArgMatches<'_>) -> Result<Uuid> {
@@ -162,6 +168,14 @@ pub fn get_synced_dirs(
     Ok(dirs?)
 }
 
+fn register_cleanup(job_id: Uuid) -> Result<()> {
+    let path = std::env::current_dir()?.join(job_id.to_string());
+    atexit::register(move || {
+        remove_dir_all(&path).expect("cleanup failed");
+    });
+    Ok(())
+}
+
 pub fn get_synced_dir(
     name: &str,
     job_id: Uuid,
@@ -204,6 +218,10 @@ pub fn build_common_config(args: &ArgMatches<'_>, generate_task_id: bool) -> Res
         PathBuf::default()
     };
 
+    if !args.is_present("keep_job_dir") {
+        register_cleanup(job_id)?;
+    }
+
     let config = CommonConfig {
         job_id,
         task_id,

src/agent/onefuzz-agent/src/local/libfuzzer.rs

@@ -20,7 +20,6 @@ use crate::{
 };
 use anyhow::Result;
 use clap::{App, SubCommand};
-
 use onefuzz::utils::try_wait_all_join_handles;
 use std::collections::HashSet;
 use tokio::task::spawn;

src/agent/onefuzz-agent/src/local/libfuzzer_coverage.rs

@@ -57,7 +57,6 @@ pub fn build_coverage_config(
 pub async fn run(args: &clap::ArgMatches<'_>) -> Result<()> {
     let common = build_common_config(args, true)?;
     let config = build_coverage_config(args, false, None, common)?;
-
     let mut task = CoverageTask::new(config);
     task.managed_run().await
 }

src/agent/onefuzz-agent/src/main.rs

@@ -40,7 +40,9 @@ fn main() -> Result<()> {
     let matches = app.get_matches();
 
     let mut rt = tokio::runtime::Runtime::new()?;
-    rt.block_on(run(matches))
+    let result = rt.block_on(run(matches));
+    atexit::execute();
+    result
 }
 
 async fn run(args: ArgMatches<'_>) -> Result<()> {