diff --git a/core/store/src/trie/insert_delete.rs b/core/store/src/trie/insert_delete.rs index 64ec0c45712..a80ad46dfca 100644 --- a/core/store/src/trie/insert_delete.rs +++ b/core/store/src/trie/insert_delete.rs @@ -1,4 +1,5 @@ use super::TrieRefcountDeltaMap; +use crate::trie::mem::updating::GenericTrieUpdate; use crate::trie::nibble_slice::NibbleSlice; use crate::trie::{ Children, NodeHandle, RawTrieNode, RawTrieNodeWithSize, StorageHandle, StorageValueHandle, @@ -11,7 +12,7 @@ use near_primitives::state::{ValueRef, ValueUpdate}; pub(crate) struct NodesStorage<'a> { nodes: Vec>, - values: Vec>>, + pub(crate) values: Vec>>, pub(crate) refcount_changes: TrieRefcountDeltaMap, pub(crate) trie: &'a Trie, } @@ -58,19 +59,6 @@ impl<'a> NodesStorage<'a> { StorageHandle(self.nodes.len() - 1) } - pub(crate) fn store_value(&mut self, value: ValueUpdate) -> StorageValueHandle { - let ValueUpdate::MemtrieAndDisk(value) = value else { - unimplemented!( - "NodesStorage for Trie doesn't support value {value:?} \ - because disk updates must be generated." - ); - }; - - let value_len = value.len(); - self.values.push(Some(value)); - StorageValueHandle(self.values.len() - 1, value_len) - } - pub(crate) fn value_ref(&self, handle: StorageValueHandle) -> &[u8] { self.values .get(handle.0) @@ -111,11 +99,8 @@ impl Trie { let children_memory_usage = memory_usage - node.memory_usage_direct(memory); match node { TrieNode::Empty => { - let value_handle = memory.store_value(value); - let leaf_node = TrieNode::Leaf( - partial.encoded(true).into_vec(), - ValueHandle::InMemory(value_handle), - ); + let value_handle = memory.generic_store_value(value); + let leaf_node = TrieNode::Leaf(partial.encoded(true).into_vec(), value_handle); let memory_usage = leaf_node.memory_usage_direct(memory); memory.store_at(handle, TrieNodeWithSize { node: leaf_node, memory_usage }); break; @@ -123,12 +108,11 @@ impl Trie { TrieNode::Branch(mut children, existing_value) => { // If the key ends here, store the value in branch's value. if partial.is_empty() { - if let Some(value) = &existing_value { - self.delete_value(memory, value)?; + if let Some(value) = existing_value { + memory.generic_delete_value(value)?; } - let value_handle = memory.store_value(value); - let new_node = - TrieNode::Branch(children, Some(ValueHandle::InMemory(value_handle))); + let value_handle = memory.generic_store_value(value); + let new_node = TrieNode::Branch(children, Some(value_handle)); let new_memory_usage = children_memory_usage + new_node.memory_usage_direct(memory); memory.store_at(handle, TrieNodeWithSize::new(new_node, new_memory_usage)); @@ -160,9 +144,9 @@ impl Trie { let common_prefix = partial.common_prefix(&existing_key); if common_prefix == existing_key.len() && common_prefix == partial.len() { // Equivalent leaf. - self.delete_value(memory, &existing_value)?; - let value_handle = memory.store_value(value); - let node = TrieNode::Leaf(key, ValueHandle::InMemory(value_handle)); + memory.generic_delete_value(existing_value)?; + let value_handle = memory.generic_store_value(value); + let node = TrieNode::Leaf(key, value_handle); let memory_usage = node.memory_usage_direct(memory); memory.store_at(handle, TrieNodeWithSize { node, memory_usage }); break; @@ -342,7 +326,7 @@ impl Trie { } TrieNode::Leaf(key, value) => { if NibbleSlice::from_encoded(&key).0 == partial { - self.delete_value(memory, &value)?; + memory.generic_delete_value(value)?; memory.store_at(handle, TrieNodeWithSize::empty()); break; } else { @@ -367,7 +351,7 @@ impl Trie { key_deleted = false; break; } - self.delete_value(memory, &value.unwrap())?; + memory.generic_delete_value(value.unwrap())?; Trie::calc_memory_usage_and_store( memory, handle, diff --git a/core/store/src/trie/mem/loading.rs b/core/store/src/trie/mem/loading.rs index 35b2b3c4cda..376a5e48caa 100644 --- a/core/store/src/trie/mem/loading.rs +++ b/core/store/src/trie/mem/loading.rs @@ -158,9 +158,9 @@ pub fn load_trie_from_flat_state_and_delta( for (key, value) in changes.0 { match value { Some(value) => { - trie_update.insert_memtrie_only(&key, value); + trie_update.insert_memtrie_only(&key, value)?; } - None => trie_update.delete(&key), + None => trie_update.delete(&key)?, }; } diff --git a/core/store/src/trie/mem/resharding.rs b/core/store/src/trie/mem/resharding.rs index e29047106be..452bd2ab5ef 100644 --- a/core/store/src/trie/mem/resharding.rs +++ b/core/store/src/trie/mem/resharding.rs @@ -349,7 +349,7 @@ mod tests { let mut memtries = MemTries::new(ShardUId::single_shard()); let mut update = memtries.update(Trie::EMPTY_ROOT, false).unwrap(); for (key, value) in initial_entries { - update.insert(&key, value); + update.insert(&key, value).unwrap(); } let memtrie_changes = update.to_mem_trie_changes_only(); let state_root = memtries.apply_memtrie_changes(0, &memtrie_changes); diff --git a/core/store/src/trie/mem/updating.rs b/core/store/src/trie/mem/updating.rs index 1190c745d0c..77c01eed21e 100644 --- a/core/store/src/trie/mem/updating.rs +++ b/core/store/src/trie/mem/updating.rs @@ -13,8 +13,8 @@ use super::metrics::MEM_TRIE_NUM_NODES_CREATED_FROM_UPDATES; use super::node::{InputMemTrieNode, MemTrieNodeId, MemTrieNodeView}; use crate::trie::insert_delete::NodesStorage; use crate::trie::{ - Children, MemTrieChanges, NodeHandle, StorageHandle, TrieNode, TrieNodeWithSize, - TrieRefcountDeltaMap, ValueHandle, TRIE_COSTS, + Children, MemTrieChanges, NodeHandle, StorageHandle, StorageValueHandle, TrieNode, + TrieNodeWithSize, TrieRefcountDeltaMap, ValueHandle, TRIE_COSTS, }; use crate::{NibbleSlice, RawTrieNode, RawTrieNodeWithSize, TrieChanges}; use near_primitives::errors::StorageError; @@ -202,6 +202,10 @@ pub(crate) trait GenericTrieUpdate<'a, GenericTrieNodePtr, GenericValueHandle> { /// Squashes a node to ensure uniqueness of the trie structure. /// TODO(#12324): should be implemented using the methods above. fn generic_squash_node(&mut self, node_id: GenericUpdatedNodeId) -> Result<(), StorageError>; + + fn generic_store_value(&mut self, value: ValueUpdate) -> GenericValueHandle; + + fn generic_delete_value(&mut self, value: GenericValueHandle) -> Result<(), StorageError>; } /// Keeps values and internal nodes accessed on updating memtrie. @@ -338,6 +342,43 @@ impl<'a, M: ArenaMemory> GenericTrieUpdate<'a, MemTrieNodeId, FlatStateValue> self.squash_node(node_id); Ok(()) } + + fn generic_store_value(&mut self, value: ValueUpdate) -> FlatStateValue { + // First, set the value which will be stored in memtrie. + let flat_value = match &value { + ValueUpdate::MemtrieOnly(value) => return value.clone(), + ValueUpdate::MemtrieAndDisk(value) => FlatStateValue::on_disk(value.as_slice()), + }; + + // Then, record disk changes if needed. + let Some(tracked_node_changes) = self.tracked_trie_changes.as_mut() else { + return flat_value; + }; + let ValueUpdate::MemtrieAndDisk(value) = value else { + return flat_value; + }; + tracked_node_changes + .refcount_inserted_values + .entry(value) + .and_modify(|rc| *rc += 1) + .or_insert(1); + + flat_value + } + + fn generic_delete_value(&mut self, value: FlatStateValue) -> Result<(), StorageError> { + if let Some(tracked_node_changes) = self.tracked_trie_changes.as_mut() { + let hash = value.to_value_ref().hash; + tracked_node_changes.accesses.values.insert(hash, value); + tracked_node_changes + .refcount_deleted_hashes + .entry(hash) + .and_modify(|rc| *rc += 1) + .or_insert(1); + } + + Ok(()) + } } pub(crate) type TrieStorageNodePtr = CryptoHash; @@ -455,6 +496,32 @@ impl<'a> GenericTrieUpdate<'a, TrieStorageNodePtr, ValueHandle> for NodesStorage let trie = self.trie; trie.squash_node(self, StorageHandle(index)) } + + fn generic_store_value(&mut self, value: ValueUpdate) -> ValueHandle { + let ValueUpdate::MemtrieAndDisk(value) = value else { + unimplemented!( + "NodesStorage for Trie doesn't support value {value:?} \ + because disk updates must be generated." + ); + }; + + let value_len = value.len(); + self.values.push(Some(value)); + ValueHandle::InMemory(StorageValueHandle(self.values.len() - 1, value_len)) + } + + fn generic_delete_value(&mut self, value: ValueHandle) -> Result<(), StorageError> { + match value { + ValueHandle::HashAndSize(value) => { + self.trie.internal_retrieve_trie_node(&value.hash, true, true)?; + self.refcount_changes.subtract(value.hash, 1); + } + ValueHandle::InMemory(_) => { + // do nothing + } + } + Ok(()) + } } impl<'a, M: ArenaMemory> MemTrieUpdate<'a, M> { @@ -543,41 +610,19 @@ impl<'a, M: ArenaMemory> MemTrieUpdate<'a, M> { } } - fn add_refcount_to_value(&mut self, value: ValueUpdate) { - let Some(tracked_node_changes) = self.tracked_trie_changes.as_mut() else { - return; - }; - let ValueUpdate::MemtrieAndDisk(value) = value else { - return; - }; - tracked_node_changes - .refcount_inserted_values - .entry(value) - .and_modify(|rc| *rc += 1) - .or_insert(1); - } - - fn subtract_refcount_for_value(&mut self, value: FlatStateValue) { - if let Some(tracked_node_changes) = self.tracked_trie_changes.as_mut() { - let hash = value.to_value_ref().hash; - tracked_node_changes.accesses.values.insert(hash, value); - tracked_node_changes - .refcount_deleted_hashes - .entry(hash) - .and_modify(|rc| *rc += 1) - .or_insert(1); - } - } - /// Inserts the given key value pair into the trie. - pub fn insert(&mut self, key: &[u8], value: Vec) { - self.insert_impl(key, ValueUpdate::MemtrieAndDisk(value)); + pub fn insert(&mut self, key: &[u8], value: Vec) -> Result<(), StorageError> { + self.insert_impl(key, ValueUpdate::MemtrieAndDisk(value)) } /// Inserts the given key value pair into the trie, but the value may be a reference. /// This is used to update the in-memory trie only, without caring about on-disk changes. - pub fn insert_memtrie_only(&mut self, key: &[u8], value: FlatStateValue) { - self.insert_impl(key, ValueUpdate::MemtrieOnly(value)); + pub fn insert_memtrie_only( + &mut self, + key: &[u8], + value: FlatStateValue, + ) -> Result<(), StorageError> { + self.insert_impl(key, ValueUpdate::MemtrieOnly(value)) } /// Insertion logic. We descend from the root down to whatever node corresponds to @@ -585,13 +630,9 @@ impl<'a, M: ArenaMemory> MemTrieUpdate<'a, M> { /// the way to achieve that. This takes care of refcounting changes for existing /// nodes as well as values, but will not yet increment refcount for any newly /// created nodes - that's done at the end. - fn insert_impl(&mut self, key: &[u8], value: ValueUpdate) { + fn insert_impl(&mut self, key: &[u8], value: ValueUpdate) -> Result<(), StorageError> { let mut node_id = 0; // root let mut partial = NibbleSlice::new(key); - let flat_value = match &value { - ValueUpdate::MemtrieOnly(value) => value.clone(), - ValueUpdate::MemtrieAndDisk(value) => FlatStateValue::on_disk(value.as_slice()), - }; loop { // Take out the current node; we'd have to change it no matter what. @@ -599,27 +640,27 @@ impl<'a, M: ArenaMemory> MemTrieUpdate<'a, M> { match node { UpdatedMemTrieNode::Empty => { // There was no node here, create a new leaf. + let value_handle = self.generic_store_value(value); self.place_node( node_id, UpdatedMemTrieNode::Leaf { extension: partial.encoded(true).into_vec().into_boxed_slice(), - value: flat_value, + value: value_handle, }, ); - self.add_refcount_to_value(value); break; } UpdatedMemTrieNode::Branch { children, value: old_value } => { if partial.is_empty() { // This branch node is exactly where the value should be added. if let Some(value) = old_value { - self.subtract_refcount_for_value(value); + self.generic_delete_value(value)?; } + let value_handle = self.generic_store_value(value); self.place_node( node_id, - UpdatedMemTrieNode::Branch { children, value: Some(flat_value) }, + UpdatedMemTrieNode::Branch { children, value: Some(value_handle) }, ); - self.add_refcount_to_value(value); break; } else { // Continue descending into the branch, possibly adding a new child. @@ -644,12 +685,12 @@ impl<'a, M: ArenaMemory> MemTrieUpdate<'a, M> { let common_prefix = partial.common_prefix(&existing_key); if common_prefix == existing_key.len() && common_prefix == partial.len() { // We're at the exact leaf. Rewrite the value at this leaf. - self.subtract_refcount_for_value(old_value); + self.generic_delete_value(old_value)?; + let value_handle = self.generic_store_value(value); self.place_node( node_id, - UpdatedMemTrieNode::Leaf { extension, value: flat_value }, + UpdatedMemTrieNode::Leaf { extension, value: value_handle }, ); - self.add_refcount_to_value(value); break; } else if common_prefix == 0 { // Convert the leaf to an equivalent branch. We are not adding @@ -757,6 +798,8 @@ impl<'a, M: ArenaMemory> MemTrieUpdate<'a, M> { } } } + + Ok(()) } /// Deletes a key from the trie. @@ -766,7 +809,7 @@ impl<'a, M: ArenaMemory> MemTrieUpdate<'a, M> { /// consistent by changing the types of any nodes along the way. /// /// Deleting a non-existent key is allowed, and is a no-op. - pub fn delete(&mut self, key: &[u8]) { + pub fn delete(&mut self, key: &[u8]) -> Result<(), StorageError> { let mut node_id = 0; // root let mut partial = NibbleSlice::new(key); let mut path = vec![]; // for squashing at the end. @@ -779,17 +822,17 @@ impl<'a, M: ArenaMemory> MemTrieUpdate<'a, M> { UpdatedMemTrieNode::Empty => { // Nothing to delete. self.place_node(node_id, UpdatedMemTrieNode::Empty); - return; + return Ok(()); } UpdatedMemTrieNode::Leaf { extension, value } => { if NibbleSlice::from_encoded(&extension).0 == partial { - self.subtract_refcount_for_value(value); + self.generic_delete_value(value)?; self.place_node(node_id, UpdatedMemTrieNode::Empty); break; } else { // Key being deleted doesn't exist. self.place_node(node_id, UpdatedMemTrieNode::Leaf { extension, value }); - return; + return Ok(()); } } UpdatedMemTrieNode::Branch { children: old_children, value } => { @@ -800,9 +843,9 @@ impl<'a, M: ArenaMemory> MemTrieUpdate<'a, M> { node_id, UpdatedMemTrieNode::Branch { children: old_children, value }, ); - return; + return Ok(()); }; - self.subtract_refcount_for_value(value.unwrap()); + self.generic_delete_value(value.unwrap())?; self.place_node( node_id, UpdatedMemTrieNode::Branch { children: old_children, value: None }, @@ -820,7 +863,7 @@ impl<'a, M: ArenaMemory> MemTrieUpdate<'a, M> { node_id, UpdatedMemTrieNode::Branch { children: old_children, value }, ); - return; + return Ok(()); } }; let new_child_id = self.ensure_updated(old_child_id); @@ -859,7 +902,7 @@ impl<'a, M: ArenaMemory> MemTrieUpdate<'a, M> { node_id, UpdatedMemTrieNode::Extension { extension, child }, ); - return; + return Ok(()); } } } @@ -869,6 +912,7 @@ impl<'a, M: ArenaMemory> MemTrieUpdate<'a, M> { for node_id in path.into_iter().rev() { self.squash_node(node_id); } + Ok(()) } /// When we delete keys, it may be necessary to change types of some nodes, @@ -1302,9 +1346,9 @@ mod tests { }); for (key, value) in changes { if let Some(value) = value { - update.insert(&key, value); + update.insert(&key, value).unwrap(); } else { - update.delete(&key); + update.delete(&key).unwrap(); } } update.to_trie_changes().0 @@ -1319,9 +1363,9 @@ mod tests { }); for (key, value) in changes { if let Some(value) = value { - update.insert_memtrie_only(&key, FlatStateValue::on_disk(&value)); + update.insert_memtrie_only(&key, FlatStateValue::on_disk(&value)).unwrap(); } else { - update.delete(&key); + update.delete(&key).unwrap(); } } update.to_mem_trie_changes_only() @@ -1699,9 +1743,9 @@ mod tests { for (key, value) in changes { if let Some(value) = value { - update.insert_memtrie_only(&key, FlatStateValue::on_disk(&value)); + update.insert_memtrie_only(&key, FlatStateValue::on_disk(&value)).unwrap(); } else { - update.delete(&key); + update.delete(&key).unwrap(); } } diff --git a/core/store/src/trie/mod.rs b/core/store/src/trie/mod.rs index 86951272be2..cf014461801 100644 --- a/core/store/src/trie/mod.rs +++ b/core/store/src/trie/mod.rs @@ -888,23 +888,6 @@ impl Trie { memory_usage } - fn delete_value( - &self, - memory: &mut NodesStorage, - value: &ValueHandle, - ) -> Result<(), StorageError> { - match value { - ValueHandle::HashAndSize(value) => { - self.internal_retrieve_trie_node(&value.hash, true, true)?; - memory.refcount_changes.subtract(value.hash, 1); - } - ValueHandle::InMemory(_) => { - // do nothing - } - } - Ok(()) - } - /// Prints the trie nodes starting from `hash`, up to `max_depth` depth. The node hash can be any node in the trie. /// Depending on arguments provided, can limit output to no more than `limit` entries, /// show only subtree for a given `record_type`, or skip subtrees where `AccountId` is less than `from` or greater than `to`. @@ -1623,8 +1606,8 @@ impl Trie { let mut trie_update = guard.update(self.root, true)?; for (key, value) in changes { match value { - Some(arr) => trie_update.insert(&key, arr), - None => trie_update.delete(&key), + Some(arr) => trie_update.insert(&key, arr)?, + None => trie_update.delete(&key)?, } } let (trie_changes, trie_accesses) = trie_update.to_trie_changes();