fix: file open handling, remove upon success

graphops · Mar 25, 2024 · bbede0c · bbede0c
1 parent ee20806
commit bbede0c
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 17 deletions.
diff --git a/file-exchange/src/download_client/mod.rs b/file-exchange/src/download_client/mod.rs
@@ -119,7 +119,7 @@ impl Downloader {
         let store = Store::new(&args.storage_method).expect("Create store");
 
         let target_chunks = if let Some(file_path) = &args.progress_file {
-            let map = read_json_to_map(&file_path).expect("Progress cache ill-formatted");
+            let map = read_json_to_map(file_path).expect("Progress cache ill-formatted");
             Arc::new(StdMutex::new(map))
         } else {
             Arc::new(StdMutex::new(HashMap::new()))
@@ -156,7 +156,7 @@ impl Downloader {
     pub fn init_target_chunks(&self, bundle: &Bundle) {
         if let Some(file_path) = &self.config.progress_file {
             let mut target_chunks = self.target_chunks.lock().unwrap();
-            *target_chunks = read_json_to_map(&file_path).expect("Progress cache ill-formatted");
+            *target_chunks = read_json_to_map(file_path).expect("Progress cache ill-formatted");
         }
         for file_manifest_meta in &bundle.file_manifests {
             let mut target_chunks = self.target_chunks.lock().unwrap();
@@ -207,12 +207,17 @@ impl Downloader {
             tracing::warn!(msg);
             // store progress into a json file: {hash: missing_chunk_indices}
             if let Some(file_path) = &self.config.progress_file {
-                store_map_as_json(&incomplete_progresses, &file_path)?;
+                store_map_as_json(&incomplete_progresses, file_path)?;
             };
             return Err(Error::DataUnavailable(msg));
-        } else {
-            tracing::info!("File manifests download completed");
         }
+
+        tracing::info!("File manifests download completed");
+
+        if let Some(file_path) = &self.config.progress_file {
+            let _ = fs::remove_file(file_path);
+        };
+
         Ok(())
     }
 

diff --git a/file-exchange/src/util.rs b/file-exchange/src/util.rs
@@ -247,21 +247,27 @@ pub fn store_map_as_json(
 
 // Reads a JSON file and converts it into a HashMap<String, HashSet<u64>>.
 pub fn read_json_to_map(file_path: &str) -> Result<HashMap<String, HashSet<u64>>, Error> {
-    let mut file = File::open(file_path).map_err(Error::FileIOError)?;
+    match File::open(file_path) {
+        Ok(mut file) => {
+            let mut contents = String::new();
+            file.read_to_string(&mut contents)
+                .map_err(Error::FileIOError)?;
 
-    let mut contents = String::new();
-    file.read_to_string(&mut contents)
-        .map_err(Error::FileIOError)?;
-
-    let temp_map: HashMap<String, Vec<u64>> =
-        serde_json::from_str(&contents).map_err(Error::JsonError)?;
+            let temp_map: HashMap<String, Vec<u64>> =
+                serde_json::from_str(&contents).map_err(Error::JsonError)?;
 
-    let map: HashMap<String, HashSet<u64>> = temp_map
-        .into_iter()
-        .map(|(key, vec)| (key, vec.into_iter().collect::<HashSet<u64>>()))
-        .collect();
+            let map: HashMap<String, HashSet<u64>> = temp_map
+                .into_iter()
+                .map(|(key, vec)| (key, vec.into_iter().collect::<HashSet<u64>>()))
+                .collect();
 
-    Ok(map)
+            Ok(map)
+        }
+        Err(e) => {
+            tracing::info!("Failed to open file at '{}'. Error: {}.", file_path, e);
+            Ok(HashMap::new())
+        }
+    }
 }
 
 #[cfg(test)]