Repository: cberner/redb
Pull request #161 (2022-02-07 04:54:18)

Return removed value from .remove()


Patch:
diff --git a/src/db.rs b/src/db.rs
index 45fe0d6..b60a6af 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -275,11 +275,15 @@ impl<'a> DatabaseTransaction<'a> {
         Ok(found)
     }
 
+    /// List all the tables
+    // TODO: should return an iterator of &str, once GATs are available
     pub fn list_tables(&self) -> Result<impl Iterator<Item = String> + '_, Error> {
         self.storage
             .list_tables(TableType::Normal, self.root_page.get())
     }
 
+    /// List all the multimap tables
+    // TODO: should return an iterator of &str, once GATs are available
     pub fn list_multimap_tables(&self) -> Result<impl Iterator<Item = String> + '_, Error> {
         self.storage
             .list_tables(TableType::Multimap, self.root_page.get())
@@ -404,10 +408,14 @@ impl<'a> ReadOnlyDatabaseTransaction<'a> {
         ))
     }
 
+    /// List all the tables
+    // TODO: should return an iterator of &str, once GATs are available
     pub fn list_tables(&self) -> Result<impl Iterator<Item = String> + '_, Error> {
         self.storage.list_tables(TableType::Normal, self.root_page)
     }
 
+    /// List all the multimap tables
+    // TODO: should return an iterator of &str, once GATs are available
     pub fn list_multimap_tables(&self) -> Result<impl Iterator<Item = String> + '_, Error> {
         self.storage
             .list_tables(TableType::Multimap, self.root_page)
diff --git a/src/multimap_table.rs b/src/multimap_table.rs
index f387fd6..f7592ac 100644
--- a/src/multimap_table.rs
+++ b/src/multimap_table.rs
@@ -349,14 +349,14 @@ impl<'s, K: RedbKey + ?Sized, V: RedbKey + ?Sized> MultimapTable<'s, K, V> {
         // Tables can only be opened mutably in one location (see Error::TableAlreadyOpen),
         // and we borrow &mut self.
         let (root_page, found) = unsafe {
-            self.storage.remove::<MultimapKVPair<K, V>>(
+            self.storage.remove::<MultimapKVPair<K, V>, [u8]>(
                 &kv,
                 self.transaction_id,
                 self.table_root.get(),
             )?
         };
         self.table_root.set(root_page);
-        Ok(found)
+        Ok(found.is_some())
     }
 
     // TODO: maybe this should return all the removed values as an iter?
@@ -369,13 +369,13 @@ impl<'s, K: RedbKey + ?Sized, V: RedbKey + ?Sized> MultimapTable<'s, K, V> {
             // Tables can only be opened mutably in one location (see Error::TableAlreadyOpen),
             // and we borrow &mut self.
             let (new_root, found) = unsafe {
-                self.storage.remove::<MultimapKVPair<K, V>>(
+                self.storage.remove::<MultimapKVPair<K, V>, [u8]>(
                     &key_only,
                     self.transaction_id,
                     self.table_root.get(),
                 )?
             };
-            if !found {
+            if found.is_none() {
                 break;
             }
             self.table_root.set(new_root);
diff --git a/src/table.rs b/src/table.rs
index c2a73df..8a4792c 100644
--- a/src/table.rs
+++ b/src/table.rs
@@ -76,13 +76,12 @@ impl<'s, K: RedbKey + ?Sized, V: RedbValue + ?Sized> Table<'s, K, V> {
         Ok(guard)
     }
 
-    // TODO: maybe this should return the removed value?
-    pub fn remove(&mut self, key: &K) -> Result<bool, Error> {
+    pub fn remove(&mut self, key: &K) -> Result<Option<AccessGuard<V>>, Error> {
         // Safety: No other references to this table can exist.
         // Tables can only be opened mutably in one location (see Error::TableAlreadyOpen),
         // and we borrow &mut self.
         let (root_page, found) = unsafe {
-            self.storage.remove::<K>(
+            self.storage.remove::<K, V>(
                 key.as_bytes().as_ref(),
                 self.transaction_id,
                 self.table_root.get(),
diff --git a/src/tree_store/btree_base.rs b/src/tree_store/btree_base.rs
index ae420e2..f47b251 100644
--- a/src/tree_store/btree_base.rs
+++ b/src/tree_store/btree_base.rs
@@ -1,6 +1,6 @@
-use crate::tree_store::page_store::{Page, PageMut};
+use crate::tree_store::page_store::{Page, PageImpl, PageMut, TransactionalMemory};
 use crate::tree_store::PageNumber;
-use crate::types::RedbKey;
+use crate::types::{RedbKey, RedbValue, WithLifetime};
 use std::cmp::Ordering;
 use std::marker::PhantomData;
 use std::mem::size_of;
@@ -14,6 +14,53 @@ pub trait BtreeEntry<'a: 'b, 'b> {
     fn value(&'b self) -> &'a [u8];
 }
 
+pub struct AccessGuard<'a, V: RedbValue + ?Sized> {
+    page: PageImpl<'a>,
+    offset: usize,
+    len: usize,
+    free_uncommitted_on_drop: bool,
+    mem: &'a TransactionalMemory,
+    _value_type: PhantomData<V>,
+}
+
+impl<'a, V: RedbValue + ?Sized> AccessGuard<'a, V> {
+    // Safety: if free_on_drop is true, caller must guarantee that no other references to page exist,
+    // and that no references will be created until this AccessGuard is dropped
+    pub(in crate::tree_store) unsafe fn new(
+        page: PageImpl<'a>,
+        offset: usize,
+        len: usize,
+        free_on_drop: bool,
+        mem: &'a TransactionalMemory,
+    ) -> Self {
+        Self {
+            page,
+            offset,
+            len,
+            free_uncommitted_on_drop: free_on_drop,
+            mem,
+            _value_type: Default::default(),
+        }
+    }
+
+    // TODO: implement Deref instead of this to_value() method, when GAT is stable
+    pub fn to_value(&self) -> <<V as RedbValue>::View as WithLifetime>::Out {
+        V::from_bytes(&self.page.memory()[self.offset..(self.offset + self.len)])
+    }
+}
+
+impl<'a, V: RedbValue + ?Sized> Drop for AccessGuard<'a, V> {
+    fn drop(&mut self) {
+        if self.free_uncommitted_on_drop {
+            let page_number = self.page.get_page_number();
+            // Safety: caller to new() guaranteed that no other references to this page exist
+            unsafe {
+                assert!(self.mem.free_if_uncommitted(page_number).unwrap());
+            }
+        }
+    }
+}
+
 pub struct AccessGuardMut<'a> {
     page: PageMut<'a>,
     offset: usize,
@@ -26,6 +73,7 @@ impl<'a> AccessGuardMut<'a> {
     }
 }
 
+// TODO: this should return a RedbValue typed reference
 impl<'a> AsMut<[u8]> for AccessGuardMut<'a> {
     fn as_mut(&mut self) -> &mut [u8] {
         &mut self.page.memory_mut()[self.offset..(self.offset + self.len)]
diff --git a/src/tree_store/btree_utils.rs b/src/tree_store/btree_utils.rs
index f074d5f..1d8ff8b 100644
--- a/src/tree_store/btree_utils.rs
+++ b/src/tree_store/btree_utils.rs
@@ -1,11 +1,11 @@
 use crate::tree_store::btree_base::{
-    InternalAccessor, InternalBuilder, InternalMutator, LeafAccessor, LeafBuilder, BTREE_ORDER,
-    INTERNAL, LEAF,
+    AccessGuard, InternalAccessor, InternalBuilder, InternalMutator, LeafAccessor, LeafBuilder,
+    BTREE_ORDER, INTERNAL, LEAF,
 };
 use crate::tree_store::btree_utils::DeletionResult::{PartialInternal, PartialLeaf, Subtree};
 use crate::tree_store::page_store::{Page, PageImpl, PageNumber, TransactionalMemory};
 use crate::tree_store::{AccessGuardMut, BtreeEntry};
-use crate::types::RedbKey;
+use crate::types::{RedbKey, RedbValue};
 use crate::Error;
 use std::cmp::max;
 
@@ -102,13 +102,21 @@ pub(in crate) fn print_tree<'a>(page: PageImpl<'a>, manager: &'a TransactionalMe
 
 // Returns the new root, bool indicating if the key existed, and a list of freed pages
 // Safety: see tree_delete_helper()
-pub(in crate) unsafe fn tree_delete<'a, K: RedbKey + ?Sized>(
+#[allow(clippy::type_complexity)]
+pub(in crate) unsafe fn tree_delete<'a, K: RedbKey + ?Sized, V: RedbValue + ?Sized>(
     page: PageImpl<'a>,
     key: &[u8],
     manager: &'a TransactionalMemory,
-) -> Result<(Option<PageNumber>, bool, Vec<PageNumber>), Error> {
+) -> Result<
+    (
+        Option<PageNumber>,
+        Option<AccessGuard<'a, V>>,
+        Vec<PageNumber>,
+    ),
+    Error,
+> {
     let mut freed = vec![];
-    let (deletion_result, found) = tree_delete_helper::<K>(page, key, &mut freed, manager)?;
+    let (deletion_result, found) = tree_delete_helper::<K, V>(page, key, &mut freed, manager)?;
     let result = match deletion_result {
         DeletionResult::Subtree(page) => Some(page),
         DeletionResult::PartialLeaf(entries) => {
@@ -389,19 +397,19 @@ fn max_key(page: PageImpl, manager: &TransactionalMemory) -> Vec<u8> {
 // If key is not found, guaranteed not to modify the tree
 //
 // Safety: caller must ensure that no references to uncommitted pages in this table exist
-unsafe fn tree_delete_helper<'a, K: RedbKey + ?Sized>(
+unsafe fn tree_delete_helper<'a, K: RedbKey + ?Sized, V: RedbValue + ?Sized>(
     page: PageImpl<'a>,
     key: &[u8],
     freed: &mut Vec<PageNumber>,
     manager: &'a TransactionalMemory,
-) -> Result<(DeletionResult, bool), Error> {
+) -> Result<(DeletionResult, Option<AccessGuard<'a, V>>), Error> {
     let node_mem = page.memory();
     match node_mem[0] {
         LEAF => {
             let accessor = LeafAccessor::new(&page);
             let (position, found) = accessor.position::<K>(key);
             if !found {
-                return Ok((Subtree(page.get_page_number()), false));
+                return Ok((Subtree(page.get_page_number()), None));
             }
             if accessor.num_pairs() < BTREE_ORDER / 2 {
                 let mut partial = Vec::with_capacity(accessor.num_pairs());
@@ -412,10 +420,15 @@ unsafe fn tree_delete_helper<'a, K: RedbKey + ?Sized>(
                     let entry = accessor.entry(i).unwrap();
                     partial.push((entry.key().to_vec(), entry.value().to_vec()));
                 }
-                if !manager.free_if_uncommitted(page.get_page_number())? {
+                let uncommitted = manager.uncommitted(page.get_page_number());
+                if !uncommitted {
+                    // Won't be freed until the end of the transaction, so returning the page
+                    // in the AccessGuard below is still safe
                     freed.push(page.get_page_number());
                 }
-                Ok((PartialLeaf(partial), true))
+                let (start, end) = accessor.value_range(position).unwrap();
+                let guard = AccessGuard::new(page, start, end - start, uncommitted, manager);
+                Ok((PartialLeaf(partial), Some(guard)))
             } else {
                 let old_size = accessor.length_of_pairs(0, accessor.num_pairs());
                 let new_size =
@@ -435,26 +448,35 @@ unsafe fn tree_delete_helper<'a, K: RedbKey + ?Sized>(
                     builder.append(entry.key(), entry.value());
                 }
                 drop(builder);
-                if !manager.free_if_uncommitted(page.get_page_number())? {
+                let uncommitted = manager.uncommitted(page.get_page_number());
+                if !uncommitted {
+                    // Won't be freed until the end of the transaction, so returning the page
+                    // in the AccessGuard below is still safe
                     freed.push(page.get_page_number());
                 }
-                Ok((Subtree(new_page.get_page_number()), true))
+                let (start, end) = accessor.value_range(position).unwrap();
+                let guard = AccessGuard::new(page, start, end - start, uncommitted, manager);
+                Ok((Subtree(new_page.get_page_number()), Some(guard)))
             }
         }
         INTERNAL => {
             let accessor = InternalAccessor::new(&page);
             let original_page_number = page.get_page_number();
             let (child_index, child_page_number) = accessor.child_for_key::<K>(key);
-            let (result, found) =
-                tree_delete_helper::<K>(manager.get_page(child_page_number), key, freed, manager)?;
-            if !found {
-                return Ok((Subtree(original_page_number), false));
+            let (result, found) = tree_delete_helper::<K, V>(
+                manager.get_page(child_page_number),
+                key,
+                freed,
+                manager,
+            )?;
+            if found.is_none() {
+                return Ok((Subtree(original_page_number), None));
             }
             let final_result = match result {
                 Subtree(new_child) => {
                     if new_child == child_page_number {
                         // NO-OP. One of our descendants is uncommitted, so there was no change
-                        return Ok((Subtree(original_page_number), true));
+                        return Ok((Subtree(original_page_number), found));
                     } else if manager.uncommitted(original_page_number) {
                         drop(page);
                         // Safety: Caller guarantees there are no references to uncommitted pages,
@@ -462,7 +484,7 @@ unsafe fn tree_delete_helper<'a, K: RedbKey + ?Sized>(
                         let mut mutpage = manager.get_page_mut(original_page_number);
                         let mut mutator = InternalMutator::new(&mut mutpage);
                         mutator.write_child_page(child_index, new_child);
-                        return Ok((Subtree(original_page_number), true));
+                        return Ok((Subtree(original_page_number), found));
                     } else {
                         let mut new_page = manager.allocate(InternalBuilder::required_bytes(
                             accessor.count_children() - 1,
@@ -548,7 +570,7 @@ unsafe fn tree_delete_helper<'a, K: RedbKey + ?Sized>(
                 freed.push(original_page_number);
             }
 
-            Ok((final_result, true))
+            Ok((final_result, found))
         }
         _ => unreachable!(),
     }
@@ -976,23 +998,25 @@ unsafe fn tree_insert_helper<'a, K: RedbKey + ?Sized>(
 }
 
 // Returns the (offset, len) of the value for the queried key, if present
-pub(in crate) fn lookup_in_raw<'a, K: RedbKey + ?Sized>(
+pub(in crate) fn find_key<'a, K: RedbKey + ?Sized, V: RedbValue + ?Sized>(
     page: PageImpl<'a>,
     query: &[u8],
     manager: &'a TransactionalMemory,
-) -> Option<(PageImpl<'a>, usize, usize)> {
+) -> Option<AccessGuard<'a, V>> {
     let node_mem = page.memory();
     match node_mem[0] {
         LEAF => {
             let accessor = LeafAccessor::new(&page);
             let entry_index = accessor.find_key::<K>(query)?;
             let (start, end) = accessor.value_range(entry_index).unwrap();
-            Some((page, start, end - start))
+            // Safety: free_on_drop is false
+            let guard = unsafe { AccessGuard::new(page, start, end - start, false, manager) };
+            Some(guard)
         }
         INTERNAL => {
             let accessor = InternalAccessor::new(&page);
             let (_, child_page) = accessor.child_for_key::<K>(query);
-            return lookup_in_raw::<K>(manager.get_page(child_page), query, manager);
+            return find_key::<K, V>(manager.get_page(child_page), query, manager);
         }
         _ => unreachable!(),
     }
diff --git a/src/tree_store/mod.rs b/src/tree_store/mod.rs
index 4c0a24e..3b5202b 100644
--- a/src/tree_store/mod.rs
+++ b/src/tree_store/mod.rs
@@ -4,8 +4,8 @@ mod btree_utils;
 mod page_store;
 mod storage;
 
+pub use btree_base::AccessGuard;
 pub(crate) use btree_base::{AccessGuardMut, BtreeEntry};
 pub(crate) use btree_iters::BtreeRangeIter;
 pub(crate) use page_store::{expand_db_size, get_db_size, PageNumber};
-pub use storage::AccessGuard;
 pub(crate) use storage::{DbStats, Storage, TableType, FREED_TABLE};
diff --git a/src/tree_store/storage.rs b/src/tree_store/storage.rs
index 7071d97..f09f511 100644
--- a/src/tree_store/storage.rs
+++ b/src/tree_store/storage.rs
@@ -1,13 +1,14 @@
+use crate::tree_store::btree_base::AccessGuard;
 use crate::tree_store::btree_iters::{
     find_iter_start, find_iter_start_reversed, find_iter_unbounded_reversed,
     find_iter_unbounded_start, page_numbers_iter_start_state, AllPageNumbersBtreeIter,
 };
 use crate::tree_store::btree_utils::{
-    lookup_in_raw, make_mut_single_leaf, print_tree, tree_delete, tree_height, tree_insert,
+    find_key, make_mut_single_leaf, print_tree, tree_delete, tree_height, tree_insert,
 };
-use crate::tree_store::page_store::{Page, PageImpl, PageNumber, TransactionalMemory};
+use crate::tree_store::page_store::{PageNumber, TransactionalMemory};
 use crate::tree_store::{AccessGuardMut, BtreeEntry, BtreeRangeIter};
-use crate::types::{RedbKey, RedbValue, WithLifetime};
+use crate::types::{RedbKey, RedbValue};
 use crate::Error;
 use memmap2::MmapRaw;
 use std::borrow::Borrow;
@@ -15,7 +16,6 @@ use std::cell::{Cell, RefCell};
 use std::cmp::{max, min};
 use std::collections::{BTreeSet, Bound};
 use std::convert::TryInto;
-use std::marker::PhantomData;
 use std::mem::size_of;
 use std::ops::{RangeBounds, RangeFull, RangeTo};
 use std::panic;
@@ -310,8 +310,8 @@ impl Storage {
 
             // Safety: References into the master table are never returned to the user
             let (new_root, found) =
-                unsafe { self.remove::<str>(name.as_bytes(), transaction_id, root_page)? };
-            return Ok((new_root, found));
+                unsafe { self.remove::<str, [u8]>(name.as_bytes(), transaction_id, root_page)? };
+            return Ok((new_root, found.is_some()));
         }
 
         Ok((root_page, false))
@@ -509,7 +509,7 @@ impl Storage {
         for key in to_remove {
             // Safety: all references to the freed table above have already been dropped
             let (new_root, _) =
-                unsafe { self.remove::<[u8]>(&key, transaction_id, freed_table.table_root)? };
+                unsafe { self.remove::<[u8], [u8]>(&key, transaction_id, freed_table.table_root)? };
             freed_table.table_root = new_root;
         }
         // Safety: References into the master table are never returned to the user
@@ -650,9 +650,7 @@ impl Storage {
     ) -> Result<Option<AccessGuard<V>>, Error> {
         if let Some(handle) = root_page_handle {
             let root_page = self.mem.get_page(handle);
-            if let Some((page, offset, len)) = lookup_in_raw::<K>(root_page, key, &self.mem) {
-                return Ok(Some(AccessGuard::new(page, offset, len)));
-            }
+            return Ok(find_key::<K, V>(root_page, key, &self.mem));
         }
         Ok(None)
     }
@@ -726,44 +724,21 @@ impl Storage {
     // Returns the new root page, and a bool indicating whether the entry existed
     // Safety: caller must ensure that no references to uncommitted data in this transaction exist
     // TODO: this method could be made safe, if the transaction_id was not copy and was borrowed mut
-    pub(in crate) unsafe fn remove<K: RedbKey + ?Sized>(
+    pub(in crate) unsafe fn remove<K: RedbKey + ?Sized, V: RedbValue + ?Sized>(
         &self,
         key: &[u8],
         transaction_id: u128,
         root_handle: Option<PageNumber>,
-    ) -> Result<(Option<PageNumber>, bool), Error> {
+    ) -> Result<(Option<PageNumber>, Option<AccessGuard<V>>), Error> {
         assert_eq!(transaction_id, self.live_write_transaction.get().unwrap());
         if let Some(handle) = root_handle {
             let root_page = self.mem.get_page(handle);
-            let (new_root, found, freed) = tree_delete::<K>(root_page, key, &self.mem)?;
+            let (new_root, found, freed) = tree_delete::<K, V>(root_page, key, &self.mem)?;
             self.pending_freed_pages
                 .borrow_mut()
                 .extend_from_slice(&freed);
             return Ok((new_root, found));
         }
-        Ok((root_handle, false))
-    }
-}
-
-pub struct AccessGuard<'a, V: RedbValue + ?Sized> {
-    page: PageImpl<'a>,
-    offset: usize,
-    len: usize,
-    _value_type: PhantomData<V>,
-}
-
-impl<'a, V: RedbValue + ?Sized> AccessGuard<'a, V> {
-    fn new(page: PageImpl<'a>, offset: usize, len: usize) -> Self {
-        Self {
-            page,
-            offset,
-            len,
-            _value_type: Default::default(),
-        }
-    }
-
-    // TODO: implement Deref instead of this to_value() method, when GAT is stable
-    pub fn to_value(&self) -> <<V as RedbValue>::View as WithLifetime>::Out {
-        V::from_bytes(&self.page.memory()[self.offset..(self.offset + self.len)])
+        Ok((root_handle, None))
     }
 }


Test patch:
diff --git a/tests/basic_tests.rs b/tests/basic_tests.rs
index c8cdd99..4fdcc02 100644
--- a/tests/basic_tests.rs
+++ b/tests/basic_tests.rs
@@ -181,7 +181,11 @@ fn delete() {
 
     let write_txn = db.begin_write().unwrap();
     let mut table: Table<[u8], [u8]> = write_txn.open_table("x").unwrap();
-    table.remove(b"hello").unwrap();
+    assert_eq!(
+        b"world",
+        table.remove(b"hello").unwrap().unwrap().to_value()
+    );
+    assert!(table.remove(b"hello").unwrap().is_none());
     write_txn.commit().unwrap();
 
     let read_txn = db.begin_read().unwrap();