sled_playground: Use batching and transactions.

script/research/sled_playground/src/main.rs

@@ -16,200 +16,74 @@
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
 
-use sled::{Config, Db, IVec};
+use sled::{transaction::ConflictableTransactionError, Config, Transactional};
 
-#[derive(Debug, PartialEq)]
-struct MyBullshitError;
+pub mod overlay;
+use overlay::SledOverlay;
 
-struct SledCache {
-    // A sequence to execute
-    sequence: Vec<(IVec, Option<IVec>)>
-}
-
-impl SledCache {
-    fn new() -> Self {
-        Self {sequence: vec![]}
-    }
-    
-    // Check if a key is cached
-    fn get<K: AsRef<[u8]> + Copy>(&self, key: K) -> Option<IVec> 
-    where 
-        K: Into<IVec>
-    {
-        let key = key.into();
-        // TODO: optimize this
-        for pair in &self.sequence {
-            if pair.0 == key {
-                return pair.1.clone()
-            }
-        }
-        None
-    }
-    
-    // Set a key to a new value
-    fn push<K>(&mut self, key: K, value: Option<IVec>)
-    where
-        K: Into<IVec>,
-    {
-        self.sequence.push((key.into(), value));
-    }
-    
-    // Set a key to a new value at front of sequence
-    fn insert<K>(&mut self, key: K, value: Option<IVec>)
-    where
-        K: Into<IVec>,
-    {
-        self.sequence.insert(0, (key.into(), value));
-    }
-}
-
-struct SledOverlay {
-    // Actual sled storage
-    sled: Db,    
-    // Cached writes sequence
-    // NOTE: this should be per Tree
-    writes: SledCache,
-}
-
-impl SledOverlay {
-    fn new(sled: Db) -> Self {
-        Self {sled, writes: SledCache::new()}
-    }
-    
-    // Check if a key is cached for changes, otherwise get it
-    // from sled directly
-    fn get<K: AsRef<[u8]> + Copy>(&self, key: K) -> Result<Option<IVec>, sled::Error> 
-    where 
-        K: Into<IVec>
-    {
-        if let Some(v) = self.writes.get(&key.into()) {
-            return Ok(Some(v.clone()))
-        }
-        
-        self.sled.get(key)
-    }
-    
-    // Set a key to a new value
-    fn insert<K, V>(&mut self, key: K, value: V)
-    where
-        K: Into<IVec>,
-        V: Into<IVec>,
-    {
-        self.writes.push(key, Some(value.into()));
-    }
-    
-    // Apply all writes and flush sled db
-    fn flush(&mut self) -> Result<(), sled::Error> {
-        let mut rollback = SledCache::new();
-        for (k, v_opt) in &self.writes.sequence {
-            if let Some(v) = v_opt {
-                let old = self.sled.insert(k, v)?;
-                rollback.insert(k, old);
-            } else {
-                let old = self.sled.remove(k)?;
-                rollback.insert(k, old);
-            }
-        }
-        
-        self.sled.flush()?;
-        self.writes = SledCache::new();
-        
-        Ok(())
-    }
-    
-    // Apply all writes, rollback to previous and flush sled db
-    fn flush_rollback(&mut self) -> Result<(), sled::Error> {
-        // Rollback keeps track of writes in reverse order
-        let mut rollback = SledCache::new();
-        let res = {
-            for (k, v_opt) in &self.writes.sequence {
-                if let Some(v) = v_opt {
-                    let old = self.sled.insert(k, v)?;
-                    rollback.insert(k, old);
-                } else {
-                    let old = self.sled.remove(k)?;
-                    rollback.insert(k, old);
-                }
-            }
-            
-            MyBullshitError
-        };
-        
-        // Execute rollback in case of error
-        if res == MyBullshitError {
-            for (k, v_opt) in &rollback.sequence {
-                if let Some(v) = v_opt {
-                    let _old = self.sled.insert(k, v)?;
-                } else {
-                    let _old = self.sled.remove(k)?;
-                }
-            }
-        }
-        
-        self.sled.flush()?;
-        self.writes = SledCache::new();
-        
-        Ok(())
-    }
-}
+const TREE_1: &str = "_tree1";
+const TREE_2: &str = "_tree2";
 
-// This examples showcases a serial apply of writes
-// where if one fails we can detect it and rollback.
 fn main() -> Result<(), sled::Error> {
     // Initialize database overlay
     let config = Config::new().temporary(true);
     let db = config.open()?;
-    let mut overlay = SledOverlay::new(db.clone());
-    
-    // Insert some values to cache
-    overlay.insert("key_a", "val_a");
-    overlay.insert("key_b", "val_b");
-    overlay.insert("key_c", "val_c");
-    
-    assert_eq!(overlay.get("key_a")?, Some("val_a".into()));
-    assert_eq!(overlay.get("key_b")?, Some("val_b".into()));
-    assert_eq!(overlay.get("key_c")?, Some("val_c".into()));
-    
+
+    let tree_1 = db.open_tree(TREE_1)?;
+    let tree_2 = db.open_tree(TREE_2)?;
+    let mut overlay_1 = SledOverlay::new(&tree_1);
+    let mut overlay_2 = SledOverlay::new(&tree_2);
+
+    // Insert some values to the overlays
+    overlay_1.insert(b"key_a", b"val_a")?;
+    overlay_1.insert(b"key_b", b"val_b")?;
+    overlay_1.insert(b"key_c", b"val_c")?;
+
+    overlay_2.insert(b"key_d", b"val_d")?;
+    overlay_2.insert(b"key_e", b"val_e")?;
+    overlay_2.insert(b"key_f", b"val_f")?;
+
+    // Verify they are in the overlays
+    assert_eq!(overlay_1.get(b"key_a")?, Some(b"val_a".into()));
+    assert_eq!(overlay_1.get(b"key_b")?, Some(b"val_b".into()));
+    assert_eq!(overlay_1.get(b"key_c")?, Some(b"val_c".into()));
+
+    assert_eq!(overlay_2.get(b"key_d")?, Some(b"val_d".into()));
+    assert_eq!(overlay_2.get(b"key_e")?, Some(b"val_e".into()));
+    assert_eq!(overlay_2.get(b"key_f")?, Some(b"val_f".into()));
+
     // Verify they are not in sled
-    assert_eq!(db.get(b"key_a")?, None);
-    assert_eq!(db.get(b"key_b")?, None);
-    assert_eq!(db.get(b"key_c")?, None);
-    
-    // Now we write them to sled
-    overlay.flush()?;
-    
+    assert_eq!(tree_1.get(b"key_a")?, None);
+    assert_eq!(tree_1.get(b"key_b")?, None);
+    assert_eq!(tree_1.get(b"key_c")?, None);
+
+    assert_eq!(tree_2.get(b"key_d")?, None);
+    assert_eq!(tree_2.get(b"key_e")?, None);
+    assert_eq!(tree_2.get(b"key_f")?, None);
+
+    // Aggregate all the batches for writing
+    let mut batches = vec![];
+    batches.push(overlay_1.aggregate());
+    batches.push(overlay_2.aggregate());
+
+    // Now we write them to sled (this should be wrapped in a macro maybe)
+    (&tree_1, &tree_2)
+        .transaction(|(tree_1, tree_2)| {
+            tree_1.apply_batch(&batches[0])?;
+            tree_2.apply_batch(&batches[1])?;
+
+            Ok::<(), ConflictableTransactionError<sled::Error>>(())
+        })
+        .unwrap();
+
     // Verify sled contains keys
-    assert_eq!(db.get("key_a")?, Some("val_a".into()));
-    assert_eq!(db.get("key_b")?, Some("val_b".into()));
-    assert_eq!(db.get("key_c")?, Some("val_c".into()));
-    
-    // Perform the same steps, but assume an error occured
-    // during flashing, so the sled performs a rollback.
-    let config = Config::new().temporary(true);
-    let db = config.open()?;
-    let mut overlay = SledOverlay::new(db.clone());
-    
-    // Insert some values to cache
-    overlay.insert("key_a", "val_a");
-    overlay.insert("key_b", "val_b");
-    overlay.insert("key_c", "val_c");
-    
-    assert_eq!(overlay.get("key_a")?, Some("val_a".into()));
-    assert_eq!(overlay.get("key_b")?, Some("val_b".into()));
-    assert_eq!(overlay.get("key_c")?, Some("val_c".into()));
-    
-    // Verify they are not in sled
-    assert_eq!(db.get(b"key_a")?, None);
-    assert_eq!(db.get(b"key_b")?, None);
-    assert_eq!(db.get(b"key_c")?, None);
-    
-    // Now we write them to sled
-    overlay.flush_rollback()?;
-    
-    // Verify they are not in sled
-    assert_eq!(db.get(b"key_a")?, None);
-    assert_eq!(db.get(b"key_b")?, None);
-    assert_eq!(db.get(b"key_c")?, None);
+    assert_eq!(tree_1.get(b"key_a")?, Some(b"val_a".into()));
+    assert_eq!(tree_1.get(b"key_b")?, Some(b"val_b".into()));
+    assert_eq!(tree_1.get(b"key_c")?, Some(b"val_c".into()));
+
+    assert_eq!(tree_2.get(b"key_d")?, Some(b"val_d".into()));
+    assert_eq!(tree_2.get(b"key_e")?, Some(b"val_e".into()));
+    assert_eq!(tree_2.get(b"key_f")?, Some(b"val_f".into()));
 
     Ok(())
 }

script/research/sled_playground/src/overlay.rs

@@ -16,9 +16,9 @@
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
 
-use std::collections::BTreeMap;
+use std::collections::{btree_map::Iter, BTreeMap};
 
-use sled::IVec;
+use sled::{Batch, IVec};
 
 struct SledCache(BTreeMap<IVec, IVec>);
 
@@ -42,11 +42,15 @@ impl SledCache {
     fn remove(&mut self, key: &IVec) -> Option<IVec> {
         self.0.remove(key)
     }
+
+    fn iter(&self) -> Iter<'_, IVec, IVec> {
+        self.0.iter()
+    }
 }
 
 /// We instantiate an overlay on top of a `sled::Tree` directly.
 pub struct SledOverlay {
-    tree: sled::Tree,
+    pub tree: sled::Tree,
     cache: SledCache,
     removed: BTreeMap<IVec, IVec>,
 }
@@ -104,4 +108,18 @@ impl SledOverlay {
 
         Ok(self.cache.remove(&key.into()))
     }
+
+    pub fn aggregate(&self) -> sled::Batch {
+        let mut batch = Batch::default();
+
+        for (k, v) in self.cache.iter() {
+            batch.insert(k, v);
+        }
+
+        for k in self.removed.keys() {
+            batch.remove(k);
+        }
+
+        batch
+    }
 }