Implement snapshots

Cargo.toml

@@ -8,6 +8,7 @@ members = [
     "crates/http",
     "crates/tui",
     "crates/grpc",
+    "crates/snapshot",
 ]
 
 [workspace.package]
@@ -51,5 +52,6 @@ grpc = { path = "crates/grpc" }
 http = { path = "crates/http" }
 index = { path = "crates/index" }
 server = { path = "crates/server" }
+snapshot = { path = "crates/snapshot" }
 storage = { path = "crates/storage" }
 tui = { path = "crates/tui" }

crates/api/Cargo.toml

@@ -10,6 +10,7 @@ license.workspace = true
 defs.workspace = true
 index.workspace = true
 snafu.workspace = true
+snapshot.workspace = true
 storage.workspace = true
 tempfile.workspace = true
 uuid.workspace = true

crates/api/src/lib.rs

@@ -1,13 +1,15 @@
-use defs::{Dimension, IndexedVector, Similarity};
-
+use defs::{DbError, Dimension, IndexedVector, Similarity, SnapshottableDb};
 use defs::{DenseVector, Payload, Point, PointId};
 use index::hnsw::HnswIndex;
-use std::path::PathBuf;
+use index::kd_tree::index::KDTree;
+use std::path::{Path, PathBuf};
+use tempfile::tempdir;
 // use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::{Arc, RwLock};
 
-use index::flat::FlatIndex;
+use index::flat::index::FlatIndex;
 use index::{IndexType, VectorIndex};
+use snapshot::Snapshot;
 use storage::rocks_db::RocksDbStorage;
 use storage::{StorageEngine, StorageType, VectorPage};
 
@@ -152,6 +154,33 @@ impl VectorDb {
     }
 }
 
+impl SnapshottableDb for VectorDb {
+    fn create_snapshot(&self, dir_path: &Path) -> Result<PathBuf, DbError> {
+        if !dir_path.is_dir() {
+            return Err(DbError::SnapshotError(format!(
+                "Invalid path: {}",
+                dir_path.display()
+            )));
+        }
+
+        let index_snapshot = self
+            .index
+            .read()
+            .map_err(|_| DbError::LockError)?
+            .snapshot()?;
+
+        let tempdir = tempdir().unwrap();
+        let storage_checkpoint = self.storage.checkpoint_at(tempdir.path()).map_err(|e| {
+            DbError::StorageCheckpointError(format!("Could not create storage checkpoint: {e}"))
+        })?;
+
+        let snapshot = Snapshot::new(index_snapshot, storage_checkpoint, self.dimension)?;
+        let snapshot_path = snapshot.save(dir_path)?;
+
+        Ok(snapshot_path)
+    }
+}
+
 #[derive(Debug)]
 pub struct DbConfig {
     pub storage_type: StorageType,
@@ -161,6 +190,28 @@ pub struct DbConfig {
     pub similarity: Similarity,
 }
 
+#[derive(Debug)]
+pub struct DbRestoreConfig {
+    pub data_path: PathBuf,
+    pub snapshot_path: PathBuf,
+}
+
+impl DbRestoreConfig {
+    pub fn new(data_path: PathBuf, snapshot_path: PathBuf) -> Self {
+        Self {
+            data_path,
+            snapshot_path,
+        }
+    }
+}
+
+pub fn restore_from_snapshot(config: &DbRestoreConfig) -> Result<VectorDb, DbError> {
+    // restore the index from the snapshot
+    let (storage_engine, index, dimensions) =
+        Snapshot::load(&config.snapshot_path, &config.data_path)?;
+    Ok(VectorDb::_new(storage_engine, index, dimensions))
+}
+
 pub fn init_api(config: DbConfig) -> Result<VectorDb> {
     // Initialize the storage engine
     let storage = match config.storage_type {
@@ -171,11 +222,11 @@ pub fn init_api(config: DbConfig) -> Result<VectorDb> {
     // Initialize the vector index
     let index: Arc<RwLock<dyn VectorIndex>> = match config.index_type {
         IndexType::Flat => Arc::new(RwLock::new(FlatIndex::new())),
+        IndexType::KDTree => Arc::new(RwLock::new(KDTree::build_empty(config.dimension))),
         IndexType::HNSW => Arc::new(RwLock::new(HnswIndex::new(
             config.similarity,
             config.dimension,
         ))),
-        _ => Arc::new(RwLock::new(FlatIndex::new())),
     };
 
     // Init the db
@@ -192,8 +243,11 @@ mod tests {
 
     // TODO: Add more exhaustive tests
 
+    use std::sync::Mutex;
+
     use super::*;
     use defs::ContentType;
+    use snapshot::{engine::SnapshotEngine, registry::local::LocalRegistry};
     use tempfile::{TempDir, tempdir};
 
     // Helper function to create a test database
@@ -419,4 +473,143 @@ mod tests {
         let inserted = db.build_index().unwrap();
         assert_eq!(inserted, 10);
     }
+
+    #[test]
+    fn test_create_and_load_snapshot() {
+        let (old_db, temp_dir) = create_test_db();
+
+        let v1 = vec![0.0, 1.0, 2.0];
+        let v2 = vec![3.0, 4.0, 5.0];
+        let v3 = vec![6.0, 7.0, 8.0];
+
+        let id1 = old_db
+            .insert(
+                v1.clone(),
+                Payload {
+                    content_type: ContentType::Text,
+                    content: "test".to_string(),
+                },
+            )
+            .unwrap();
+
+        let id2 = old_db
+            .insert(
+                v2.clone(),
+                Payload {
+                    content_type: ContentType::Text,
+                    content: "test".to_string(),
+                },
+            )
+            .unwrap();
+
+        let temp_snapshot_dir = tempdir().unwrap();
+        let snapshot_path = old_db.create_snapshot(temp_snapshot_dir.path()).unwrap();
+
+        // insert v3 after snapshot
+        let id3 = old_db
+            .insert(
+                v3.clone(),
+                Payload {
+                    content_type: ContentType::Text,
+                    content: "test".to_string(),
+                },
+            )
+            .unwrap();
+
+        let reload_config = DbRestoreConfig {
+            data_path: temp_dir.path().to_path_buf(),
+            snapshot_path,
+        };
+
+        std::mem::drop(old_db);
+        let loaded_db = restore_from_snapshot(&reload_config).unwrap();
+
+        assert!(loaded_db.get(id1).unwrap_or(None).is_some());
+        assert!(loaded_db.get(id2).unwrap_or(None).is_some());
+        assert!(loaded_db.get(id3).unwrap_or(None).is_none()); // v3 was inserted after snapshot was taken
+
+        // vector restore check
+        assert!(loaded_db.get(id1).unwrap().unwrap().vector.unwrap() == v1);
+        assert!(loaded_db.get(id2).unwrap().unwrap().vector.unwrap() == v2);
+    }
+
+    #[test]
+    fn test_snapshot_engine() {
+        let (_db, _temp_dir) = create_test_db();
+        let db = Arc::new(Mutex::new(_db));
+
+        let registry_tempdir = tempdir().unwrap();
+
+        let registry = Arc::new(Mutex::new(
+            LocalRegistry::new(registry_tempdir.path()).unwrap(),
+        ));
+
+        let last_k = 4;
+        let mut se = SnapshotEngine::new(last_k, db.clone(), registry.clone());
+
+        let v1 = vec![0.0, 1.0, 2.0];
+        let v2 = vec![3.0, 4.0, 5.0];
+        let v3 = vec![6.0, 7.0, 8.0];
+
+        let test_vectors = vec![v1.clone(), v2.clone(), v3.clone()];
+        let mut inserted_ids = Vec::new();
+
+        for (i, vector) in test_vectors.clone().into_iter().enumerate() {
+            se.snapshot().unwrap();
+            let id = db
+                .lock()
+                .unwrap()
+                .insert(
+                    vector.clone(),
+                    Payload {
+                        content_type: ContentType::Text,
+                        content: format!("{}", i),
+                    },
+                )
+                .unwrap();
+            inserted_ids.push(id);
+        }
+        se.snapshot().unwrap();
+        let snapshots = se.list_alive_snapshots().unwrap();
+
+        // asserting these cases:
+        // snapshot 0 : no vectors
+        // snapshot 1 : v1
+        // snapshot 2 : v1, v2
+        // snapshot 3 : v1, v2, v3
+
+        std::mem::drop(db);
+        std::mem::drop(se);
+
+        for (i, snapshot) in snapshots.iter().enumerate() {
+            let temp_dir = tempdir().unwrap();
+            let db = restore_from_snapshot(&DbRestoreConfig {
+                data_path: temp_dir.path().to_path_buf(),
+                snapshot_path: snapshot.path.clone(),
+            })
+            .unwrap();
+            for j in 0..i {
+                // test if point is present
+                assert!(db.get(inserted_ids[j]).unwrap_or(None).is_some());
+                // test vector restore
+                assert!(
+                    db.get(inserted_ids[j]).unwrap().unwrap().vector.unwrap() == test_vectors[j]
+                );
+                // test payload restore
+                assert!(
+                    db.get(inserted_ids[j])
+                        .unwrap()
+                        .unwrap()
+                        .payload
+                        .unwrap()
+                        .content
+                        == format!("{}", j)
+                );
+            }
+            for absent_id in inserted_ids.iter().skip(i) {
+                assert!(db.get(*absent_id).unwrap_or(None).is_none());
+            }
+            std::mem::drop(db);
+        }
+    }
 }

crates/defs/src/error.rs

@@ -1,5 +1,29 @@
 use std::io;
 
+use crate::{Dimension, PointId};
+#[derive(Debug, PartialEq, Eq)]
+pub enum DbError {
+    ParseError,
+    StorageError(String),
+    SerializationError(String),
+    DeserializationError,
+    IndexError(String),
+    LockError,
+    IndexInitError, //TODO: Change this
+    UnsupportedSimilarity,
+    DimensionMismatch,
+    SnapshotError(String),
+    StorageInitializationError,
+    StorageCheckpointError(String),
+    InvalidMagicBytes(String),
+    VectorNotFound(uuid::Uuid),
+    SnapshotRegistryError(String),
+    SnapshotEngineError(String),
+    InvalidDimension { expected: Dimension, got: Dimension },
+    PointAlreadyExists { id: PointId },
+    PointNotFound { id: PointId },
+}
+
 #[derive(Debug)]
 pub enum ServerError {
     Bind(io::Error),

crates/defs/src/lib.rs

@@ -3,4 +3,10 @@ pub mod types;
 
 // Without re-exports, users would need to write defs::types::SomeType instead of just defs::SomeType. Re-exports simplify the API by flattening the module hierarchy. The * means "everything public" from that module.
 pub use error::*;
+use std::path::{Path, PathBuf};
 pub use types::*;
+
+// hoisted trait so it can be used by the snapshots crate
+pub trait SnapshottableDb: Send + Sync {
+    fn create_snapshot(&self, dir_path: &Path) -> Result<PathBuf, DbError>;
+}

crates/defs/src/types.rs

@@ -15,6 +15,8 @@ pub type Dimension = usize;
 // Sparse vector implementation not supported yet. Refer lib/sparse/src/common/sparse_vector.rs
 pub type DenseVector = Vec<Element>;
 
+pub type Magic = [u8; 4];
+
 pub enum StoredVector {
     Dense(DenseVector),
 }
@@ -45,7 +47,7 @@ pub struct IndexedVector {
     pub vector: DenseVector,
 }
 
-#[derive(Debug, Deserialize, Copy, Clone)]
+#[derive(Debug, Serialize, Deserialize, Copy, Clone)]
 pub enum Similarity {
     Euclidean,
     Manhattan,

crates/grpc/src/error.rs

@@ -120,6 +120,21 @@ impl From<storage::error::StorageError> for GrpcError {
             StorageError::Deserialization { id, source: _ } => GrpcError::Internal {
                 message: format!("failed to deserialize point {}", id),
             },
+            StorageError::RocksDbInitialization {} => GrpcError::Internal {
+                message: "failed to initialize storage".to_string(),
+            },
+            StorageError::RocksDbCheckpointMsg { msg } => GrpcError::Internal {
+                message: format!("checkpoint error: {}", msg),
+            },
+            StorageError::RocksDbCheckpointIo { msg, source: _ } => GrpcError::Internal {
+                message: format!("checkpoint io error: {}", msg),
+            },
+            StorageError::RocksDbCheckpoint { source: _ } => GrpcError::Internal {
+                message: "checkpoint error".to_string(),
+            },
+            StorageError::RocksDbFlush { source: _ } => GrpcError::Internal {
+                message: "flush error".to_string(),
+            },
         }
     }
 }

crates/http/src/handler.rs

@@ -126,7 +126,12 @@ fn api_error_to_response(err: &ApiError) -> (StatusCode, String) {
             | StorageError::RocksDbDelete { .. }
             | StorageError::RocksDbIteration { .. }
             | StorageError::Serialization { .. }
-            | StorageError::Deserialization { .. } => {
+            | StorageError::Deserialization { .. }
+            | StorageError::RocksDbCheckpoint { .. }
+            | StorageError::RocksDbFlush { .. }
+            | StorageError::RocksDbInitialization { .. }
+            | StorageError::RocksDbCheckpointMsg { .. }
+            | StorageError::RocksDbCheckpointIo { .. } => {
                 (StatusCode::INTERNAL_SERVER_ERROR, source.to_string())
             }
         },

crates/index/Cargo.toml

@@ -7,7 +7,10 @@ edition.workspace = true
 license.workspace = true
 
 [dependencies]
+bincode.workspace = true
 defs.workspace = true
 rand.workspace = true
+serde.workspace = true
 snafu.workspace = true
+storage.workspace = true
 uuid.workspace = true