From 4f12612abc189bf5925584ea7393fe5f9c315c6f Mon Sep 17 00:00:00 2001 From: "hinto.janai" Date: Sun, 15 Dec 2024 12:09:39 -0500 Subject: [PATCH] update call-sites --- storage/blockchain/src/ops/alt_block/chain.rs | 25 +++-------- storage/blockchain/src/ops/output.rs | 11 ++--- storage/database/src/backend/heed/database.rs | 14 +++++++ storage/database/src/backend/tests.rs | 4 +- storage/database/src/database.rs | 8 ++++ storage/database/src/entry/entry.rs | 42 +++++++++++++------ storage/database/src/entry/occupied_entry.rs | 19 +++++++-- storage/txpool/src/service/write.rs | 7 ++-- 8 files changed, 83 insertions(+), 47 deletions(-) diff --git a/storage/blockchain/src/ops/alt_block/chain.rs b/storage/blockchain/src/ops/alt_block/chain.rs index 676fd7f..6c19ecc 100644 --- a/storage/blockchain/src/ops/alt_block/chain.rs +++ b/storage/blockchain/src/ops/alt_block/chain.rs @@ -28,10 +28,10 @@ pub fn update_alt_chain_info( Err(e) => return Err(e), }; - // try update the info if one exists for this chain. - let update = tables + tables .alt_chain_infos_mut() - .update(&alt_block_height.chain_id, |mut info| { + .entry(&alt_block_height.chain_id)? + .and_update(|info| { if info.chain_height < alt_block_height.height + 1 { // If the chain height is increasing we only need to update the chain height. info.chain_height = alt_block_height.height + 1; @@ -43,25 +43,12 @@ pub fn update_alt_chain_info( } info.chain_height = alt_block_height.height + 1; - Some(info) - }); - - match update { - Ok(()) => return Ok(()), - Err(RuntimeError::KeyNotFound) => (), - Err(e) => return Err(e), - } - - // If one doesn't already exist add it. - - tables.alt_chain_infos_mut().put( - &alt_block_height.chain_id, - &AltChainInfo { + })? + .or_insert_with(|| AltChainInfo { parent_chain: parent_chain.into(), common_ancestor_height: alt_block_height.height.checked_sub(1).unwrap(), chain_height: alt_block_height.height + 1, - }, - ) + }) } /// Get the height history of an alt-chain in reverse chronological order. diff --git a/storage/blockchain/src/ops/output.rs b/storage/blockchain/src/ops/output.rs index 96d94bb..ff8d14b 100644 --- a/storage/blockchain/src/ops/output.rs +++ b/storage/blockchain/src/ops/output.rs @@ -67,14 +67,9 @@ pub fn remove_output( // `btree_map::Entry`-like API, fix `trait DatabaseRw`. tables .num_outputs_mut() - .update(&pre_rct_output_id.amount, |num_outputs| { - // INVARIANT: Should never be 0. - if num_outputs == 1 { - None - } else { - Some(num_outputs - 1) - } - })?; + .entry(&pre_rct_output_id.amount)? + .and_remove(|num_outputs| *num_outputs == 1)? + .and_update(|num_outputs| *num_outputs -= 1)?; // Delete the output data itself. tables.outputs_mut().delete(pre_rct_output_id) diff --git a/storage/database/src/backend/heed/database.rs b/storage/database/src/backend/heed/database.rs index 1740c3f..41ac149 100644 --- a/storage/database/src/backend/heed/database.rs +++ b/storage/database/src/backend/heed/database.rs @@ -189,6 +189,20 @@ impl DatabaseRw for HeedTableRw<'_, '_, T> { Ok(()) } + #[inline] + fn take(&mut self, key: &T::Key) -> DbResult { + // LMDB/heed does not return the value on deletion. + // So, fetch it first - then delete. + let value = get::(&self.db, &self.tx_rw.borrow(), key)?; + match self.db.delete(&mut self.tx_rw.borrow_mut(), key) { + Ok(true) => Ok(value), + Err(e) => Err(e.into()), + // We just `get()`'ed the value - it is + // incorrect for it to suddenly not exist. + Ok(false) => unreachable!(), + } + } + #[inline] fn pop_first(&mut self) -> DbResult<(T::Key, T::Value)> { let tx_rw = &mut self.tx_rw.borrow_mut(); diff --git a/storage/database/src/backend/tests.rs b/storage/database/src/backend/tests.rs index 24713d0..de1275b 100644 --- a/storage/database/src/backend/tests.rs +++ b/storage/database/src/backend/tests.rs @@ -293,11 +293,11 @@ fn db_read_write() { assert_value(value); } - // Assert `Entry` returns the correct value. + // Assert `take()` works. { let mut key = KEY; key += 1; - let value = table.entry(&key).unwrap().and_remove().unwrap(); + let value = table.take(&key).unwrap(); assert_eq!(value, VALUE); let get = table.get(&KEY); diff --git a/storage/database/src/database.rs b/storage/database/src/database.rs index 06a4b47..4dec239 100644 --- a/storage/database/src/database.rs +++ b/storage/database/src/database.rs @@ -171,6 +171,14 @@ pub trait DatabaseRw: DatabaseRo + Sized { /// This will never [`RuntimeError::KeyExists`]. fn delete(&mut self, key: &T::Key) -> DbResult<()>; + /// Delete and return a key-value pair in the database. + /// + /// This is the same as [`DatabaseRw::delete`], however, + /// it will serialize the `T::Value` and return it. + /// + #[doc = doc_database!()] + fn take(&mut self, key: &T::Key) -> DbResult; + /// Removes and returns the first `(key, value)` pair in the database. /// #[doc = doc_database!()] diff --git a/storage/database/src/entry/entry.rs b/storage/database/src/entry/entry.rs index ab0aa1e..4aa9680 100644 --- a/storage/database/src/entry/entry.rs +++ b/storage/database/src/entry/entry.rs @@ -29,6 +29,16 @@ where T: Table, D: DatabaseRw, { + /// TODO + pub const fn is_occupied(&self) -> bool { + matches!(self, Self::Occupied(_)) + } + + /// TODO + pub const fn is_vacant(&self) -> bool { + matches!(self, Self::Vacant(_)) + } + /// Ensures a value is in the entry by inserting the `default` if empty. /// /// This only inserts if the entry is [`VacantEntry`]. @@ -44,24 +54,24 @@ where /// This only inserts if the entry is [`VacantEntry`]. pub fn or_insert_with(self, default: F) -> DbResult<()> where - F: FnOnce() -> &'a T::Value, + F: FnOnce() -> T::Value, { match self { Self::Occupied(_) => Ok(()), - Self::Vacant(entry) => entry.insert(default()), + Self::Vacant(entry) => entry.insert(&default()), } } /// Same as [`Self::or_insert_with`] but gives access to the key. pub fn or_insert_with_key(self, default: F) -> DbResult<()> where - F: FnOnce(&'a T::Key) -> &'a T::Value, + F: FnOnce(&'a T::Key) -> T::Value, { match self { Self::Occupied(_) => Ok(()), Self::Vacant(entry) => { let key = entry.key; - entry.insert(default(key)) + entry.insert(&default(key)) } } } @@ -96,15 +106,23 @@ where } } - /// [`OccupiedEntry::remove`] the value if it already exists, else do nothing. + /// Conditionally [`OccupiedEntry::remove`] the value if it already exists. /// - /// # Errors - /// This returns [`RuntimeError::KeyNotFound`] if the entry is [`VacantEntry`]. - pub fn and_remove(self) -> DbResult { - match self { - Self::Occupied(entry) => entry.remove(), - Self::Vacant(_) => Err(RuntimeError::KeyNotFound), - } + /// This functions does nothing if the entry is [`VacantEntry`]. + pub fn and_remove(self, f: F) -> DbResult + where + F: FnOnce(&T::Value) -> bool, + { + Ok(match self { + Self::Occupied(entry) => { + if f(&entry.value) { + entry.remove()?.0 + } else { + Self::Occupied(entry) + } + } + Self::Vacant(entry) => Self::Vacant(entry), + }) } /// [`OccupiedEntry::update`] the value if it already exists diff --git a/storage/database/src/entry/occupied_entry.rs b/storage/database/src/entry/occupied_entry.rs index af3f756..44d5b5d 100644 --- a/storage/database/src/entry/occupied_entry.rs +++ b/storage/database/src/entry/occupied_entry.rs @@ -2,6 +2,8 @@ use crate::{DatabaseRw, DbResult, Table}; +use super::{Entry, VacantEntry}; + /// A view into an occupied entry in a [`DatabaseRw`]. It is part of [`crate::entry::Entry`]. pub struct OccupiedEntry<'a, T, D> where @@ -13,7 +15,7 @@ where pub(crate) value: T::Value, } -impl OccupiedEntry<'_, T, D> +impl<'a, T, D> OccupiedEntry<'a, T, D> where T: Table, D: DatabaseRw, @@ -52,7 +54,18 @@ where } /// Remove this entry. - pub fn remove(self) -> DbResult { - DatabaseRw::delete(self.db, self.key).map(|()| Ok(self.value))? + /// + /// The returns values are: + /// - An [`Entry::VacantEntry`] + /// - The value that was removed + pub fn remove(self) -> DbResult<(Entry<'a, T, D>, T::Value)> { + DatabaseRw::delete(self.db, self.key)?; + Ok(( + Entry::Vacant(VacantEntry { + db: self.db, + key: self.key, + }), + self.value, + )) } } diff --git a/storage/txpool/src/service/write.rs b/storage/txpool/src/service/write.rs index 23c5a8a..4b06bd4 100644 --- a/storage/txpool/src/service/write.rs +++ b/storage/txpool/src/service/write.rs @@ -114,10 +114,11 @@ fn promote(env: &ConcreteEnv, tx_hash: &TransactionHash) -> DbResult(&tx_rw)?; - tx_infos.update(tx_hash, |mut info| { + tx_infos.entry(tx_hash)?.and_update(|info| { info.flags.remove(TxStateFlags::STATE_STEM); - Some(info) - }) + })?; + + Ok(()) }; if let Err(e) = res() {