cleanup local job_id directory upon clean exit (#738)

src/agent/Cargo.lock

@@ -1702,6 +1702,7 @@ dependencies = [
  "anyhow",
  "appinsights",
  "async-trait",
+ "atexit",
  "backoff",
  "clap",
  "env_logger",
@@ -1713,6 +1714,7 @@ dependencies = [
  "onefuzz",
  "onefuzz-telemetry",
  "path-absolutize",
+ "remove_dir_all 0.7.0",
  "reqwest",
  "reqwest-retry",
  "serde",
@@ -2236,6 +2238,19 @@ dependencies = [
  "winapi 0.3.9",
 ]
 
+[[package]]
+name = "remove_dir_all"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "882f368737489ea543bc5c340e6f3d34a28c39980bd9a979e47322b26f60ac40"
+dependencies = [
+ "libc",
+ "log",
+ "num_cpus",
+ "rayon",
+ "winapi 0.3.9",
+]
+
 [[package]]
 name = "reqwest"
 version = "0.10.10"
@@ -2723,7 +2738,7 @@ dependencies = [
  "libc",
  "rand 0.8.3",
  "redox_syscall",
- "remove_dir_all",
+ "remove_dir_all 0.5.3",
  "winapi 0.3.9",
 ]
 

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,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<()> {