feat: capture SIGINT to avoid corrupting downloads

src/main.rs

@@ -10,6 +10,8 @@ use tokio::time::{sleep, Duration};
 
 use std::io::Write;
 use std::process::ExitCode;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
 
 const BAR_LENGTH: u64 = 8;
 
@@ -20,12 +22,23 @@ async fn main() -> ExitCode {
     let uri_tags = &args.tags.join("+");
     let _ = std::fs::create_dir(uri_tags);
 
+    let running = Arc::new(AtomicBool::new(true));
+    let running_t = running.clone();
+    ctrlc::set_handler(move || {
+        running_t.store(false, Ordering::SeqCst);
+    })
+    .unwrap();
+
     let client = Client::builder()
         .user_agent(&args.user_agent)
         .build()
         .unwrap();
 
     for page in args.page - 1.. {
+        if !running.load(Ordering::SeqCst) {
+            return ExitCode::FAILURE;
+        }
+
         println!("now scraping page {} (https://rule34.xxx/index.php?page=post&s=list&tags={uri_tags}&pid={})", page + 1, page * 42);
 
         let post_html = async |client: &Client| {
@@ -51,6 +64,10 @@ async fn main() -> ExitCode {
                 println!("no urls found, retrying in 5 seconds...");
                 sleep(Duration::from_millis(5000)).await;
 
+                if !running.load(Ordering::SeqCst) {
+                    return ExitCode::FAILURE;
+                }
+
                 urls = post_html(&client).await;
                 if !urls.is_empty() {
                     println!("urls found! continuing...");
@@ -71,12 +88,17 @@ async fn main() -> ExitCode {
             .map(|(i, url)| {
                 let i = i + 1;
                 let client = &client;
+                let running_t = running.clone();
                 let this_bar = indicatif::ProgressBar::new(BAR_LENGTH);
                 this_bar.set_style(indicatif::ProgressStyle::with_template("[{bar}] {msg}").unwrap().progress_chars("=> "));
                 let this_bar = multi_prog.insert(i, this_bar);
                 async move {
                     // "thread"
                     loop {
+                        if !running_t.load(Ordering::SeqCst) {
+                            return;
+                        }
+
                         this_bar.set_message(format!("\x1b[37m[{i: >4}/{urls_amount}] \x1b[36mscraping {url}\x1b[0m"));
                         let resp = client.get(&url).send().await.unwrap();
                         if let Ok(img_url) = extract_img_url(&resp.text().await.unwrap()) {
@@ -85,7 +107,7 @@ async fn main() -> ExitCode {
                                     "\x1b[37m[{i: >4}/{urls_amount}] \x1b[1;31mimage url not found\x1b[0m"
                                 ));
                             } else {
-                                download_file(&img_url, this_bar, i, urls_amount, uri_tags).await;
+                                download_file(running_t, &img_url, this_bar, i, urls_amount, uri_tags).await;
                             }
                             break;
                         }
@@ -129,6 +151,7 @@ fn extract_img_url(html: &str) -> Result<String, &'static str> {
 }
 
 async fn download_file(
+    running: Arc<AtomicBool>,
     img_url: &str,
     this_bar: ProgressBar,
     i: usize,
@@ -151,7 +174,7 @@ async fn download_file(
         ));
         return;
     } else {
-        std::fs::File::create(file_path).unwrap()
+        std::fs::File::create(&file_path).unwrap()
     };
 
     let mut res = Client::new()
@@ -163,6 +186,14 @@ async fn download_file(
     let file_length = res.content_length().unwrap();
     let mut written = 0;
     while let Some(img_chunk) = res.chunk().await.unwrap() {
+        if !running.load(Ordering::SeqCst) {
+            this_bar.abandon_with_message(format!(
+                "\x1b[37m[{i: >4}/{urls_amount}] \x1b[33mcancelling {img_url}\x1b[0m"
+            ));
+            drop(file);
+            std::fs::remove_file(&file_path).unwrap();
+            return;
+        }
         file.write_all(&img_chunk).unwrap();
         written += img_chunk.len();
         this_bar.set_position((written as f64 / file_length as f64 * BAR_LENGTH as f64) as u64);