Preparation for storage upgrade (#4)

src/cache/cache_container.rs

@@ -10,7 +10,7 @@ use crate::iterator::ScanDirection;
 use crate::schema::{KeyCodec, ValueCodec};
 use crate::{Operation, RawDbIter, ReadOnlyLock, Schema, SchemaKey, SchemaValue, DB};
 
-/// Holds collection of [`ChangeSet`]'s associated with particular Snapshot
+/// Holds a collection of [`ChangeSet`]'s associated with particular Snapshot
 /// and knows how to traverse them.
 /// Managed externally.
 /// Should be managed carefully, because discrepancy between `snapshots` and `to_parent` leads to panic
@@ -71,7 +71,7 @@ impl CacheContainer {
         let snapshot = self.snapshots.remove(snapshot_id).ok_or_else(|| {
             anyhow::anyhow!("Attempt to commit unknown snapshot id={}", snapshot_id)
         })?;
-        self.db.write_schemas(snapshot.into())
+        self.db.write_schemas(&snapshot.into())
     }
 
     /// Indicates, if CacheContainer has any snapshots in memory.
@@ -481,7 +481,7 @@ mod tests {
         let mut db_data = SchemaBatch::new();
         db_data.put::<S>(&one, &one).unwrap();
         db_data.put::<S>(&three, &three).unwrap();
-        db.write_schemas(db_data).unwrap();
+        db.write_schemas(&db_data).unwrap();
 
         let snapshot_manager = Arc::new(RwLock::new(CacheContainer::new(db, to_parent.into())));
 
@@ -529,7 +529,7 @@ mod tests {
 
         let mut db_data = SchemaBatch::new();
         db_data.put::<S>(&f1, &f1).unwrap();
-        db.write_schemas(db_data).unwrap();
+        db.write_schemas(&db_data).unwrap();
 
         let snapshot_manager = Arc::new(RwLock::new(CacheContainer::new(db, to_parent.into())));
 
@@ -674,7 +674,7 @@ mod tests {
         db_data.put::<S>(&f4, &f1).unwrap();
         db_data.put::<S>(&f0, &f1).unwrap();
         db_data.put::<S>(&f11, &f9).unwrap();
-        db.write_schemas(db_data).unwrap();
+        db.write_schemas(&db_data).unwrap();
         // DB Data
         // | key | value |
         // |   3 |     9 |

src/lib.rs

@@ -152,7 +152,7 @@ impl DB {
         // Used in tests only anyway.
         let mut batch = SchemaBatch::new();
         batch.put::<S>(key, value)?;
-        self.write_schemas(batch)
+        self.write_schemas(&batch)
     }
 
     /// Delete a single key from the database.
@@ -161,7 +161,7 @@ impl DB {
         // Used in tests only anyway.
         let mut batch = SchemaBatch::new();
         batch.delete::<S>(key)?;
-        self.write_schemas(batch)
+        self.write_schemas(&batch)
     }
 
     /// Removes the database entries in the range `["from", "to")` using default write options.
@@ -229,7 +229,7 @@ impl DB {
     }
 
     /// Writes a group of records wrapped in a [`SchemaBatch`].
-    pub fn write_schemas(&self, batch: SchemaBatch) -> anyhow::Result<()> {
+    pub fn write_schemas(&self, batch: &SchemaBatch) -> anyhow::Result<()> {
         let _timer = SCHEMADB_BATCH_COMMIT_LATENCY_SECONDS
             .with_label_values(&[self.name])
             .start_timer();

src/schema_batch.rs

@@ -82,7 +82,9 @@ impl SchemaBatch {
             .unwrap_or_default()
     }
 
-    pub(crate) fn merge(&mut self, other: SchemaBatch) {
+    /// Merge other [`SchemaBatch`] on top of this one.
+    /// Keys from other will overwrite keys in self.
+    pub fn merge(&mut self, other: SchemaBatch) {
         for (cf_name, other_cf_map) in other.last_writes {
             let cf_map = self.last_writes.entry(cf_name).or_default();
             cf_map.extend(other_cf_map);
@@ -291,4 +293,47 @@ mod tests {
                 .for_each(drop);
         }
     }
+
+    mod merge {
+        use super::*;
+
+        #[test]
+        fn test_simple_merge() {
+            let field_1 = TestField(1);
+            let field_2 = TestField(2);
+            let field_3 = TestField(3);
+            let field_4 = TestField(4);
+            let field_5 = TestField(5);
+
+            let mut batch1 = SchemaBatch::new();
+            batch1.put::<TestSchema1>(&field_1, &field_2).unwrap();
+            batch1.put::<TestSchema1>(&field_3, &field_1).unwrap();
+            batch1.put::<TestSchema1>(&field_5, &field_4).unwrap();
+
+            let mut batch2 = SchemaBatch::new();
+            batch2.put::<TestSchema1>(&field_1, &field_3).unwrap();
+            batch2.put::<TestSchema1>(&field_4, &field_2).unwrap();
+            batch2.delete::<TestSchema1>(&field_5).unwrap();
+
+            batch1.merge(batch2);
+
+            let get_value = |field: &TestField| -> Option<TestField> {
+                batch1
+                    .get::<TestSchema1>(field)
+                    .unwrap()
+                    .unwrap()
+                    .decode_value::<TestSchema1>()
+                    .unwrap()
+            };
+
+            assert_eq!(Some(field_3), get_value(&field_1), "key (1) wasn't updated");
+            assert_eq!(
+                Some(field_1),
+                get_value(&field_3),
+                "key (3) has been be changed, when it shouldn't"
+            );
+            assert_eq!(Some(field_2), get_value(&field_4), "key (4) wasn't added");
+            assert_eq!(None, get_value(&field_5), "key (5) wasn't deleted");
+        }
+    }
 }

tests/db_test.rs

@@ -171,7 +171,7 @@ fn test_single_schema_batch() {
         .put::<TestSchema2>(&TestField(5), &TestField(5))
         .unwrap();
 
-    db.write_schemas(db_batch).unwrap();
+    db.write_schemas(&db_batch).unwrap();
 
     assert_eq!(
         collect_values::<TestSchema1>(&db),
@@ -198,7 +198,7 @@ fn test_two_schema_batches() {
         .put::<TestSchema1>(&TestField(2), &TestField(2))
         .unwrap();
     db_batch1.delete::<TestSchema1>(&TestField(2)).unwrap();
-    db.write_schemas(db_batch1).unwrap();
+    db.write_schemas(&db_batch1).unwrap();
 
     assert_eq!(
         collect_values::<TestSchema1>(&db),
@@ -216,7 +216,7 @@ fn test_two_schema_batches() {
     db_batch2
         .put::<TestSchema2>(&TestField(5), &TestField(5))
         .unwrap();
-    db.write_schemas(db_batch2).unwrap();
+    db.write_schemas(&db_batch2).unwrap();
 
     assert_eq!(
         collect_values::<TestSchema1>(&db),
@@ -292,7 +292,7 @@ fn test_report_size() {
         db_batch
             .put::<TestSchema2>(&TestField(i), &TestField(i))
             .unwrap();
-        db.write_schemas(db_batch).unwrap();
+        db.write_schemas(&db_batch).unwrap();
     }
 
     db.flush_cf("TestCF1").unwrap();