Repository: cberner/redb
Pull request #7 (2021-07-06 18:25:59)

Linear-scan implementation with mmap
Storage is now persistent, but insert()/get() are both O(n)


Patch:
diff --git a/Cargo.toml b/Cargo.toml
index edd0085..eb1225b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -15,6 +15,7 @@ crate-type = ["cdylib", "rlib"]
 [dependencies]
 rand = "0.8"
 chrono = "0.4.6"
+memmap2 = "0.3.0"
 pyo3 = {version = "0.13", features=["extension-module", "abi3-py36"], optional = true }
 
 [dev-dependencies]
diff --git a/benches/lmdb_benchmark.rs b/benches/lmdb_benchmark.rs
index 51db14a..0d08fc3 100644
--- a/benches/lmdb_benchmark.rs
+++ b/benches/lmdb_benchmark.rs
@@ -192,7 +192,7 @@ fn redb_bench(path: &Path) {
                 let (key, value) = &pairs[*i % pairs.len()];
                 let mut mut_key = key.clone();
                 mut_key.extend_from_slice(&i.to_be_bytes());
-                let result: &[u8] = txn.get(&mut_key).unwrap().unwrap();
+                let result: &[u8] = &txn.get(&mut_key).unwrap().unwrap();
                 checksum += result[0] as u64;
                 expected_checksum += value[0] as u64;
             }
diff --git a/deny.toml b/deny.toml
index 3f2384b..3260717 100644
--- a/deny.toml
+++ b/deny.toml
@@ -72,7 +72,7 @@ unlicensed = "deny"
 allow = [
     "MIT",
 #    "BSD-2-Clause",
-    "BSD-3-Clause",
+#    "BSD-3-Clause",
     "Apache-2.0",
     #"Apache-2.0 WITH LLVM-exception",
 ]
diff --git a/src/db.rs b/src/db.rs
index 3ea94ca..0471fd0 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -1,20 +1,41 @@
+use crate::storage::Storage;
 use crate::table::Table;
 use crate::Error;
+use memmap2::MmapMut;
+use std::fs::OpenOptions;
 use std::path::Path;
 
-pub struct Database {}
+pub struct Database {
+    storage: Storage,
+}
 
 impl Database {
+    /// Opens the specified file as a redb database.
+    /// * if the file does not exist, or is an empty file, a new database will be initialized in it
+    /// * if the file is a valid redb database, it will be opened
+    /// * otherwise this function will return an error
     ///
     /// # Safety
     ///
     /// The file referenced by `path` must only be concurrently modified by compatible versions
     /// of redb
-    pub unsafe fn open(_path: &Path) -> Result<Database, Error> {
-        Ok(Database {})
+    pub unsafe fn open(path: &Path) -> Result<Database, Error> {
+        let file = OpenOptions::new()
+            .read(true)
+            .write(true)
+            .create(true)
+            .open(path)?;
+
+        // TODO: what value should we use here?
+        file.set_len(4 * 1024 * 1024 * 1024)?;
+
+        let mmap = MmapMut::map_mut(&file)?;
+        let storage = Storage::new(mmap);
+        storage.initialize()?;
+        Ok(Database { storage })
     }
 
     pub fn open_table(&self, _name: &str) -> Result<Table, Error> {
-        Table::new()
+        Table::new(&self.storage)
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
index 14599a8..bc9b2ea 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -2,11 +2,13 @@ mod db;
 mod error;
 #[cfg(feature = "python")]
 mod python;
+mod storage;
 mod table;
 mod transactions;
 
 pub use db::Database;
 pub use error::Error;
+pub use storage::AccessGuard;
 pub use table::Table;
 pub use transactions::{ReadOnlyTransaction, WriteTransaction};
 
diff --git a/src/storage.rs b/src/storage.rs
new file mode 100644
index 0000000..fa48196
--- /dev/null
+++ b/src/storage.rs
@@ -0,0 +1,127 @@
+use crate::Error;
+use memmap2::MmapMut;
+use std::cell::{Ref, RefCell};
+use std::convert::TryInto;
+use std::ops::Deref;
+
+const MAGICNUMBER: [u8; 4] = [b'r', b'e', b'd', b'b'];
+const ENTRIES_OFFSET: usize = MAGICNUMBER.len();
+const DATA_OFFSET: usize = ENTRIES_OFFSET + 8;
+
+pub(in crate) struct Storage {
+    mmap: RefCell<MmapMut>,
+}
+
+impl Storage {
+    pub(in crate) fn new(mmap: MmapMut) -> Storage {
+        Storage {
+            mmap: RefCell::new(mmap),
+        }
+    }
+
+    pub(in crate) fn initialize(&self) -> Result<(), Error> {
+        let mut mmap = self.mmap.borrow_mut();
+        if mmap[0..MAGICNUMBER.len()] == MAGICNUMBER {
+            return Ok(());
+        }
+        mmap[ENTRIES_OFFSET..(ENTRIES_OFFSET + 8)].copy_from_slice(&0u64.to_be_bytes());
+        mmap.flush()?;
+        // Write the magic number only after the data structure is initialized and written to disk
+        // to ensure that it's crash safe
+        mmap[0..MAGICNUMBER.len()].copy_from_slice(&MAGICNUMBER);
+        mmap.flush()?;
+
+        Ok(())
+    }
+
+    pub(in crate) fn append(&self, key: &[u8], value: &[u8]) -> Result<(), Error> {
+        let mut mmap = self.mmap.borrow_mut();
+        let mut entries = u64::from_be_bytes(
+            mmap[ENTRIES_OFFSET..(ENTRIES_OFFSET + 8)]
+                .try_into()
+                .unwrap(),
+        ) as usize;
+
+        let mut index = DATA_OFFSET;
+        let mut entry = 0;
+        while entry < entries {
+            let key_len = u64::from_be_bytes(mmap[index..(index + 8)].try_into().unwrap()) as usize;
+            index += 8;
+            if key == &mmap[index..(index + key_len)] {
+                todo!("support overwriting existing keys");
+            }
+            index += key_len;
+            let value_len =
+                u64::from_be_bytes(mmap[index..(index + 8)].try_into().unwrap()) as usize;
+            index += 8 + value_len;
+            entry += 1;
+        }
+
+        // Does not exist, append the new key & value
+        mmap[index..(index + 8)].copy_from_slice(&(key.len() as u64).to_be_bytes());
+        index += 8;
+        mmap[index..(index + key.len())].copy_from_slice(key);
+        index += key.len();
+        mmap[index..(index + 8)].copy_from_slice(&(value.len() as u64).to_be_bytes());
+        index += 8;
+        mmap[index..(index + value.len())].copy_from_slice(value);
+
+        entries += 1;
+        mmap[ENTRIES_OFFSET..(ENTRIES_OFFSET + 8)].copy_from_slice(&entries.to_be_bytes());
+        Ok(())
+    }
+
+    pub(in crate) fn fsync(&self) -> Result<(), Error> {
+        self.mmap.borrow().flush()?;
+        Ok(())
+    }
+
+    pub(in crate) fn get(&self, key: &[u8]) -> Result<Option<AccessGuard>, Error> {
+        let mmap = self.mmap.borrow();
+
+        let entries = u64::from_be_bytes(
+            mmap[ENTRIES_OFFSET..(ENTRIES_OFFSET + 8)]
+                .try_into()
+                .unwrap(),
+        ) as usize;
+
+        let mut index = DATA_OFFSET;
+        let mut entry = 0;
+        while entry < entries {
+            let key_len = u64::from_be_bytes(mmap[index..(index + 8)].try_into().unwrap()) as usize;
+            index += 8;
+            let value_len = u64::from_be_bytes(
+                mmap[(index + key_len)..(index + key_len + 8)]
+                    .try_into()
+                    .unwrap(),
+            ) as usize;
+            if key == &mmap[index..(index + key_len)] {
+                index += key_len + 8;
+                let guard = AccessGuard {
+                    mmap_ref: mmap,
+                    offset: index,
+                    len: value_len,
+                };
+                return Ok(Some(guard));
+            }
+            index += key_len + 8 + value_len;
+            entry += 1;
+        }
+
+        Ok(None)
+    }
+}
+
+pub struct AccessGuard<'mmap> {
+    mmap_ref: Ref<'mmap, MmapMut>,
+    offset: usize,
+    len: usize,
+}
+
+impl<'mmap> Deref for AccessGuard<'mmap> {
+    type Target = [u8];
+
+    fn deref(&self) -> &[u8] {
+        &self.mmap_ref[self.offset..(self.offset + self.len)]
+    }
+}
diff --git a/src/table.rs b/src/table.rs
index 34230ba..6589ebd 100644
--- a/src/table.rs
+++ b/src/table.rs
@@ -1,25 +1,22 @@
 use crate::error::Error;
+use crate::storage::Storage;
 use crate::transactions::WriteTransaction;
 use crate::ReadOnlyTransaction;
-use std::cell::RefCell;
-use std::collections::HashMap;
 
-pub struct Table {
-    data: RefCell<HashMap<Vec<u8>, Vec<u8>>>,
+pub struct Table<'mmap> {
+    storage: &'mmap Storage,
 }
 
-impl Table {
-    pub(in crate) fn new() -> Result<Table, Error> {
-        Ok(Table {
-            data: RefCell::new(HashMap::new()),
-        })
+impl<'mmap> Table<'mmap> {
+    pub(in crate) fn new(storage: &'mmap Storage) -> Result<Table<'mmap>, Error> {
+        Ok(Table { storage })
     }
 
     pub fn begin_write(&mut self) -> Result<WriteTransaction, Error> {
-        Ok(WriteTransaction::new(&self.data))
+        Ok(WriteTransaction::new(self.storage))
     }
 
     pub fn read_transaction(&self) -> Result<ReadOnlyTransaction, Error> {
-        Ok(ReadOnlyTransaction::new(self.data.borrow().clone()))
+        Ok(ReadOnlyTransaction::new(self.storage))
     }
 }
diff --git a/src/transactions.rs b/src/transactions.rs
index c9534c7..ee25844 100644
--- a/src/transactions.rs
+++ b/src/transactions.rs
@@ -1,16 +1,18 @@
 use crate::error::Error;
-use std::cell::RefCell;
+use crate::storage::{AccessGuard, Storage};
 use std::collections::HashMap;
 
-pub struct WriteTransaction<'a> {
-    table: &'a RefCell<HashMap<Vec<u8>, Vec<u8>>>,
+pub struct WriteTransaction<'mmap> {
+    storage: &'mmap Storage,
     data: HashMap<Vec<u8>, Vec<u8>>,
 }
 
-impl<'a> WriteTransaction<'a> {
-    pub(in crate) fn new(table: &'a RefCell<HashMap<Vec<u8>, Vec<u8>>>) -> WriteTransaction<'a> {
-        let data = table.borrow().clone();
-        WriteTransaction { table, data }
+impl<'mmap> WriteTransaction<'mmap> {
+    pub(in crate) fn new(storage: &'mmap Storage) -> WriteTransaction<'mmap> {
+        WriteTransaction {
+            storage,
+            data: HashMap::new(),
+        }
     }
 
     pub fn insert(&mut self, key: &[u8], value: &[u8]) -> Result<(), Error> {
@@ -19,21 +21,24 @@ impl<'a> WriteTransaction<'a> {
     }
 
     pub fn commit(self) -> Result<(), Error> {
-        self.table.replace(self.data);
+        for (key, value) in self.data.iter() {
+            self.storage.append(key, value)?;
+        }
+        self.storage.fsync()?;
         Ok(())
     }
 }
 
-pub struct ReadOnlyTransaction {
-    data: HashMap<Vec<u8>, Vec<u8>>,
+pub struct ReadOnlyTransaction<'mmap> {
+    storage: &'mmap Storage,
 }
 
-impl ReadOnlyTransaction {
-    pub(in crate) fn new(data: HashMap<Vec<u8>, Vec<u8>>) -> ReadOnlyTransaction {
-        ReadOnlyTransaction { data }
+impl<'mmap> ReadOnlyTransaction<'mmap> {
+    pub(in crate) fn new(storage: &'mmap Storage) -> ReadOnlyTransaction<'mmap> {
+        ReadOnlyTransaction { storage }
     }
 
-    pub fn get(&self, key: &[u8]) -> Result<Option<&[u8]>, Error> {
-        Ok(self.data.get(key).map(|x| x.as_slice()))
+    pub fn get(&self, key: &[u8]) -> Result<Option<AccessGuard<'mmap>>, Error> {
+        self.storage.get(key)
     }
 }


Test patch:
diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs
new file mode 100644
index 0000000..bc0e300
--- /dev/null
+++ b/tests/integration_tests.rs
@@ -0,0 +1,56 @@
+use tempfile::NamedTempFile;
+
+use rand::prelude::SliceRandom;
+use rand::Rng;
+use redb::Database;
+
+const ELEMENTS: usize = 100;
+
+/// Returns pairs of key, value
+fn gen_data(count: usize, key_size: usize, value_size: usize) -> Vec<(Vec<u8>, Vec<u8>)> {
+    let mut pairs = vec![];
+
+    for _ in 0..count {
+        let key: Vec<u8> = (0..key_size).map(|_| rand::thread_rng().gen()).collect();
+        let value: Vec<u8> = (0..value_size).map(|_| rand::thread_rng().gen()).collect();
+        pairs.push((key, value));
+    }
+
+    pairs
+}
+
+#[test]
+fn persistence() {
+    let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();
+
+    let db = unsafe { Database::open(tmpfile.path()).unwrap() };
+    let mut table = db.open_table("").unwrap();
+
+    let pairs = gen_data(100, 16, 20);
+
+    let mut txn = table.begin_write().unwrap();
+    {
+        for i in 0..ELEMENTS {
+            let (key, value) = &pairs[i % pairs.len()];
+            txn.insert(key, value).unwrap();
+        }
+    }
+    txn.commit().unwrap();
+
+    drop(table);
+    drop(db);
+    let db = unsafe { Database::open(tmpfile.path()).unwrap() };
+    let table = db.open_table("").unwrap();
+
+    let mut key_order: Vec<usize> = (0..ELEMENTS).collect();
+    key_order.shuffle(&mut rand::thread_rng());
+
+    let txn = table.read_transaction().unwrap();
+    {
+        for i in &key_order {
+            let (key, value) = &pairs[*i % pairs.len()];
+            let result: &[u8] = &txn.get(key).unwrap().unwrap();
+            assert_eq!(result, value);
+        }
+    }
+}