cleanup local job_id directory upon clean exit

src/agent/atexit/src/lib.rs

@@ -25,6 +25,13 @@ pub fn exit_process(code: i32) -> ! {
     ATEXIT.exit_process(code)
 }
 
+/// Runs the registered functions but does *not* terminate the process
+///
+/// This function is not called automatically (e.g. via `drop`).
+pub fn execute() {
+    ATEXIT.execute()
+}
+
 impl AtExit {
     fn new() -> Arc<Self> {
         let result = Arc::new(AtExit {
@@ -56,9 +63,13 @@ impl AtExit {
     }
 
     fn exit_process(&self, code: i32) -> ! {
+        self.execute();
+        std::process::exit(code);
+    }
+
+    fn execute(&self) {
         for function in self.functions.write().unwrap().iter_mut() {
             function();
         }
-        std::process::exit(code);
     }
 }

src/agent/onefuzz-agent/Cargo.toml

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

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

@@ -1,21 +1,20 @@
 use crate::tasks::config::CommonConfig;
 use crate::tasks::utils::parse_key_value;
-use anyhow::Result;
+use anyhow::{Context, 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 reqwest::Url;
+use std::task::Poll;
 use std::{
     collections::HashMap,
     path::{Path, PathBuf},
     time::Duration,
 };
 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 +127,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 +167,30 @@ pub fn get_synced_dirs(
     Ok(dirs?)
 }
 
+fn recursive_remove(path: &Path) -> Result<()> {
+    if path.metadata()?.is_dir() {
+        for child in std::fs::read_dir(&path)? {
+            let path = child?.path();
+            // allow children remove to fail, as tempdirs are cleaning up at this
+            // point.  the follow on remove_dir_all will catch errors.
+            recursive_remove(&path).ok();
+        }
+        std::fs::remove_dir_all(&path)
+            .with_context(|| format!("removing directory {}", path.display()))?;
+    } else {
+        std::fs::remove_file(&path).with_context(|| format!("removing {}", path.display()))?;
+    }
+    Ok(())
+}
+
+fn register_cleanup(job_id: Uuid) -> Result<()> {
+    let path = std::env::current_dir()?.join(job_id.to_string());
+    atexit::register(move || {
+        recursive_remove(&path).expect("cleanup failed");
+    });
+    Ok(())
+}
+
 pub fn get_synced_dir(
     name: &str,
     job_id: Uuid,
@@ -204,6 +233,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<()> {