Add more reliability

Cargo.lock

@@ -389,13 +389,12 @@ dependencies = [
 
 [[package]]
 name = "nhentairs"
-version = "0.5.4"
+version = "0.5.5"
 dependencies = [
  "reqwest",
  "serde",
  "serde_json",
  "tokio",
- "tokio-stream",
 ]
 
 [[package]]

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "nhentairs"
-version = "0.5.4"
+version = "0.5.5"
 authors = ["blank X <theblankx@protonmail.com>"]
 edition = "2018"
 
@@ -12,6 +12,5 @@ lto = true
 [dependencies]
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
-reqwest = { version = "0.11", features = ["stream"] }
-tokio = { version = "1.0", features = ["rt-multi-thread", "sync"] }
-tokio-stream = "0.1"
+reqwest = "0.11"
+tokio = { version = "1.0", features = ["rt-multi-thread", "sync", "time"] }

src/commands/download.rs

@@ -7,11 +7,13 @@ use std::path::Path;
 use std::process::exit;
 use tokio::sync::Mutex;
 use tokio::task::JoinHandle;
+use tokio::time::{sleep, Duration};
 use std::fs::{rename, create_dir, write};
 extern crate tokio;
 extern crate reqwest;
 
 const DOWNLOAD_WORKERS: usize = 5;
+const FAIL_DOWNLOAD_WAIT_TIME: u64 = 5000;
 
 pub async fn run(args: env::Args) {
     let sauces = utils::get_arg_sauces(args).unwrap();
@@ -90,7 +92,17 @@ pub async fn run(args: env::Args) {
                 eprintln!("[DW{}] Downloading {} to {}", worker_id, url, file_path);
                 let mut tmp_file_path = file_path.clone();
                 tmp_file_path.push_str(".tmp");
-                utils::download_file(cloned_client, &url, &tmp_file_path).await.unwrap();
+                loop {
+                    match utils::download_file(cloned_client.clone(), &url, &tmp_file_path).await {
+                        Ok(success) => {
+                            if success {
+                                break;
+                            }
+                        },
+                        Err(err) => eprintln!("[DW{}] Failed to download {} due to {}, sleeping for {}ms", worker_id, file_path, err, FAIL_DOWNLOAD_WAIT_TIME)
+                    };
+                    sleep(Duration::from_millis(FAIL_DOWNLOAD_WAIT_TIME)).await;
+                }
                 rename(&tmp_file_path, &file_path).unwrap();
                 eprintln!("[DW{}] {} downloaded", worker_id, file_path);
             }

src/utils.rs

@@ -3,7 +3,6 @@ use crate::structs;
 use std::env;
 use std::fs::File;
 use std::io::Write;
-use tokio_stream::StreamExt;
 extern crate serde_json;
 extern crate reqwest;
 
@@ -25,16 +24,19 @@ pub async fn get_search_info(client: reqwest::Client, search_query: &str) -> Res
     Ok(serde_json::from_str(&resp.text().await?).unwrap())
 }
 
-pub async fn download_file(client: reqwest::Client, url: &str, file_name: &str) -> Result<(), reqwest::Error> {
+pub async fn download_file(client: reqwest::Client, url: &str, file_name: &str) -> Result<bool, reqwest::Error> {
     let resp = client.get(url)
         .send()
         .await?;
-    let mut file = File::create(&file_name).unwrap();
-    let mut stream = resp.bytes_stream();
-    while let Some(item) = stream.next().await {
-        file.write(&item?).unwrap();
-    }
-    Ok(())
+    Ok(match resp.headers().get("Content-Type") {
+        Some(header) if header.to_str().unwrap_or_default().starts_with("image/") => {
+            let bytes = resp.bytes().await?;
+            let mut file = File::create(&file_name).unwrap();
+            file.write_all(&bytes).unwrap();
+            true
+        },
+        _ => false
+    })
 }
 
 pub fn get_arg_sauces(args: env::Args) -> Result<Vec<i32>, String> {