From 16e3bfe2c5a685625bf2fca5d89ae786942ae126 Mon Sep 17 00:00:00 2001
From: james58899 <james59988@gmail.com>
Date: Thu, 9 Nov 2023 12:15:51 +0000
Subject: [PATCH] Saving download state before create temp file

---
 src/main.rs        |  2 +-
 src/route/cache.rs | 18 +++++++++++++-----
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index c69ef83..43e5d6d 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -137,7 +137,7 @@ struct Args {
     proxy: Option<String>,
 }
 
-type DownloadState = RwLock<HashMap<[u8; 20], (Arc<TempPath>, watch::Receiver<u64>)>>;
+type DownloadState = RwLock<HashMap<[u8; 20], (watch::Receiver<Option<Arc<TempPath>>>, watch::Receiver<u64>)>>;
 
 struct AppState {
     runtime: Handle,
diff --git a/src/route/cache.rs b/src/route/cache.rs
index 9ac4bab..6454806 100644
--- a/src/route/cache.rs
+++ b/src/route/cache.rs
@@ -83,14 +83,22 @@ async fn hath(
 
     // Check if the file is already downloading
     let download_state = data.download_state.read().get(&info.hash()).cloned();
-    let (temp_path, mut rx) = if let Some(state) = download_state {
-        state
+    let (temp_path, mut rx) = if let Some((mut tempfile, progress)) = download_state {
+        let tempfile = tempfile.wait_for(Option::is_some).await;
+        if let Err(err) = tempfile {
+            error!("Waiting tempfile create error: {}", err);
+            data.download_state.write().remove(&info.hash());
+            return HttpResponse::NotFound().body("An error has occurred. (404)");
+        }
+        (tempfile.unwrap().as_ref().unwrap().clone(), progress)
     } else {
-        let temp_path = Arc::new(data.cache_manager.create_temp_file().await);
+        // Tracking download progress
+        let (temp_tx, temp_rx) = watch::channel(None); // Tempfile
         let (tx, rx) = watch::channel(0); // Download progress
+        data.download_state.write().insert(info.hash(), (temp_rx.clone(), rx.clone()));
 
-        // Tracking download progress
-        data.download_state.write().insert(info.hash(), (temp_path.clone(), rx.clone()));
+        let temp_path = Arc::new(data.cache_manager.create_temp_file().await);
+        temp_tx.send_replace(Some(temp_path.clone()));
 
         let sources = match data.rpc.sr_fetch(file_index, xres, &file_id).await {
             Some(v) => v,