fix(snapshots): remove unwraps and handle errors

Author: TanmayArya-1p

crates/defs/src/error.rs

@@ -18,7 +18,7 @@ pub enum DbError {
     InvalidMagicBytes(String),
     VectorNotFound(uuid::Uuid),
     SnapshotRegistryError(String),
-    StorageEngineError(String),
+    SnapshotEngineError(String),
     InvalidDimension { expected: Dimension, got: Dimension },
     PointAlreadyExists { id: PointId },
     PointNotFound { id: PointId },

crates/snapshot/src/engine/mod.rs

@@ -36,7 +36,7 @@ impl SnapshotEngine {
         // acquire lock for worker_running
         let mut worker_running = self.worker_running.lock().map_err(|_| DbError::LockError)?;
         if !*worker_running {
-            return Err(DbError::StorageEngineError(
+            return Err(DbError::SnapshotEngineError(
                 "Worker thread not running".to_string(),
             ));
         }
@@ -50,7 +50,7 @@ impl SnapshotEngine {
         // acquire lock for worker_running
         let worker_running = self.worker_running.lock().map_err(|_| DbError::LockError)?;
         if !*worker_running {
-            return Err(DbError::StorageEngineError(
+            return Err(DbError::SnapshotEngineError(
                 "Worker thread not running".to_string(),
             ));
         }
@@ -82,7 +82,7 @@ impl SnapshotEngine {
         // acquire lock for worker_running
         let mut worker_running = self.worker_running.lock().map_err(|_| DbError::LockError)?;
         if *worker_running {
-            return Err(DbError::StorageEngineError(
+            return Err(DbError::SnapshotEngineError(
                 "Worker thread already running".to_string(),
             ));
         }
@@ -119,25 +119,48 @@ impl SnapshotEngine {
     ) -> Result<(), DbError> {
         let snapshot_path = db
             .lock()
-            .unwrap()
-            .create_snapshot(registry.lock().unwrap().dir().as_path())
-            .unwrap();
-        let snapshot_metadata = Metadata::parse(&snapshot_path).unwrap();
+            .map_err(|_| DbError::LockError)?
+            .create_snapshot(
+                registry
+                    .lock()
+                    .map_err(|_| DbError::LockError)?
+                    .dir()
+                    .as_path(),
+            )
+            .map_err(|err| {
+                DbError::SnapshotEngineError(format!("Could not create snapshot : {}", err))
+            })?;
+        let snapshot_metadata = Metadata::parse(&snapshot_path).map_err(|err| {
+            DbError::SnapshotEngineError(format!("Could not parse snapshot metadata: {}", err))
+        })?;
 
         // add the snapshot to registry
         registry
             .lock()
-            .unwrap()
+            .map_err(|_| DbError::LockError)?
             .add_snapshot(&snapshot_path)
-            .unwrap();
+            .map_err(|err| {
+                DbError::SnapshotEngineError(format!("Could not add snapshot to registry: {}", err))
+            })?;
 
         {
-            let mut queue = snapshot_queue.lock().unwrap();
+            let mut queue = snapshot_queue.lock().map_err(|_| DbError::LockError)?;
             queue.push_back(snapshot_metadata);
 
             while queue.len() > last_k {
-                let old = queue.pop_front().unwrap();
-                registry.lock().unwrap().mark_dead(old.small_id).unwrap();
+                let old = queue.pop_front().ok_or_else(|| {
+                    DbError::SnapshotEngineError("Snapshot metadata queue is empty".to_string())
+                })?;
+                registry
+                    .lock()
+                    .map_err(|_| DbError::LockError)?
+                    .mark_dead(old.small_id)
+                    .map_err(|err| {
+                        DbError::SnapshotEngineError(format!(
+                            "Could not mark snapshot as dead in registry: {}",
+                            err
+                        ))
+                    })?;
             }
             // drop queue lock
         }
@@ -156,10 +179,7 @@ impl SnapshotEngine {
     ) {
         loop {
             // acquire the lock and exit if its false
-            let worker_running = worker_running
-                .lock()
-                .map_err(|_| DbError::LockError)
-                .unwrap();
+            let worker_running = worker_running.lock().unwrap();
             if !*worker_running {
                 break;
             }

crates/snapshot/src/lib.rs

@@ -106,7 +106,9 @@ impl Snapshot {
                 "Storage checkpoint was not properly made".to_string(),
             ))?
             .to_str()
-            .unwrap()
+            .ok_or(DbError::SnapshotError(
+                "Storage checkpoint filename is not valid UTF-8".to_string(),
+            ))?
             .to_string();
 
         // create manifest file

crates/snapshot/src/util.rs

@@ -1,7 +1,7 @@
 use data_encoding::HEXLOWER;
 use sha2::{Digest, Sha256};
 use std::fs::File;
-use std::io::{BufReader, Error, Read};
+use std::io::{BufReader, Error, ErrorKind, Read};
 use std::path::PathBuf;
 
 use defs::{DbError, Magic};
@@ -92,7 +92,10 @@ pub fn compress_archive(path: &Path, files: &[&Path]) -> Result<(), Error> {
     let mut tar = Builder::new(enc);
 
     for file in files {
-        let rel_path = file.file_name().unwrap();
+        let rel_path = file.file_name().ok_or(Error::new(
+            ErrorKind::InvalidFilename,
+            "Could not create archive : invalid snapshot file",
+        ))?;
         let mut f = File::open(file)?;
         tar.append_file(rel_path, &mut f)?;
     }