feat: Listen for ctrl-c

bootandy · bootandy · commit b54a2158055e · 2024-03-25T21:33:40.000Z
before ctrl-c would not cancel a running job.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -40,6 +40,7 @@ config-file = "0.2"
 serde = { version = "1.0", features = ["derive"] }
 directories = "4"
 sysinfo = "0.27"
+ctrlc = "3.4"
 
 [target.'cfg(windows)'.dependencies]
 winapi-util = "0.1"
diff --git a/src/dir_walker.rs b/src/dir_walker.rs
@@ -130,6 +130,9 @@ fn ignore_file(entry: &DirEntry, walk_data: &WalkData) -> bool {
 fn walk(dir: PathBuf, walk_data: &WalkData, depth: usize) -> Option<Node> {
     let prog_data = &walk_data.progress_data;
     let errors = &walk_data.errors;
+    if errors.lock().unwrap().abort {
+        return None;
+    }
 
     let children = if dir.is_dir() {
         let read_dir = fs::read_dir(&dir);
diff --git a/src/main.rs b/src/main.rs
@@ -105,6 +105,17 @@ fn main() {
     let options = build_cli().get_matches();
     let config = get_config();
 
+    let errors = RuntimeErrors::default();
+    let error_listen_for_ctrlc = Arc::new(Mutex::new(errors));
+    let errors_for_rayon = error_listen_for_ctrlc.clone();
+    let errors_final = error_listen_for_ctrlc.clone();
+
+    ctrlc::set_handler(move || {
+        error_listen_for_ctrlc.lock().unwrap().abort = true;
+        println!("\nAborting");
+    })
+    .expect("Error setting Ctrl-C handler");
+
     let target_dirs = match options.get_many::<String>("params") {
         Some(values) => values.map(|v| v.as_str()).collect::<Vec<&str>>(),
         None => vec!["."],
@@ -197,7 +208,7 @@ fn main() {
         ignore_hidden,
         follow_links,
         progress_data: indicator.data.clone(),
-        errors: Arc::new(Mutex::new(RuntimeErrors::default())),
+        errors: errors_for_rayon,
     };
     let stack_size = config.get_custom_stack_size(&options);
     init_rayon(&stack_size);
@@ -222,6 +233,10 @@ fn main() {
     // Must have stopped indicator before we print to stderr
     indicator.stop();
 
+    if errors_final.lock().unwrap().abort {
+        return;
+    }
+
     let final_errors = walk_data.errors.lock().unwrap();
     let failed_permissions = final_errors.no_permissions;
     if !final_errors.file_not_found.is_empty() {
diff --git a/src/progress.rs b/src/progress.rs
@@ -73,6 +73,7 @@ pub struct RuntimeErrors {
     pub no_permissions: bool,
     pub file_not_found: HashSet<String>,
     pub unknown_error: HashSet<String>,
+    pub abort: bool,
 }
 
 /* -------------------------------------------------------------------------- */

Original file line number	Diff line number	Diff line change
`@@ -73,6 +73,7 @@ pub struct RuntimeErrors {`
`73`	`73`	`pub no_permissions: bool,`
`74`	`74`	`pub file_not_found: HashSet<String>,`
`75`	`75`	`pub unknown_error: HashSet<String>,`
	`76`	`+ pub abort: bool,`
`76`	`77`	`}`
`77`	`78`
`78`	`79`	`/* -------------------------------------------------------------------------- */`