[Memtrie] (3/n) Implement memtrie update logic. (#10015)

This is a complete implementation of the insertion and deletion logic on
top of the in-memory trie. Originally coming from @Longarithm's
prototype, it's been cleaned up, thoroughly reviewed, and thoroughly
tested.

The MemTrieChanges data structure, similar to the existing TrieChanges
data structure, describes all that is needed to apply changes to an
in-memory trie. Additionally, MemTrieChanges can be embedded within
TrieChanges to apply both in-memory and on-disk changes together.

When an in-memory trie is updated, the mutations (insertions and
deletions of key-value pairs) go through only the MemTrieUpdate<'a>
struct. This struct is responsible for calculating both the
MemTrieChanges as well as the TrieChanges. By accessing only the
existing in-memory trie nodes, it is able to compute both in-memory
changes and on-disk changes. (There is also an option to compute only
in-memory changes, which is useful during initial loading when we need
to apply flat state deltas).

Extensive testing has been implemented to assert that the following are
all completely consistent with one another:
- Computing the on-disk trie changes using the previously existing logic
implemented in Trie::update;
- Computing the in-memory trie changes using the newly implemented
MemTrieUpdate;
- Computing both in-memory and on-disk trie changes together using
MemTrieUpdate.

Also, a manual test plus a randomized test both independently achieve
effectively 100% coverage of all code paths in updating.rs.

core/primitives/src/state.rs

@@ -99,4 +99,11 @@ impl FlatStateValue {
             Self::Inlined(value) => ValueRef::new(value),
         }
     }
+
+    pub fn value_len(&self) -> usize {
+        match self {
+            Self::Ref(value_ref) => value_ref.len(),
+            Self::Inlined(value) => value.len(),
+        }
+    }
 }

core/store/src/trie/insert_delete.rs

@@ -610,7 +610,13 @@ impl Trie {
             last_hash = key;
         }
         let (insertions, deletions) = memory.refcount_changes.into_changes();
-        Ok(TrieChanges { old_root: *old_root, new_root: last_hash, insertions, deletions })
+        Ok(TrieChanges {
+            old_root: *old_root,
+            new_root: last_hash,
+            insertions,
+            deletions,
+            mem_trie_changes: None,
+        })
     }
 
     fn flatten_value(memory: &mut NodesStorage, value: ValueHandle) -> ValueRef {

core/store/src/trie/mem/metrics.rs

@@ -1,4 +1,6 @@
-use near_o11y::metrics::{try_create_int_gauge_vec, IntGaugeVec};
+use near_o11y::metrics::{
+    try_create_int_counter_vec, try_create_int_gauge_vec, IntCounterVec, IntGaugeVec,
+};
 use once_cell::sync::Lazy;
 
 pub static MEM_TRIE_NUM_ROOTS: Lazy<IntGaugeVec> = Lazy::new(|| {
@@ -9,3 +11,12 @@ pub static MEM_TRIE_NUM_ROOTS: Lazy<IntGaugeVec> = Lazy::new(|| {
     )
     .unwrap()
 });
+
+pub static MEM_TRIE_NUM_NODES_CREATED_FROM_UPDATES: Lazy<IntCounterVec> = Lazy::new(|| {
+    try_create_int_counter_vec(
+        "near_mem_trie_num_nodes_created_from_updates",
+        "Number of trie nodes created by applying updates",
+        &["shard_uid"],
+    )
+    .unwrap()
+});

core/store/src/trie/mem/mod.rs

@@ -1,6 +1,8 @@
 use self::arena::Arena;
 use self::metrics::MEM_TRIE_NUM_ROOTS;
 use self::node::{MemTrieNodeId, MemTrieNodePtr};
+use self::updating::MemTrieUpdate;
+use near_primitives::errors::StorageError;
 use near_primitives::hash::CryptoHash;
 use near_primitives::shard_layout::ShardUId;
 use near_primitives::types::{BlockHeight, StateRoot};
@@ -13,6 +15,7 @@ pub mod loading;
 pub mod lookup;
 mod metrics;
 pub mod node;
+pub mod updating;
 
 /// Check this, because in the code we conveniently assume usize is 8 bytes.
 /// In-memory trie can't possibly work under 32-bit anyway.
@@ -136,6 +139,33 @@ impl MemTries {
     pub fn num_roots(&self) -> usize {
         self.heights.iter().map(|(_, v)| v.len()).sum()
     }
+
+    pub fn update(
+        &self,
+        root: CryptoHash,
+        track_disk_changes: bool,
+    ) -> Result<MemTrieUpdate, StorageError> {
+        let root_id = if root == CryptoHash::default() {
+            None
+        } else {
+            let root_id = self
+                .get_root(&root)
+                .ok_or_else(|| {
+                    StorageError::StorageInconsistentState(format!(
+                        "Failed to find root node {:?} in memtrie",
+                        root
+                    ))
+                })?
+                .id();
+            Some(root_id)
+        };
+        Ok(MemTrieUpdate::new(
+            root_id,
+            &self.arena.memory(),
+            self.shard_uid.to_string(),
+            track_disk_changes,
+        ))
+    }
 }
 
 #[cfg(test)]

core/store/src/trie/mem/node/encoding.rs

@@ -36,8 +36,8 @@ pub(crate) struct NonLeafHeader {
 }
 
 impl NonLeafHeader {
-    pub(crate) fn new(memory_usage: u64) -> Self {
-        Self { hash: CryptoHash::default(), memory_usage }
+    pub(crate) fn new(memory_usage: u64, node_hash: Option<CryptoHash>) -> Self {
+        Self { hash: node_hash.unwrap_or_default(), memory_usage }
     }
 }
 
@@ -103,7 +103,11 @@ impl BorshFixedSize for BranchWithValueHeader {
 
 impl MemTrieNodeId {
     /// Encodes the data.
-    pub(crate) fn new_impl(arena: &mut Arena, node: InputMemTrieNode) -> Self {
+    pub(crate) fn new_impl(
+        arena: &mut Arena,
+        node: InputMemTrieNode,
+        node_hash: Option<CryptoHash>,
+    ) -> Self {
         // We add reference to all the children when creating the node.
         // As for the refcount of this newly created node, it starts at 0.
         // It is expected that either our parent will increment our own
@@ -187,7 +191,7 @@ impl MemTrieNodeId {
                 );
                 data.encode(ExtensionHeader {
                     common: CommonHeader { refcount: 0, kind: NodeKind::Extension },
-                    nonleaf: NonLeafHeader::new(memory_usage),
+                    nonleaf: NonLeafHeader::new(memory_usage, node_hash),
                     child: child.pos,
                     extension: extension_header,
                 });
@@ -202,7 +206,7 @@ impl MemTrieNodeId {
                 );
                 data.encode(BranchHeader {
                     common: CommonHeader { refcount: 0, kind: NodeKind::Branch },
-                    nonleaf: NonLeafHeader::new(memory_usage),
+                    nonleaf: NonLeafHeader::new(memory_usage, node_hash),
                     children: children_header,
                 });
                 data.encode_flexible(&children_header, children);
@@ -219,7 +223,7 @@ impl MemTrieNodeId {
                 );
                 data.encode(BranchWithValueHeader {
                     common: CommonHeader { refcount: 0, kind: NodeKind::BranchWithValue },
-                    nonleaf: NonLeafHeader::new(memory_usage),
+                    nonleaf: NonLeafHeader::new(memory_usage, node_hash),
                     children: children_header,
                     value: value_header,
                 });

core/store/src/trie/mem/node/mod.rs

@@ -29,7 +29,11 @@ pub struct MemTrieNodeId {
 
 impl MemTrieNodeId {
     pub fn new(arena: &mut Arena, input: InputMemTrieNode) -> Self {
-        Self::new_impl(arena, input)
+        Self::new_impl(arena, input, None)
+    }
+
+    pub fn new_with_hash(arena: &mut Arena, input: InputMemTrieNode, hash: CryptoHash) -> Self {
+        Self::new_impl(arena, input, Some(hash))
     }
 
     pub fn as_ptr<'a>(&self, arena: &'a ArenaMemory) -> MemTrieNodePtr<'a> {