Merge pull request deven96#72 from deven96/diretnan/fix-persistence

Fixing persistence issues from predicate index

Author: Iamdavidonuh

ahnlich/db/src/engine/predicate.rs

@@ -14,7 +14,8 @@ use serde::Serialize;
 use std::collections::HashSet as StdHashSet;
 use std::mem::size_of_val;
 
-type InnerPredicateIndex = ConcurrentHashMap<MetadataValue, ConcurrentHashSet<StoreKeyId>>;
+type InnerPredicateIndexVal = ConcurrentHashSet<StoreKeyId>;
+type InnerPredicateIndex = ConcurrentHashMap<MetadataValue, InnerPredicateIndexVal>;
 type InnerPredicateIndices = ConcurrentHashMap<MetadataKey, PredicateIndex>;
 
 /// Predicate indices are all the indexes referenced by their names
@@ -200,7 +201,35 @@ impl PredicateIndices {
 /// ids. This is essential in helping us filter down the entire dataset using a predicate before
 /// performing similarity algorithmic search
 #[derive(Debug, Serialize, Deserialize)]
-struct PredicateIndex(InnerPredicateIndex);
+struct PredicateIndex(#[serde(with = "custom_metadata_map")] InnerPredicateIndex);
+
+mod custom_metadata_map {
+    use super::*;
+    use serde::{self, Deserialize, Deserializer, Serialize, Serializer};
+
+    pub fn serialize<S>(map: &InnerPredicateIndex, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        let vec: Vec<(MetadataValue, InnerPredicateIndexVal)> = map
+            .iter(&map.guard())
+            .map(|(k, v)| (k.clone(), v.clone()))
+            .collect();
+        vec.serialize(serializer)
+    }
+
+    pub fn deserialize<'de, D>(deserializer: D) -> Result<InnerPredicateIndex, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        let vec: Vec<(MetadataValue, InnerPredicateIndexVal)> = Vec::deserialize(deserializer)?;
+        let map = ConcurrentHashMap::new();
+        for (k, v) in vec {
+            map.insert(k, v, &map.guard());
+        }
+        Ok(map)
+    }
+}
 
 impl PredicateIndex {
     fn size(&self) -> usize {

ahnlich/db/src/tests/server_test.rs

@@ -421,7 +421,13 @@ async fn test_server_with_persistence() {
             store: StoreName("Main".to_string()),
             inputs: vec![
                 (StoreKey(array![1.0, 1.1, 1.2, 1.3]), HashMap::new()),
-                (StoreKey(array![1.1, 1.2, 1.3, 1.4]), HashMap::new()),
+                (
+                    StoreKey(array![1.1, 1.2, 1.3, 1.4]),
+                    HashMap::from_iter([(
+                        MetadataKey::new("medal".into()),
+                        MetadataValue::Binary(vec![1, 2, 3]),
+                    )]),
+                ),
             ],
         },
         DBQuery::ListStores,
@@ -449,7 +455,7 @@ async fn test_server_with_persistence() {
         StoreInfo {
             name: StoreName("Main".to_string()),
             len: 2,
-            size_in_bytes: 1880,
+            size_in_bytes: 1936,
         },
     ]))));
     expected.push(Err(
@@ -460,7 +466,7 @@ async fn test_server_with_persistence() {
         StoreInfo {
             name: StoreName("Main".to_string()),
             len: 1,
-            size_in_bytes: 1808,
+            size_in_bytes: 1864,
         },
     ]))));
     let stream = TcpStream::connect(address).await.unwrap();
@@ -478,7 +484,11 @@ async fn test_server_with_persistence() {
     let address = server.local_addr().expect("Could not get local addr");
     let _ = tokio::spawn(async move { server.start().await });
     // Allow some time for the server to start
+    let file_metadata =
+        std::fs::metadata(&CONFIG_WITH_PERSISTENCE.persist_location.clone().unwrap()).unwrap();
+    assert!(file_metadata.len() > 0, "The persistence file is empty");
     tokio::time::sleep(Duration::from_millis(100)).await;
+    // check peristence was not overriden
     let message = ServerDBQuery::from_queries(&[
         // should error as store was loaded from previous persistence and main still exists
         DBQuery::CreateStore {
@@ -504,7 +514,10 @@ async fn test_server_with_persistence() {
     expected.push(Ok(ServerResponse::Del(0)));
     expected.push(Ok(ServerResponse::Get(vec![(
         StoreKey(array![1.1, 1.2, 1.3, 1.4]),
-        HashMap::new(),
+        HashMap::from_iter([(
+            MetadataKey::new("medal".into()),
+            MetadataValue::Binary(vec![1, 2, 3]),
+        )]),
     )])));
     let stream = TcpStream::connect(address).await.unwrap();
     let mut reader = BufReader::new(stream);

ahnlich/utils/src/persistence.rs

@@ -1,6 +1,7 @@
 use serde::de::DeserializeOwned;
 use serde::Serialize;
 use std::fs::File;
+use std::fs::OpenOptions;
 use std::io::BufReader;
 use std::path::Path;
 use std::sync::atomic::AtomicBool;
@@ -42,7 +43,10 @@ impl<T: Serialize + DeserializeOwned> Persistence<T> {
         persist_location: &std::path::PathBuf,
         persist_object: T,
     ) -> Self {
-        let _ = File::create(persist_location)
+        let _ = OpenOptions::new()
+            .append(true)
+            .create(true)
+            .open(persist_location)
             .expect("Persistence enabled but could not open peristence file");
         Self {
             write_flag,