script/research/sled_playground: overlay2 added, wrapping everything in there

script/research/sled_playground/src/main.rs

@@ -21,6 +21,9 @@ use sled::{transaction::ConflictableTransactionError, Config, Transactional};
 pub mod overlay;
 use overlay::SledOverlay;
 
+pub mod overlay2;
+use overlay2::SledOverlay2;
+
 const TREE_1: &str = "_tree1";
 const TREE_2: &str = "_tree2";
 
@@ -87,5 +90,53 @@ fn main() -> Result<(), sled::Error> {
     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()));
 
+    // Testing overlay2
+    // Initialize database overlay
+    let config = Config::new().temporary(true);
+    let db = config.open()?;
+    let mut overlay = SledOverlay2::new(&db, &[&TREE_1, &TREE_2])?;
+    // We keep seperate trees for validation
+    let tree_1 = db.open_tree(TREE_1)?;
+    let tree_2 = db.open_tree(TREE_2)?;
+
+    // Insert some values to the overlays
+    overlay.insert(TREE_1, b"key_a", b"val_a")?;
+    overlay.insert(TREE_1, b"key_b", b"val_b")?;
+    overlay.insert(TREE_1, b"key_c", b"val_c")?;
+
+    overlay.insert(TREE_2, b"key_d", b"val_d")?;
+    overlay.insert(TREE_2, b"key_e", b"val_e")?;
+    overlay.insert(TREE_2, b"key_f", b"val_f")?;
+
+    // Verify they are in the overlay
+    assert_eq!(overlay.get(TREE_1, b"key_a")?, Some(b"val_a".into()));
+    assert_eq!(overlay.get(TREE_1, b"key_b")?, Some(b"val_b".into()));
+    assert_eq!(overlay.get(TREE_1, b"key_c")?, Some(b"val_c".into()));
+
+    assert_eq!(overlay.get(TREE_2, b"key_d")?, Some(b"val_d".into()));
+    assert_eq!(overlay.get(TREE_2, b"key_e")?, Some(b"val_e".into()));
+    assert_eq!(overlay.get(TREE_2, b"key_f")?, Some(b"val_f".into()));
+
+    // Verify they are not in sled
+    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);
+
+    // Now execute all tree baches in the overlay
+    assert_eq!(overlay.execute(), Ok(()));
+
+    // Verify sled contains keys
+    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/overlay2.rs

@@ -0,0 +1,219 @@
+/* This file is part of DarkFi (https://dark.fi)
+ *
+ * Copyright (C) 2020-2023 Dyne.org foundation
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as
+ * published by the Free Software Foundation, either version 3 of the
+ * License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+use std::collections::{btree_map::Iter, BTreeMap};
+
+use sled::{
+    transaction::{ConflictableTransactionError, TransactionError},
+    Batch, IVec, Transactional,
+};
+
+#[derive(Debug, PartialEq)]
+struct CacheNotFoundError;
+
+struct TreeCache(BTreeMap<IVec, IVec>);
+
+impl TreeCache {
+    fn new() -> Self {
+        Self(BTreeMap::new())
+    }
+
+    fn contains_key(&self, key: &IVec) -> bool {
+        self.0.contains_key(key)
+    }
+
+    fn get(&self, key: &IVec) -> Option<IVec> {
+        self.0.get(key).cloned()
+    }
+
+    fn insert(&mut self, key: IVec, value: IVec) -> Option<IVec> {
+        self.0.insert(key, value)
+    }
+
+    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 TreeOverlay {
+    tree: sled::Tree,
+    cache: TreeCache,
+    removed: BTreeMap<IVec, IVec>,
+}
+
+impl TreeOverlay {
+    pub fn new(db: &sled::Tree) -> Self {
+        Self { tree: db.clone(), cache: TreeCache::new(), removed: BTreeMap::new() }
+    }
+
+    pub fn contains_key(&self, key: &[u8]) -> Result<bool, sled::Error> {
+        if self.removed.contains_key::<IVec>(&key.into()) {
+            return Ok(false)
+        }
+
+        if self.cache.contains_key(&key.into()) || self.tree.contains_key(key)? {
+            return Ok(true)
+        }
+
+        Ok(false)
+    }
+
+    pub fn get(&self, key: &[u8]) -> Result<Option<IVec>, sled::Error> {
+        if self.removed.contains_key::<IVec>(&key.into()) {
+            return Ok(None)
+        }
+
+        if let Some(v) = self.cache.get(&key.into()) {
+            return Ok(Some(v.clone()))
+        }
+
+        self.tree.get(key)
+    }
+
+    pub fn insert(&mut self, key: &[u8], value: &[u8]) -> Result<Option<IVec>, sled::Error> {
+        let mut prev: Option<IVec> = self.cache.insert(key.into(), value.into());
+
+        if self.removed.contains_key::<IVec>(&key.into()) {
+            self.removed.remove(key);
+            return Ok(None)
+        }
+
+        if prev.is_none() {
+            prev = self.tree.get::<IVec>(key.into())?;
+        }
+
+        Ok(prev)
+    }
+
+    pub fn remove(&mut self, key: &[u8]) -> Result<Option<IVec>, sled::Error> {
+        if self.removed.contains_key::<IVec>(&key.into()) {
+            return Ok(None)
+        }
+
+        self.removed.insert(key.into(), vec![].into());
+
+        Ok(self.cache.remove(&key.into()))
+    }
+
+    pub fn aggregate(&self) -> Option<sled::Batch> {
+        if self.cache.0.is_empty() && self.removed.is_empty() {
+            return None
+        }
+        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);
+        }
+
+        Some(batch)
+    }
+}
+
+/// We instantiate overlays on top of requested
+/// sled tree keys.
+pub struct SledOverlay2 {
+    db: sled::Db,
+    trees: BTreeMap<IVec, sled::Tree>,
+    caches: BTreeMap<IVec, TreeOverlay>,
+}
+
+impl SledOverlay2 {
+    pub fn new(db: &sled::Db, trees_keys: &[&str]) -> Result<Self, sled::Error> {
+        let mut trees = BTreeMap::new();
+        let mut caches = BTreeMap::new();
+        for tree_key in trees_keys {
+            let tree = db.open_tree(tree_key)?;
+            let cache = TreeOverlay::new(&tree);
+            trees.insert(tree_key.clone().into(), tree.clone());
+            caches.insert(tree_key.clone().into(), cache);
+        }
+        Ok(Self { db: db.clone(), trees, caches })
+    }
+
+    fn get_cache(&self, tree_key: IVec) -> Result<&TreeOverlay, sled::Error> {
+        if let Some(v) = self.caches.get(&tree_key) {
+            return Ok(v)
+        }
+        Err(sled::Error::CollectionNotFound(tree_key.clone()))
+    }
+
+    fn get_cache_mut(&mut self, tree_key: IVec) -> Result<&mut TreeOverlay, sled::Error> {
+        if let Some(v) = self.caches.get_mut(&tree_key) {
+            return Ok(v)
+        }
+        Err(sled::Error::CollectionNotFound(tree_key.clone()))
+    }
+
+    pub fn contains_key(&self, tree_key: &str, key: &[u8]) -> Result<bool, sled::Error> {
+        let cache = self.get_cache(tree_key.clone().into())?;
+        cache.contains_key(key)
+    }
+
+    pub fn get(&self, tree_key: &str, key: &[u8]) -> Result<Option<IVec>, sled::Error> {
+        let cache = self.get_cache(tree_key.clone().into())?;
+        cache.get(key)
+    }
+
+    pub fn insert(
+        &mut self,
+        tree_key: &str,
+        key: &[u8],
+        value: &[u8],
+    ) -> Result<Option<IVec>, sled::Error> {
+        let cache = self.get_cache_mut(tree_key.clone().into())?;
+        cache.insert(key, value)
+    }
+
+    pub fn remove(&mut self, tree_key: &str, key: &[u8]) -> Result<Option<IVec>, sled::Error> {
+        let cache = self.get_cache_mut(tree_key.clone().into())?;
+        cache.remove(key)
+    }
+
+    pub fn execute(&mut self) -> Result<(), TransactionError<sled::Error>> {
+        let mut trees = vec![];
+        let mut batches = vec![];
+        for (key, tree) in &self.trees {
+            let cache = self.get_cache(key.clone())?;
+            if let Some(batch) = cache.aggregate() {
+                trees.push(tree);
+                batches.push(batch);
+            }
+        }
+
+        trees.transaction(|trees| {
+            for (index, tree) in trees.iter().enumerate() {
+                tree.apply_batch(&batches[index])?;
+            }
+
+            Ok::<(), ConflictableTransactionError<sled::Error>>(())
+        })?;
+
+        self.db.flush()?;
+
+        Ok(())
+    }
+}