retry failed chunk downloads; Continues #1

src/downloader.rs

@@ -1,5 +1,6 @@
 use crate::state::AppState;
-use std::sync::Arc;
+use std::{sync::Arc, time::Duration};
+use tokio::time::timeout;
 use tokio_stream::StreamExt;
 
 pub async fn spawn_download_task(url: &str, state: Arc<AppState>) {
@@ -18,22 +19,42 @@ async fn download_stream(url: &str, state: Arc<AppState>) -> Result<(), reqwest:
 
     let response = reqwest::Client::new().get(url).send().await?;
     let mut stream = response.bytes_stream();
-    let mut total_size = 0;
 
-    while let Some(chunk) = stream.next().await {
-        let bytes = chunk?;
-        total_size += bytes.len();
+    loop {
+        let Ok(bytes) = timeout(Duration::from_secs(5), stream.next()).await else {
+            println!("Error fetching stream, trying again...");
+            continue;
+        };
 
-        state.add_chunk(bytes).await;
-
-        if total_size % (1024 * 1024) == 0 {
-            // Log every MB
-            println!("Downloaded: {} MB", total_size / (1024 * 1024));
+        if handle_bytes(bytes, state.clone()).await {
+            break;
         }
     }
 
     state.mark_complete().await;
-    println!("Download complete! Total: {total_size} bytes");
+    println!("Download complete!");
 
     Ok(())
 }
+
+/// Returns if the end (= no more bytes) has been received
+async fn handle_bytes(
+    bytes: Option<Result<bytes::Bytes, reqwest::Error>>,
+    state: Arc<AppState>,
+) -> bool {
+    // TODO: switch checks, such that `.transpose()` is not needed anymore
+    let bytes = match bytes.transpose() {
+        Ok(bytes) => bytes,
+        Err(e) => {
+            println!("Error fetching chunk from stream: {e:?}");
+            return false;
+        }
+    };
+
+    if let Some(bytes) = bytes {
+        state.add_chunk(bytes).await;
+        false
+    } else {
+        true
+    }
+}