[WIP] Implement snapshots

Cargo.toml

@@ -8,6 +8,7 @@ members = [
     "crates/http",
     "crates/tui",
     "crates/grpc",
+    "crates/snapshot",
 ]
 
 [workspace.package]
@@ -49,5 +50,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

@@ -9,6 +9,7 @@ license.workspace = true
 [dependencies]
 defs.workspace = true
 index.workspace = true
+snapshot.workspace = true
 storage.workspace = true
 tempfile.workspace = true
 uuid.workspace = true

crates/api/src/lib.rs

@@ -1,12 +1,15 @@
-use defs::{DbError, IndexedVector, Similarity};
+use defs::{DbError, IndexedVector, Similarity, SnapshottableDb};
 
 use defs::{DenseVector, Payload, Point, PointId};
-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};
 
@@ -131,6 +134,31 @@ 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())?;
+
+        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,
@@ -139,6 +167,28 @@ pub struct DbConfig {
     pub dimension: usize,
 }
 
+#[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, DbError> {
     // Initialize the storage engine
     let storage = match config.storage_type {
@@ -149,7 +199,8 @@ pub fn init_api(config: DbConfig) -> Result<VectorDb, DbError> {
     // Initialize the vector index
     let index: Arc<RwLock<dyn VectorIndex>> = match config.index_type {
         IndexType::Flat => Arc::new(RwLock::new(FlatIndex::new())),
-        _ => Arc::new(RwLock::new(FlatIndex::new())),
+        IndexType::KDTree => Arc::new(RwLock::new(KDTree::build_empty(config.dimension))),
+        _ => Arc::new(RwLock::new(FlatIndex::new())), // TODO: add hnsw here
     };
 
     // Init the db
@@ -166,25 +217,28 @@ mod tests {
 
     // TODO: Add more exhaustive tests
 
+    use std::sync::Mutex;
+
     use super::*;
     use defs::ContentType;
-    use tempfile::tempdir;
+    use snapshot::{engine::SnapshotEngine, registry::local::LocalRegistry};
+    use tempfile::{TempDir, tempdir};
 
     // Helper function to create a test database
-    fn create_test_db() -> VectorDb {
+    fn create_test_db() -> (VectorDb, TempDir) {
         let temp_dir = tempdir().unwrap();
         let config = DbConfig {
             storage_type: StorageType::RocksDb,
             index_type: IndexType::Flat,
             data_path: temp_dir.path().to_path_buf(),
             dimension: 3,
         };
-        init_api(config).unwrap()
+        (init_api(config).unwrap(), temp_dir)
     }
 
     #[test]
     fn test_insert_and_get() {
-        let db = create_test_db();
+        let (db, _temp_dir) = create_test_db();
         let vector = vec![1.0, 2.0, 3.0];
         let payload = Payload {
             content_type: ContentType::Text,
@@ -209,7 +263,7 @@ mod tests {
 
     #[test]
     fn test_dimension_mismatch() {
-        let db = create_test_db();
+        let (db, _temp_dir) = create_test_db();
         let v1 = vec![1.0, 2.0, 3.0];
         let v2 = vec![1.0, 2.0];
         let payload = defs::Payload {
@@ -228,7 +282,7 @@ mod tests {
 
     #[test]
     fn test_delete() {
-        let db = create_test_db();
+        let (db, _temp_dir) = create_test_db();
         let vector = vec![1.0, 2.0, 3.0];
         let payload = Payload {
             content_type: ContentType::Text,
@@ -251,7 +305,7 @@ mod tests {
 
     #[test]
     fn test_search() {
-        let db = create_test_db();
+        let (db, _temp_dir) = create_test_db();
 
         // Insert some points
         let vectors = vec![
@@ -280,7 +334,7 @@ mod tests {
 
     #[test]
     fn test_search_limit() {
-        let db = create_test_db();
+        let (db, _temp_dir) = create_test_db();
 
         // Insert 5 points
         let mut ids = Vec::new();
@@ -307,7 +361,7 @@ mod tests {
 
     #[test]
     fn test_empty_database() {
-        let db = create_test_db();
+        let (db, _temp_dir) = create_test_db();
 
         // Get non-existent point
         assert!(db.get(Uuid::new_v4()).unwrap().is_none());
@@ -319,7 +373,7 @@ mod tests {
 
     #[test]
     fn test_list_vectors() {
-        let db = create_test_db();
+        let (db, _temp_dir) = create_test_db();
         // insert some points
         let mut ids = Vec::new();
         for i in 0..10 {
@@ -350,7 +404,7 @@ mod tests {
 
     #[test]
     fn test_build_index() {
-        let db = create_test_db();
+        let (db, _temp_dir) = create_test_db();
 
         // insert some points
         for i in 0..10 {
@@ -370,4 +424,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

@@ -7,7 +7,16 @@ pub enum DbError {
     DeserializationError,
     IndexError(String),
     LockError,
+    IndexInitError, //TODO: Change this
+    UnsupportedSimilarity,
     DimensionMismatch,
+    SnapshotError(String),
+    StorageInitializationError,
+    StorageCheckpointError(String),
+    InvalidMagicBytes(String),
+    VectorNotFound(uuid::Uuid),
+    SnapshotRegistryError(String),
+    StorageEngineError(String),
 }
 
 #[derive(Debug)]

crates/defs/src/lib.rs

@@ -1,6 +1,13 @@
 pub mod error;
 pub mod types;
 
+use std::path::{Path, PathBuf};
+
 // 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::*;
 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

@@ -51,6 +51,8 @@ pub enum Similarity {
     Cosine,
 }
 
+pub type Magic = [u8; 4];
+
 // Struct which stores the distance between a vector and query vector and implements ordering traits
 #[derive(Copy, Clone)]
 pub struct DistanceOrderedVector<'q> {