diff --git a/core/store/src/test_utils.rs b/core/store/src/test_utils.rs index 6974ad80b9a..3aab48d8d6f 100644 --- a/core/store/src/test_utils.rs +++ b/core/store/src/test_utils.rs @@ -94,8 +94,8 @@ impl TestTriesBuilder { self } - pub fn with_flat_storage(mut self) -> Self { - self.enable_flat_storage = true; + pub fn with_flat_storage(mut self, enable: bool) -> Self { + self.enable_flat_storage = enable; self } diff --git a/core/store/src/trie/insert_delete.rs b/core/store/src/trie/insert_delete.rs index e6174d06276..a890b202abb 100644 --- a/core/store/src/trie/insert_delete.rs +++ b/core/store/src/trie/insert_delete.rs @@ -320,6 +320,7 @@ impl Trie { let mut partial = partial; let root_node = handle; let mut path: Vec = Vec::new(); + let mut key_deleted = true; loop { path.push(handle); let TrieNodeWithSize { node, memory_usage } = memory.destroy(handle); @@ -327,6 +328,7 @@ impl Trie { match node { TrieNode::Empty => { memory.store_at(handle, TrieNodeWithSize::empty()); + key_deleted = false; break; } TrieNode::Leaf(key, value) => { @@ -338,25 +340,33 @@ impl Trie { let leaf_node = TrieNode::Leaf(key, value); let memory_usage = leaf_node.memory_usage_direct(memory); memory.store_at(handle, TrieNodeWithSize::new(leaf_node, memory_usage)); + key_deleted = false; break; } } TrieNode::Branch(mut children, value) => { if partial.is_empty() { - if let Some(value) = &value { - self.delete_value(memory, value)?; - } - if children.iter().count() == 0 { - memory.store_at(handle, TrieNodeWithSize::empty()); - } else { - Trie::calc_memory_usage_and_store( - memory, + if value.is_none() { + // Key being deleted doesn't exist. + memory.store_at( handle, - children_memory_usage, - TrieNode::Branch(children, None), - None, + TrieNodeWithSize::new( + TrieNode::Branch(children, value), + memory_usage, + ), ); + key_deleted = false; + break; } + self.delete_value(memory, &value.unwrap())?; + Trie::calc_memory_usage_and_store( + memory, + handle, + children_memory_usage, + TrieNode::Branch(children, None), + None, + ); + // if needed, branch will be squashed at the end of the function. break; } else { let child = &mut children[partial.at(0)]; @@ -386,6 +396,7 @@ impl Trie { memory_usage, ), ); + key_deleted = false; break; } } @@ -415,71 +426,85 @@ impl Trie { handle, TrieNodeWithSize::new(TrieNode::Extension(key, child), memory_usage), ); + key_deleted = false; break; } } } } - self.fix_nodes(memory, path)?; + self.fix_nodes(memory, path, key_deleted)?; Ok(root_node) } + /// Iterates over nodes in `path`, changing their types where needed, + /// if `key_deleted` is true, so trie structure has to change. + /// If `key_deleted` is false, only recomputes memory usages along the path. fn fix_nodes( &self, memory: &mut NodesStorage, path: Vec, + key_deleted: bool, ) -> Result<(), StorageError> { let mut child_memory_usage = 0; for handle in path.into_iter().rev() { let TrieNodeWithSize { node, memory_usage } = memory.destroy(handle); let memory_usage = memory_usage + child_memory_usage; - match node { - TrieNode::Empty => { - memory.store_at(handle, TrieNodeWithSize::empty()); - } - TrieNode::Leaf(key, value) => { - memory.store_at( - handle, - TrieNodeWithSize::new(TrieNode::Leaf(key, value), memory_usage), - ); - } - TrieNode::Branch(mut children, value) => { - for child in children.0.iter_mut() { - if let Some(NodeHandle::InMemory(h)) = child { - if let TrieNode::Empty = memory.node_ref(*h).node { - *child = None + if key_deleted { + match node { + TrieNode::Empty => { + memory.store_at(handle, TrieNodeWithSize::empty()); + } + TrieNode::Leaf(key, value) => { + memory.store_at( + handle, + TrieNodeWithSize::new(TrieNode::Leaf(key, value), memory_usage), + ); + } + TrieNode::Branch(mut children, value) => { + for child in children.0.iter_mut() { + if let Some(NodeHandle::InMemory(h)) = child { + if let TrieNode::Empty = memory.node_ref(*h).node { + *child = None + } } } - } - let num_children = children.iter().count(); - if num_children == 0 { - if let Some(value) = value { - let empty = NibbleSlice::new(&[]).encoded(true).into_vec(); - let leaf_node = TrieNode::Leaf(empty, value); - let memory_usage = leaf_node.memory_usage_direct(memory); - memory.store_at(handle, TrieNodeWithSize::new(leaf_node, memory_usage)); + let num_children = children.iter().count(); + if num_children == 0 { + if let Some(value) = value { + let empty = NibbleSlice::new(&[]).encoded(true).into_vec(); + let leaf_node = TrieNode::Leaf(empty, value); + let memory_usage = leaf_node.memory_usage_direct(memory); + memory.store_at( + handle, + TrieNodeWithSize::new(leaf_node, memory_usage), + ); + } else { + memory.store_at(handle, TrieNodeWithSize::empty()); + } + } else if num_children == 1 && value.is_none() { + // Branch with one child becomes extension + // Extension followed by leaf becomes leaf + // Extension followed by extension becomes extension + let idx = children.iter().next().unwrap().0; + let child = children[idx].take().unwrap(); + let key = NibbleSlice::new(&[(idx << 4) as u8]) + .encoded_leftmost(1, false) + .into_vec(); + self.fix_extension_node(memory, handle, key, child)?; } else { - memory.store_at(handle, TrieNodeWithSize::empty()); + let node = TrieNodeWithSize::new( + TrieNode::Branch(children, value), + memory_usage, + ); + memory.store_at(handle, node); } - } else if num_children == 1 && value.is_none() { - // Branch with one child becomes extension - // Extension followed by leaf becomes leaf - // Extension followed by extension becomes extension - let idx = children.iter().next().unwrap().0; - let child = children[idx].take().unwrap(); - let key = NibbleSlice::new(&[(idx << 4) as u8]) - .encoded_leftmost(1, false) - .into_vec(); + } + TrieNode::Extension(key, child) => { self.fix_extension_node(memory, handle, key, child)?; - } else { - let node = - TrieNodeWithSize::new(TrieNode::Branch(children, value), memory_usage); - memory.store_at(handle, node); } } - TrieNode::Extension(key, child) => { - self.fix_extension_node(memory, handle, key, child)?; - } + } else { + memory.store_at(handle, TrieNodeWithSize { node, memory_usage }); } child_memory_usage = memory.node_ref(handle).memory_usage; } diff --git a/core/store/src/trie/mem/loading.rs b/core/store/src/trie/mem/loading.rs index a14242f5c3b..9267c3f1028 100644 --- a/core/store/src/trie/mem/loading.rs +++ b/core/store/src/trie/mem/loading.rs @@ -197,7 +197,7 @@ mod tests { use rand::{Rng, SeedableRng}; fn check(keys: Vec>) { - let shard_tries = TestTriesBuilder::new().with_flat_storage().build(); + let shard_tries = TestTriesBuilder::new().with_flat_storage(true).build(); let shard_uid = ShardUId::single_shard(); let changes = keys.iter().map(|key| (key.to_vec(), Some(key.to_vec()))).collect::>(); let changes = simplify_changes(&changes); diff --git a/core/store/src/trie/mem/updating.rs b/core/store/src/trie/mem/updating.rs index 6360c2a3143..2eb3400a5f4 100644 --- a/core/store/src/trie/mem/updating.rs +++ b/core/store/src/trie/mem/updating.rs @@ -1306,6 +1306,7 @@ mod tests { key.push(nibble0 << 4 | nibble1); } } + let mut value_length = rand::thread_rng().gen_range(0..=10); if value_length == 10 { value_length = 8000; // make a long value that is not inlined diff --git a/core/store/src/trie/mod.rs b/core/store/src/trie/mod.rs index 33fce745135..39f4a19f4b0 100644 --- a/core/store/src/trie/mod.rs +++ b/core/store/src/trie/mod.rs @@ -80,7 +80,7 @@ pub struct TrieCosts { } /// Whether a key lookup will be performed through flat storage or through iterating the trie -#[derive(PartialEq, Eq)] +#[derive(Clone, Copy, PartialEq, Eq)] pub enum KeyLookupMode { FlatStorage, Trie, @@ -751,9 +751,16 @@ impl Trie { #[cfg(test)] fn memory_usage_verify(&self, memory: &NodesStorage, handle: NodeHandle) -> u64 { + // Cannot compute memory usage naively if given only partial storage. + if self.storage.as_partial_storage().is_some() { + return 0; + } + // We don't want to impact recorded storage by retrieving nodes for + // this sanity check. if self.recorder.is_some() { return 0; } + let TrieNodeWithSize { node, memory_usage } = match handle { NodeHandle::InMemory(h) => memory.node_ref(h).clone(), NodeHandle::Hash(h) => self.retrieve_node(&h).expect("storage failure").1, @@ -1838,7 +1845,7 @@ mod tests { fn test_contains_key() { let sid = ShardUId::single_shard(); let bid = CryptoHash::default(); - let tries = TestTriesBuilder::new().with_flat_storage().build(); + let tries = TestTriesBuilder::new().with_flat_storage(true).build(); let initial = vec![ (vec![99, 44, 100, 58, 58, 49], Some(vec![1])), (vec![99, 44, 100, 58, 58, 50], Some(vec![1])), diff --git a/core/store/src/trie/state_parts.rs b/core/store/src/trie/state_parts.rs index f2482d35511..fecf2b6e03f 100644 --- a/core/store/src/trie/state_parts.rs +++ b/core/store/src/trie/state_parts.rs @@ -1147,7 +1147,7 @@ mod tests { fn get_trie_nodes_for_part_with_flat_storage() { let value_len = 1000usize; - let tries = TestTriesBuilder::new().with_flat_storage().build(); + let tries = TestTriesBuilder::new().with_flat_storage(true).build(); let shard_uid = ShardUId::single_shard(); let block_hash = CryptoHash::default(); let part_id = PartId::new(1, 3); diff --git a/core/store/src/trie/trie_recording.rs b/core/store/src/trie/trie_recording.rs index 819ef5a295d..d44260c8238 100644 --- a/core/store/src/trie/trie_recording.rs +++ b/core/store/src/trie/trie_recording.rs @@ -59,18 +59,20 @@ mod trie_recording_tests { }; use crate::trie::mem::metrics::MEM_TRIE_NUM_LOOKUPS; use crate::trie::TrieNodesCount; - use crate::{DBCol, Store, Trie}; + use crate::{DBCol, KeyLookupMode, PartialStorage, ShardTries, Store, Trie}; use borsh::BorshDeserialize; + use near_primitives::challenge::PartialState; use near_primitives::hash::{hash, CryptoHash}; use near_primitives::shard_layout::{get_block_shard_uid, ShardUId}; use near_primitives::state::ValueRef; use near_primitives::types::chunk_extra::ChunkExtra; use near_primitives::types::StateRoot; + use rand::prelude::SliceRandom; use rand::{random, thread_rng, Rng}; use std::collections::{HashMap, HashSet}; use std::num::NonZeroU32; - const NUM_ITERATIONS_PER_TEST: usize = 100; + const NUM_ITERATIONS_PER_TEST: usize = 300; /// Prepared on-disk trie and flat storage for testing. struct PreparedTrie { @@ -92,14 +94,18 @@ mod trie_recording_tests { /// storage with some dummy block info. If `use_missing_keys` is true, /// the keys to test with will also include some keys that are not in the /// trie. - fn prepare_trie(use_missing_keys: bool) -> PreparedTrie { - let tries_for_building = TestTriesBuilder::new().with_flat_storage().build(); + fn prepare_trie( + use_missing_keys: bool, + p_existing_key: f64, + p_missing_key: f64, + ) -> PreparedTrie { + let tries_for_building = TestTriesBuilder::new().with_flat_storage(true).build(); let shard_uid = ShardUId::single_shard(); let trie_changes = gen_larger_changes(&mut thread_rng(), 50); let trie_changes = simplify_changes(&trie_changes); if trie_changes.is_empty() { // try again - return prepare_trie(use_missing_keys); + return prepare_trie(use_missing_keys, p_existing_key, p_missing_key); } let state_root = test_populate_trie( &tries_for_building, @@ -131,19 +137,32 @@ mod trie_recording_tests { .iter() .map(|(key, value)| (key.clone(), value.clone().unwrap())) .collect::>(); - let (keys_to_get, keys_to_get_ref) = trie_changes - .iter() - .map(|(key, _)| { - let mut key = key.clone(); - if use_missing_keys { - key.push(100); - } - key - }) - .partition::, _>(|_| random()); - let updates = trie_changes + let existing_keys: HashSet<_> = trie_changes + .into_iter() + .map(|(key, _)| key) + .filter(|_| thread_rng().gen_bool(p_existing_key)) + .collect(); + let missing_keys = if use_missing_keys { + existing_keys + .iter() + .cloned() + .map(|mut key| { + *key.last_mut().unwrap() = 100; + key + }) + .filter(|key| !existing_keys.contains(key) && thread_rng().gen_bool(p_missing_key)) + .collect::>() + .into_iter() + .collect::>() + } else { + vec![] + }; + let mut keys: Vec<_> = + existing_keys.iter().cloned().chain(missing_keys.into_iter()).collect(); + keys.shuffle(&mut thread_rng()); + let updates = keys .iter() - .map(|(key, _)| { + .map(|key| { let value = if thread_rng().gen_bool(0.5) { Some(vec![thread_rng().gen_range(0..10) as u8]) } else { @@ -153,6 +172,8 @@ mod trie_recording_tests { }) .filter(|_| random()) .collect::>(); + let (keys_to_get, keys_to_get_ref) = + keys.into_iter().filter(|_| random()).partition::, _>(|_| random()); PreparedTrie { store: tries_for_building.get_store(), shard_uid, @@ -191,15 +212,48 @@ mod trie_recording_tests { update.commit().unwrap(); } + fn get_trie_for_shard( + tries: &ShardTries, + shard_uid: ShardUId, + state_root: StateRoot, + use_flat_storage: bool, + ) -> Trie { + if use_flat_storage { + tries.get_trie_with_block_hash_for_shard( + shard_uid, + state_root, + &CryptoHash::default(), + false, + ) + } else { + tries.get_trie_for_shard(shard_uid, state_root) + } + } + + /// Assert equality of partial storages with human-readable output. + fn assert_partial_storage(storage: &PartialStorage, other_storage: &PartialStorage) { + let PartialState::TrieValues(nodes) = &storage.nodes; + let PartialState::TrieValues(other_nodes) = &other_storage.nodes; + let nodes: HashSet> = HashSet::from_iter(nodes.into_iter().map(|key| key.to_vec())); + let other_nodes: HashSet> = + HashSet::from_iter(other_nodes.into_iter().map(|key| key.to_vec())); + let d: Vec<&Vec> = other_nodes.difference(&nodes).collect(); + assert_eq!(d, Vec::<&Vec>::default(), "Missing nodes in first storage"); + let d: Vec<&Vec> = nodes.difference(&other_nodes).collect(); + assert_eq!(d, Vec::<&Vec>::default(), "Missing nodes in second storage"); + } + /// Verifies that when operating on a trie, the results are completely consistent /// regardless of whether we're operating on the real storage (with or without chunk /// cache), while recording reads, or when operating on recorded partial storage. fn test_trie_recording_consistency( enable_accounting_cache: bool, use_missing_keys: bool, - use_in_memory_tries: bool, + use_flat_storage: bool, ) { for _ in 0..NUM_ITERATIONS_PER_TEST { + let p_existing_key = thread_rng().gen_range(0.3..1.0); + let p_missing_key = thread_rng().gen_range(0.7..1.0); let PreparedTrie { store, shard_uid, @@ -208,30 +262,39 @@ mod trie_recording_tests { keys_to_get_ref, updates, state_root, - } = prepare_trie(use_missing_keys); - let tries = if use_in_memory_tries { - TestTriesBuilder::new().with_store(store.clone()).with_in_memory_tries().build() - } else { - TestTriesBuilder::new().with_store(store.clone()).build() - }; + } = prepare_trie(use_missing_keys, p_existing_key, p_missing_key); + let tries = TestTriesBuilder::new() + .with_store(store.clone()) + .with_flat_storage(use_flat_storage) + .build(); + let lookup_mode = + if use_flat_storage { KeyLookupMode::FlatStorage } else { KeyLookupMode::Trie }; let mem_trie_lookup_counts_before = MEM_TRIE_NUM_LOOKUPS.get(); - if use_in_memory_tries { - // Delete the on-disk state so that we really know we're using - // in-memory tries. - destructively_delete_in_memory_state_from_disk(&store, &data_in_trie); + // Check that while using flat storage counters are all zero. + // Only use get_optimized_ref(), because get() will actually + // dereference values which can cause trie reads. + if use_flat_storage { + let trie = get_trie_for_shard(&tries, shard_uid, state_root, use_flat_storage); + for key in data_in_trie.keys() { + trie.get_optimized_ref(key, lookup_mode).unwrap(); + } + assert_eq!( + trie.get_trie_nodes_count(), + TrieNodesCount { db_reads: 0, mem_reads: 0 } + ); } // Let's capture the baseline node counts - this is what will happen // in production. - let trie = tries.get_trie_for_shard(shard_uid, state_root); + let trie = get_trie_for_shard(&tries, shard_uid, state_root, use_flat_storage); trie.accounting_cache.borrow_mut().set_enabled(enable_accounting_cache); for key in &keys_to_get { assert_eq!(trie.get(key).unwrap(), data_in_trie.get(key).cloned()); } for key in &keys_to_get_ref { assert_eq!( - trie.get_optimized_ref(key, crate::KeyLookupMode::Trie) + trie.get_optimized_ref(key, lookup_mode) .unwrap() .map(|value| value.into_value_ref()), data_in_trie.get(key).map(|value| ValueRef::new(&value)) @@ -243,56 +306,15 @@ mod trie_recording_tests { // Now let's do this again while recording, and make sure that the counters // we get are exactly the same. - let trie = tries.get_trie_for_shard(shard_uid, state_root).recording_reads(); - trie.accounting_cache.borrow_mut().set_enabled(enable_accounting_cache); - for key in &keys_to_get { - assert_eq!(trie.get(key).unwrap(), data_in_trie.get(key).cloned()); - } - for key in &keys_to_get_ref { - assert_eq!( - trie.get_optimized_ref(key, crate::KeyLookupMode::Trie) - .unwrap() - .map(|value| value.into_value_ref()), - data_in_trie.get(key).map(|value| ValueRef::new(&value)) - ); - } - assert_eq!(trie.get_trie_nodes_count(), baseline_trie_nodes_count); - trie.update(updates.iter().cloned()).unwrap(); - - // Now, let's check that when doing the same lookups with the captured partial storage, - // we still get the same counters. - let partial_storage = trie.recorded_storage().unwrap(); - println!( - "Partial storage has {} nodes from {} entries", - partial_storage.nodes.len(), - data_in_trie.len() - ); - let trie = Trie::from_recorded_storage(partial_storage.clone(), state_root, false); - trie.accounting_cache.borrow_mut().set_enabled(enable_accounting_cache); - for key in &keys_to_get { - assert_eq!(trie.get(key).unwrap(), data_in_trie.get(key).cloned()); - } - for key in &keys_to_get_ref { - assert_eq!( - trie.get_optimized_ref(key, crate::KeyLookupMode::Trie) - .unwrap() - .map(|value| value.into_value_ref()), - data_in_trie.get(key).map(|value| ValueRef::new(&value)) - ); - } - assert_eq!(trie.get_trie_nodes_count(), baseline_trie_nodes_count); - trie.update(updates.iter().cloned()).unwrap(); - - // Build a Trie using recorded storage and enable recording_reads on this Trie - let trie = - Trie::from_recorded_storage(partial_storage, state_root, false).recording_reads(); + let trie = get_trie_for_shard(&tries, shard_uid, state_root, use_flat_storage) + .recording_reads(); trie.accounting_cache.borrow_mut().set_enabled(enable_accounting_cache); for key in &keys_to_get { assert_eq!(trie.get(key).unwrap(), data_in_trie.get(key).cloned()); } for key in &keys_to_get_ref { assert_eq!( - trie.get_optimized_ref(key, crate::KeyLookupMode::Trie) + trie.get_optimized_ref(key, lookup_mode) .unwrap() .map(|value| value.into_value_ref()), data_in_trie.get(key).map(|value| ValueRef::new(&value)) @@ -300,136 +322,16 @@ mod trie_recording_tests { } assert_eq!(trie.get_trie_nodes_count(), baseline_trie_nodes_count); trie.update(updates.iter().cloned()).unwrap(); - - if use_in_memory_tries { - // sanity check that we did indeed use in-memory tries. - assert!(MEM_TRIE_NUM_LOOKUPS.get() > mem_trie_lookup_counts_before); - } - } - } - - #[test] - fn test_trie_recording_consistency_no_accounting_cache() { - test_trie_recording_consistency(false, false, false); - } - - #[test] - fn test_trie_recording_consistency_with_accounting_cache() { - test_trie_recording_consistency(true, false, false); - } - - #[test] - fn test_trie_recording_consistency_no_accounting_cache_with_missing_keys() { - test_trie_recording_consistency(false, true, false); - } - - #[test] - fn test_trie_recording_consistency_with_accounting_cache_and_missing_keys() { - test_trie_recording_consistency(true, true, false); - } - - #[test] - fn test_trie_recording_consistency_memtrie_no_accounting_cache() { - test_trie_recording_consistency(false, false, true); - } - - #[test] - fn test_trie_recording_consistency_memtrie_with_accounting_cache() { - test_trie_recording_consistency(true, false, true); - } - - #[test] - fn test_trie_recording_consistency_memtrie_no_accounting_cache_with_missing_keys() { - test_trie_recording_consistency(false, true, true); - } - - #[test] - fn test_trie_recording_consistency_memtrie_with_accounting_cache_and_missing_keys() { - test_trie_recording_consistency(true, true, true); - } - - /// Verifies that when operating on a trie, the results are completely consistent - /// regardless of whether we're operating on the real storage (with or without chunk - /// cache), while recording reads, or when operating on recorded partial storage. - /// This test additionally verifies this when flat storage is used. - fn test_trie_recording_consistency_with_flat_storage( - enable_accounting_cache: bool, - use_missing_keys: bool, - use_in_memory_tries: bool, - ) { - for _ in 0..NUM_ITERATIONS_PER_TEST { - let PreparedTrie { - store, - shard_uid, - data_in_trie, - keys_to_get, - keys_to_get_ref, - updates, - state_root, - } = prepare_trie(use_missing_keys); - let tries = if use_in_memory_tries { - TestTriesBuilder::new() - .with_store(store.clone()) - .with_flat_storage() - .with_in_memory_tries() - .build() - } else { - TestTriesBuilder::new().with_store(store.clone()).with_flat_storage().build() - }; - let mem_trie_lookup_counts_before = MEM_TRIE_NUM_LOOKUPS.get(); - - if use_in_memory_tries { - // Delete the on-disk state so that we really know we're using - // in-memory tries. - destructively_delete_in_memory_state_from_disk(&store, &data_in_trie); - } - // Check that the trie is using flat storage, so that counters are all zero. - // Only use get_optimized_ref(), because get() will actually dereference values which can - // cause trie reads. - let trie = tries.get_trie_with_block_hash_for_shard( - shard_uid, - state_root, - &CryptoHash::default(), - false, - ); - for key in data_in_trie.keys() { - trie.get_optimized_ref(key, crate::KeyLookupMode::FlatStorage).unwrap(); - } - assert_eq!(trie.get_trie_nodes_count(), TrieNodesCount { db_reads: 0, mem_reads: 0 }); - - // Now, let's capture the baseline node counts - this is what will happen - // in production. - let trie = tries.get_trie_with_block_hash_for_shard( - shard_uid, - state_root, - &CryptoHash::default(), - false, - ); - trie.accounting_cache.borrow_mut().set_enabled(enable_accounting_cache); - for key in &keys_to_get { - assert_eq!(trie.get(key).unwrap(), data_in_trie.get(key).cloned()); - } - for key in &keys_to_get_ref { - assert_eq!( - trie.get_optimized_ref(key, crate::KeyLookupMode::FlatStorage) - .unwrap() - .map(|value| value.into_value_ref()), - data_in_trie.get(key).map(|value| ValueRef::new(&value)) - ); - } - let baseline_trie_nodes_count = trie.get_trie_nodes_count(); - println!("Baseline trie nodes count: {:?}", baseline_trie_nodes_count); - trie.update(updates.iter().cloned()).unwrap(); - - // Let's do this again, but this time recording reads. We'll make sure - // the counters are exactly the same even when we're recording. - let trie = tries - .get_trie_with_block_hash_for_shard( - shard_uid, - state_root, - &CryptoHash::default(), - false, - ) + let baseline_partial_storage = trie.recorded_storage().unwrap(); + + // Now let's do this again with memtries enabled. Check that counters + // are the same. + assert_eq!(MEM_TRIE_NUM_LOOKUPS.get(), mem_trie_lookup_counts_before); + tries.load_mem_trie(&shard_uid, None).unwrap(); + // Delete the on-disk state so that we really know we're using + // in-memory tries. + destructively_delete_in_memory_state_from_disk(&store, &data_in_trie); + let trie = get_trie_for_shard(&tries, shard_uid, state_root, use_flat_storage) .recording_reads(); trie.accounting_cache.borrow_mut().set_enabled(enable_accounting_cache); for key in &keys_to_get { @@ -437,7 +339,7 @@ mod trie_recording_tests { } for key in &keys_to_get_ref { assert_eq!( - trie.get_optimized_ref(key, crate::KeyLookupMode::FlatStorage) + trie.get_optimized_ref(key, lookup_mode) .unwrap() .map(|value| value.into_value_ref()), data_in_trie.get(key).map(|value| ValueRef::new(&value)) @@ -449,19 +351,21 @@ mod trie_recording_tests { // Now, let's check that when doing the same lookups with the captured partial storage, // we still get the same counters. let partial_storage = trie.recorded_storage().unwrap(); + assert_partial_storage(&baseline_partial_storage, &partial_storage); println!( "Partial storage has {} nodes from {} entries", partial_storage.nodes.len(), data_in_trie.len() ); - let trie = Trie::from_recorded_storage(partial_storage.clone(), state_root, true); + let trie = + Trie::from_recorded_storage(partial_storage.clone(), state_root, use_flat_storage); trie.accounting_cache.borrow_mut().set_enabled(enable_accounting_cache); for key in &keys_to_get { assert_eq!(trie.get(key).unwrap(), data_in_trie.get(key).cloned()); } for key in &keys_to_get_ref { assert_eq!( - trie.get_optimized_ref(key, crate::KeyLookupMode::FlatStorage) + trie.get_optimized_ref(key, lookup_mode) .unwrap() .map(|value| value.into_value_ref()), data_in_trie.get(key).map(|value| ValueRef::new(&value)) @@ -471,15 +375,15 @@ mod trie_recording_tests { trie.update(updates.iter().cloned()).unwrap(); // Build a Trie using recorded storage and enable recording_reads on this Trie - let trie = - Trie::from_recorded_storage(partial_storage, state_root, true).recording_reads(); + let trie = Trie::from_recorded_storage(partial_storage, state_root, use_flat_storage) + .recording_reads(); trie.accounting_cache.borrow_mut().set_enabled(enable_accounting_cache); for key in &keys_to_get { assert_eq!(trie.get(key).unwrap(), data_in_trie.get(key).cloned()); } for key in &keys_to_get_ref { assert_eq!( - trie.get_optimized_ref(key, crate::KeyLookupMode::FlatStorage) + trie.get_optimized_ref(key, lookup_mode) .unwrap() .map(|value| value.into_value_ref()), data_in_trie.get(key).map(|value| ValueRef::new(&value)) @@ -487,8 +391,9 @@ mod trie_recording_tests { } assert_eq!(trie.get_trie_nodes_count(), baseline_trie_nodes_count); trie.update(updates.iter().cloned()).unwrap(); + assert_partial_storage(&baseline_partial_storage, &trie.recorded_storage().unwrap()); - if use_in_memory_tries { + if !keys_to_get.is_empty() || !keys_to_get_ref.is_empty() { // sanity check that we did indeed use in-memory tries. assert!(MEM_TRIE_NUM_LOOKUPS.get() > mem_trie_lookup_counts_before); } @@ -496,44 +401,42 @@ mod trie_recording_tests { } #[test] - fn test_trie_recording_consistency_with_flat_storage_no_accounting_cache() { - test_trie_recording_consistency_with_flat_storage(false, false, false); + fn test_trie_recording_consistency_no_accounting_cache() { + test_trie_recording_consistency(false, false, false); } #[test] - fn test_trie_recording_consistency_with_flat_storage_with_accounting_cache() { - test_trie_recording_consistency_with_flat_storage(true, false, false); + fn test_trie_recording_consistency_with_accounting_cache() { + test_trie_recording_consistency(true, false, false); } #[test] - fn test_trie_recording_consistency_with_flat_storage_no_accounting_cache_with_missing_keys() { - test_trie_recording_consistency_with_flat_storage(false, true, false); + fn test_trie_recording_consistency_no_accounting_cache_with_missing_keys() { + test_trie_recording_consistency(false, true, false); } #[test] - fn test_trie_recording_consistency_with_flat_storage_with_accounting_cache_and_missing_keys() { - test_trie_recording_consistency_with_flat_storage(true, true, false); + fn test_trie_recording_consistency_with_accounting_cache_and_missing_keys() { + test_trie_recording_consistency(true, true, false); } #[test] - fn test_trie_recording_consistency_with_flat_storage_memtrie_no_accounting_cache() { - test_trie_recording_consistency_with_flat_storage(false, false, true); + fn test_trie_recording_consistency_with_flat_storage_no_accounting_cache() { + test_trie_recording_consistency(false, false, true); } #[test] - fn test_trie_recording_consistency_with_flat_storage_memtrie_with_accounting_cache() { - test_trie_recording_consistency_with_flat_storage(true, false, true); + fn test_trie_recording_consistency_with_flat_storage_with_accounting_cache() { + test_trie_recording_consistency(true, false, true); } #[test] - fn test_trie_recording_consistency_with_flat_storage_memtrie_no_accounting_cache_with_missing_keys( - ) { - test_trie_recording_consistency_with_flat_storage(false, true, true); + fn test_trie_recording_consistency_with_flat_storage_no_accounting_cache_with_missing_keys() { + test_trie_recording_consistency(false, true, true); } #[test] - fn test_trie_recording_consistency_with_flat_storage_memtrie_with_accounting_cache_and_missing_keys( - ) { - test_trie_recording_consistency_with_flat_storage(true, true, true); + fn test_trie_recording_consistency_with_flat_storage_with_accounting_cache_and_missing_keys() { + test_trie_recording_consistency(true, true, true); } } diff --git a/core/store/src/trie/trie_tests.rs b/core/store/src/trie/trie_tests.rs index 7a6cf242ef4..ebae5aae21f 100644 --- a/core/store/src/trie/trie_tests.rs +++ b/core/store/src/trie/trie_tests.rs @@ -419,36 +419,20 @@ mod trie_storage_tests { assert_eq!(count_delta.mem_reads, 1); } - // Checks that when branch restructuring is triggered on updating trie, - // impacted child is still recorded. - // - // Needed when branch has two children, one of which is removed, branch - // could be converted to extension, so reading of the only remaining child - // is also required. - #[test] - fn test_memtrie_recorded_branch_restructuring() { + fn test_memtrie_and_disk_updates_consistency(updates: Vec<(Vec, Option>)>) { init_test_logger(); + let base_changes = vec![ + (vec![7], Some(vec![1])), + (vec![7, 0], Some(vec![2])), + (vec![7, 1], Some(vec![3])), + ]; let tries = TestTriesBuilder::new().build(); let shard_uid = ShardUId::single_shard(); - let state_root = test_populate_trie( - &tries, - &Trie::EMPTY_ROOT, - shard_uid, - vec![ - (vec![7], Some(vec![1])), - (vec![7, 0], Some(vec![2])), - (vec![7, 1], Some(vec![3])), - ], - ); + let state_root = + test_populate_trie(&tries, &Trie::EMPTY_ROOT, shard_uid, base_changes.clone()); let trie = tries.get_trie_for_shard(shard_uid, state_root).recording_reads(); - let changes = trie - .update(vec![ - (vec![7], Some(vec![10])), - (vec![7, 0], None), - (vec![7, 6], Some(vec![8])), - ]) - .unwrap(); + let changes = trie.update(updates.clone()).unwrap(); tracing::info!("Changes: {:?}", changes); let recorded_normal = trie.recorded_storage(); @@ -469,29 +453,14 @@ mod trie_storage_tests { let tries = TestTriesBuilder::new() .with_store(store) - .with_flat_storage() + .with_flat_storage(true) .with_in_memory_tries() .build(); let shard_uid = ShardUId::single_shard(); - let state_root = test_populate_trie( - &tries, - &Trie::EMPTY_ROOT, - shard_uid, - vec![ - (vec![7], Some(vec![1])), - (vec![7, 0], Some(vec![2])), - (vec![7, 1], Some(vec![3])), - ], - ); + let state_root = test_populate_trie(&tries, &Trie::EMPTY_ROOT, shard_uid, base_changes); let trie = tries.get_trie_for_shard(shard_uid, state_root).recording_reads(); - let changes = trie - .update(vec![ - (vec![7], Some(vec![10])), - (vec![7, 0], None), - (vec![7, 6], Some(vec![8])), - ]) - .unwrap(); + let changes = trie.update(updates).unwrap(); tracing::info!("Changes: {:?}", changes); @@ -499,4 +468,29 @@ mod trie_storage_tests { assert_eq!(recorded_normal, recorded_memtrie); } + + // Checks that when branch restructuring is triggered on updating trie, + // impacted child is recorded on memtrie. + // + // Needed when branch has two children, one of which is removed, branch + // could be converted to extension, so reading of the only remaining child + // is also required. + #[test] + fn test_memtrie_recorded_branch_restructuring() { + test_memtrie_and_disk_updates_consistency(vec![ + (vec![7], Some(vec![1])), + (vec![7, 0], Some(vec![2])), + (vec![7, 1], Some(vec![3])), + ]); + } + + // Checks that when non-existent key is removed, only nodes along the path + // to it is recorded. + // Needed because old disk trie logic was always reading neighbouring children + // along the path to recompute memory usages, which is not needed if trie + // structure doesn't change. + #[test] + fn test_memtrie_recorded_delete_non_existent_key() { + test_memtrie_and_disk_updates_consistency(vec![(vec![8], None)]); + } } diff --git a/integration-tests/src/runtime_utils.rs b/integration-tests/src/runtime_utils.rs index 3f5ef923473..acf239a8099 100644 --- a/integration-tests/src/runtime_utils.rs +++ b/integration-tests/src/runtime_utils.rs @@ -36,7 +36,7 @@ pub fn get_runtime_and_trie_from_genesis(genesis: &Genesis) -> (Runtime, ShardTr let shard_layout = &genesis.config.shard_layout; let tries = TestTriesBuilder::new() .with_shard_layout(shard_layout.version(), shard_layout.shard_ids().count() as NumShards) - .with_flat_storage() + .with_flat_storage(true) .build(); let runtime = Runtime::new(); let mut account_ids: HashSet = HashSet::new();