[Test] Refactor create_*_tries() and ShardTries::test_*() into TestTr…

…iesBuilder.

chain/chain/src/test_utils/kv_runtime.rs

@@ -7,6 +7,7 @@ use borsh::{BorshDeserialize, BorshSerialize};
 use near_epoch_manager::types::BlockHeaderInfo;
 use near_epoch_manager::{EpochManagerAdapter, RngSeed};
 use near_primitives::state_part::PartId;
+use near_store::test_utils::TestTriesBuilder;
 use num_rational::Ratio;
 
 use near_chain_configs::{ProtocolConfig, DEFAULT_GC_NUM_EPOCHS_TO_KEEP};
@@ -334,7 +335,10 @@ impl KeyValueRuntime {
         let num_shards = epoch_manager.num_shards(&EpochId::default()).unwrap();
         let epoch_length =
             epoch_manager.get_epoch_config(&EpochId::default()).unwrap().epoch_length;
-        let tries = ShardTries::test(store.clone(), num_shards);
+        let tries = TestTriesBuilder::new()
+            .with_store(store.clone())
+            .with_shard_layout(0, num_shards)
+            .build();
         let mut initial_amounts = HashMap::new();
         for (i, validator_stake) in epoch_manager
             .validators_by_valset

chain/client/src/test_utils/test_env.rs

@@ -22,7 +22,6 @@ use near_primitives::epoch_manager::RngSeed;
 use near_primitives::errors::InvalidTxError;
 use near_primitives::hash::CryptoHash;
 use near_primitives::runtime::config::RuntimeConfig;
-use near_primitives::shard_layout::ShardUId;
 use near_primitives::sharding::PartialEncodedChunk;
 use near_primitives::test_utils::create_test_signer;
 use near_primitives::transaction::{Action, FunctionCallAction, SignedTransaction};
@@ -315,13 +314,17 @@ impl TestEnv {
     }
 
     pub fn query_account(&mut self, account_id: AccountId) -> AccountView {
-        let head = self.clients[0].chain.head().unwrap();
-        let last_block = self.clients[0].chain.get_block(&head.last_block_hash).unwrap();
-        let last_chunk_header = &last_block.chunks()[0];
-        let response = self.clients[0]
+        let client = &self.clients[0];
+        let head = client.chain.head().unwrap();
+        let last_block = client.chain.get_block(&head.last_block_hash).unwrap();
+        let shard_id =
+            client.epoch_manager.account_id_to_shard_id(&account_id, &head.epoch_id).unwrap();
+        let shard_uid = client.epoch_manager.shard_id_to_uid(shard_id, &head.epoch_id).unwrap();
+        let last_chunk_header = &last_block.chunks()[shard_id as usize];
+        let response = client
             .runtime_adapter
             .query(
-                ShardUId::single_shard(),
+                shard_uid,
                 &last_chunk_header.prev_state_root(),
                 last_block.header().height(),
                 last_block.header().raw_timestamp(),
@@ -338,13 +341,17 @@ impl TestEnv {
     }
 
     pub fn query_state(&mut self, account_id: AccountId) -> Vec<StateItem> {
-        let head = self.clients[0].chain.head().unwrap();
-        let last_block = self.clients[0].chain.get_block(&head.last_block_hash).unwrap();
-        let last_chunk_header = &last_block.chunks()[0];
-        let response = self.clients[0]
+        let client = &self.clients[0];
+        let head = client.chain.head().unwrap();
+        let last_block = client.chain.get_block(&head.last_block_hash).unwrap();
+        let shard_id =
+            client.epoch_manager.account_id_to_shard_id(&account_id, &head.epoch_id).unwrap();
+        let shard_uid = client.epoch_manager.shard_id_to_uid(shard_id, &head.epoch_id).unwrap();
+        let last_chunk_header = &last_block.chunks()[shard_id as usize];
+        let response = client
             .runtime_adapter
             .query(
-                ShardUId::single_shard(),
+                shard_uid,
                 &last_chunk_header.prev_state_root(),
                 last_block.header().height(),
                 last_block.header().raw_timestamp(),

core/store/benches/trie_bench.rs

@@ -5,7 +5,7 @@ use bencher::Bencher;
 use rand::random;
 
 use near_primitives::shard_layout::ShardUId;
-use near_store::test_utils::create_tries;
+use near_store::test_utils::TestTriesBuilder;
 use near_store::Trie;
 
 fn rand_bytes() -> Vec<u8> {
@@ -14,7 +14,7 @@ fn rand_bytes() -> Vec<u8> {
 
 fn trie_lookup(bench: &mut Bencher) {
     let (changed_keys, trie) = {
-        let tries = create_tries();
+        let tries = TestTriesBuilder::new().build();
 
         let trie = tries.get_trie_for_shard(ShardUId::single_shard(), Trie::EMPTY_ROOT);
         let mut changes = vec![];
@@ -42,7 +42,7 @@ fn trie_lookup(bench: &mut Bencher) {
 }
 
 fn trie_update(bench: &mut Bencher) {
-    let tries = create_tries();
+    let tries = TestTriesBuilder::new().build();
     let trie = tries.get_trie_for_shard(ShardUId::single_shard(), Trie::EMPTY_ROOT);
     let mut changes = vec![];
     for _ in 0..100 {

core/store/src/test_utils.rs

@@ -1,22 +1,25 @@
-use std::collections::HashMap;
-use std::sync::Arc;
-
-use itertools::Itertools;
-use near_primitives::state::{FlatStateValue, ValueRef};
-use near_primitives::trie_key::TrieKey;
-use rand::seq::SliceRandom;
-use rand::Rng;
-
 use crate::db::TestDB;
-use crate::flat::{store_helper, BlockInfo, FlatStorageReadyStatus, FlatStorageStatus};
+use crate::flat::{
+    store_helper, BlockInfo, FlatStorageManager, FlatStorageReadyStatus, FlatStorageStatus,
+};
 use crate::metadata::{DbKind, DbVersion, DB_VERSION};
-use crate::{get, get_delayed_receipt_indices, DBCol, NodeStorage, ShardTries, Store};
+use crate::{
+    get, get_delayed_receipt_indices, DBCol, NodeStorage, ShardTries, StateSnapshotConfig, Store,
+    TrieConfig,
+};
+use itertools::Itertools;
 use near_primitives::account::id::AccountId;
 use near_primitives::hash::{hash, CryptoHash};
 use near_primitives::receipt::{DataReceipt, Receipt, ReceiptEnum};
 use near_primitives::shard_layout::{ShardUId, ShardVersion};
+use near_primitives::state::{FlatStateValue, ValueRef};
+use near_primitives::trie_key::TrieKey;
 use near_primitives::types::{NumShards, StateRoot};
+use rand::seq::SliceRandom;
+use rand::Rng;
+use std::collections::HashMap;
 use std::str::from_utf8;
+use std::sync::Arc;
 
 /// Creates an in-memory node storage.
 ///
@@ -61,44 +64,75 @@ pub fn create_test_store() -> Store {
     create_test_node_storage(DB_VERSION, DbKind::RPC).get_hot_store()
 }
 
-/// Creates a Trie using an in-memory database.
-pub fn create_tries() -> ShardTries {
-    create_tries_complex(0, 1)
+pub struct TestTriesBuilder {
+    store: Option<Store>,
+    shard_version: ShardVersion,
+    num_shards: NumShards,
+    enable_flat_storage: bool,
 }
 
-pub fn create_tries_complex(shard_version: ShardVersion, num_shards: NumShards) -> ShardTries {
-    let store = create_test_store();
-    ShardTries::test_shard_version(store, shard_version, num_shards)
-}
+impl TestTriesBuilder {
+    pub fn new() -> Self {
+        Self { store: None, shard_version: 0, num_shards: 1, enable_flat_storage: false }
+    }
 
-pub fn create_tries_with_flat_storage() -> ShardTries {
-    create_tries_complex_with_flat_storage(0, 1)
-}
+    pub fn with_store(mut self, store: Store) -> Self {
+        self.store = Some(store);
+        self
+    }
 
-pub fn create_tries_complex_with_flat_storage(
-    shard_version: ShardVersion,
-    num_shards: NumShards,
-) -> ShardTries {
-    let tries = create_tries_complex(shard_version, num_shards);
-    let mut store_update = tries.store_update();
-    for shard_id in 0..num_shards {
-        let shard_uid = ShardUId { version: shard_version, shard_id: shard_id.try_into().unwrap() };
-        store_helper::set_flat_storage_status(
-            &mut store_update,
-            shard_uid,
-            FlatStorageStatus::Ready(FlatStorageReadyStatus {
-                flat_head: BlockInfo::genesis(CryptoHash::default(), 0),
-            }),
-        );
+    pub fn with_shard_layout(mut self, shard_version: ShardVersion, num_shards: NumShards) -> Self {
+        self.shard_version = shard_version;
+        self.num_shards = num_shards;
+        self
     }
-    store_update.commit().unwrap();
 
-    let flat_storage_manager = tries.get_flat_storage_manager();
-    for shard_id in 0..num_shards {
-        let shard_uid = ShardUId { version: shard_version, shard_id: shard_id.try_into().unwrap() };
-        flat_storage_manager.create_flat_storage_for_shard(shard_uid).unwrap();
+    pub fn with_flat_storage(mut self) -> Self {
+        self.enable_flat_storage = true;
+        self
+    }
+
+    pub fn build(self) -> ShardTries {
+        let store = self.store.unwrap_or_else(create_test_store);
+        let shard_uids = (0..self.num_shards)
+            .map(|shard_id| ShardUId { shard_id: shard_id as u32, version: self.shard_version })
+            .collect::<Vec<_>>();
+        let flat_storage_manager = FlatStorageManager::new(store.clone());
+        let tries = ShardTries::new(
+            store,
+            TrieConfig::default(),
+            &shard_uids,
+            flat_storage_manager,
+            StateSnapshotConfig::default(),
+        );
+        if self.enable_flat_storage {
+            let mut store_update = tries.store_update();
+            for shard_id in 0..self.num_shards {
+                let shard_uid = ShardUId {
+                    version: self.shard_version,
+                    shard_id: shard_id.try_into().unwrap(),
+                };
+                store_helper::set_flat_storage_status(
+                    &mut store_update,
+                    shard_uid,
+                    FlatStorageStatus::Ready(FlatStorageReadyStatus {
+                        flat_head: BlockInfo::genesis(CryptoHash::default(), 0),
+                    }),
+                );
+            }
+            store_update.commit().unwrap();
+
+            let flat_storage_manager = tries.get_flat_storage_manager();
+            for shard_id in 0..self.num_shards {
+                let shard_uid = ShardUId {
+                    version: self.shard_version,
+                    shard_id: shard_id.try_into().unwrap(),
+                };
+                flat_storage_manager.create_flat_storage_for_shard(shard_uid).unwrap();
+            }
+        }
+        tries
     }
-    tries
 }
 
 pub fn test_populate_trie(

core/store/src/trie/iterator.rs

@@ -432,9 +432,7 @@ mod tests {
     use rand::seq::SliceRandom;
     use rand::Rng;
 
-    use crate::test_utils::{
-        create_tries, create_tries_complex, gen_changes, simplify_changes, test_populate_trie,
-    };
+    use crate::test_utils::{gen_changes, simplify_changes, test_populate_trie, TestTriesBuilder};
     use crate::trie::iterator::IterStep;
     use crate::trie::nibble_slice::NibbleSlice;
     use crate::Trie;
@@ -449,7 +447,7 @@ mod tests {
     #[test]
     fn test_visit_interval() {
         let trie_changes = vec![(b"aa".to_vec(), Some(vec![1])), (b"abb".to_vec(), Some(vec![2]))];
-        let tries = create_tries();
+        let tries = TestTriesBuilder::new().build();
         let state_root =
             test_populate_trie(&tries, &Trie::EMPTY_ROOT, ShardUId::single_shard(), trie_changes);
         let trie = tries.get_trie_for_shard(ShardUId::single_shard(), state_root);
@@ -570,7 +568,7 @@ mod tests {
         max_depth: usize,
     ) {
         let shard_uid = ShardUId::single_shard();
-        let tries = create_tries();
+        let tries = TestTriesBuilder::new().build();
         let trie_changes = keys.iter().map(|key| (key.clone(), value())).collect();
         let state_root = test_populate_trie(&tries, &Trie::EMPTY_ROOT, shard_uid, trie_changes);
         let trie = tries.get_trie_for_shard(shard_uid, state_root);
@@ -613,7 +611,7 @@ mod tests {
     fn gen_random_trie(
         rng: &mut rand::rngs::ThreadRng,
     ) -> (Vec<(Vec<u8>, Option<Vec<u8>>)>, BTreeMap<Vec<u8>, Vec<u8>>, Trie) {
-        let tries = create_tries_complex(1, 2);
+        let tries = TestTriesBuilder::new().with_shard_layout(1, 2).build();
         let shard_uid = ShardUId { version: 1, shard_id: 0 };
         let trie_changes = gen_changes(rng, 10);
         let trie_changes = simplify_changes(&trie_changes);
@@ -654,7 +652,7 @@ mod tests {
     fn test_has_value() {
         let mut rng = rand::thread_rng();
         for _ in 0..100 {
-            let tries = create_tries();
+            let tries = TestTriesBuilder::new().build();
             let trie_changes = gen_changes(&mut rng, 10);
             let trie_changes = simplify_changes(&trie_changes);
             let state_root = test_populate_trie(

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

@@ -78,7 +78,7 @@ pub fn load_trie_from_flat_state(
 #[cfg(test)]
 mod tests {
     use crate::test_utils::{
-        create_tries, simplify_changes, test_populate_flat_storage, test_populate_trie,
+        simplify_changes, test_populate_flat_storage, test_populate_trie, TestTriesBuilder,
     };
     use crate::trie::mem::loading::load_trie_from_flat_state;
     use crate::trie::mem::lookup::memtrie_lookup;
@@ -91,7 +91,7 @@ mod tests {
     use std::collections::HashSet;
 
     fn check(keys: Vec<Vec<u8>>) {
-        let shard_tries = create_tries();
+        let shard_tries = TestTriesBuilder::new().build();
         let shard_uid = ShardUId::single_shard();
         let changes = keys.iter().map(|key| (key.to_vec(), Some(key.to_vec()))).collect::<Vec<_>>();
         let changes = simplify_changes(&changes);

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

@@ -1,5 +1,6 @@
 use near_o11y::metrics::{
-    try_create_int_counter_vec, try_create_int_gauge_vec, IntCounterVec, IntGaugeVec,
+    try_create_int_counter, try_create_int_counter_vec, try_create_int_gauge_vec, IntCounter,
+    IntCounterVec, IntGaugeVec,
 };
 use once_cell::sync::Lazy;
 
@@ -20,3 +21,11 @@ pub static MEM_TRIE_NUM_NODES_CREATED_FROM_UPDATES: Lazy<IntCounterVec> = Lazy::
     )
     .unwrap()
 });
+
+pub static MEM_TRIE_NUM_LOOKUPS: Lazy<IntCounter> = Lazy::new(|| {
+    try_create_int_counter(
+        "near_mem_trie_num_lookups",
+        "Number of in-memory trie lookups (number of keys looked up)",
+    )
+    .unwrap()
+});

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

@@ -13,7 +13,7 @@ mod construction;
 mod flexible_data;
 pub mod loading;
 pub mod lookup;
-mod metrics;
+pub mod metrics;
 pub mod node;
 pub mod updating;
 
@@ -86,7 +86,7 @@ impl MemTries {
             self.roots.entry(state_root).or_default().push(mem_root);
         }
         MEM_TRIE_NUM_ROOTS
-            .with_label_values(&[&self.shard_uid.shard_id.to_string()])
+            .with_label_values(&[&self.shard_uid.to_string()])
             .set(self.roots.len() as i64);
     }
 
@@ -131,7 +131,7 @@ impl MemTries {
             debug_assert!(false, "Deleting non-existent root: {}", state_root);
         }
         MEM_TRIE_NUM_ROOTS
-            .with_label_values(&[&self.shard_uid.shard_id.to_string()])
+            .with_label_values(&[&self.shard_uid.to_string()])
             .set(self.roots.len() as i64);
     }
 

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

@@ -869,7 +869,7 @@ pub fn apply_memtrie_changes(memtries: &mut MemTries, changes: &MemTrieChanges)
 
 #[cfg(test)]
 mod tests {
-    use crate::test_utils::create_test_store;
+    use crate::test_utils::TestTriesBuilder;
     use crate::trie::mem::lookup::memtrie_lookup;
     use crate::trie::mem::updating::apply_memtrie_changes;
     use crate::trie::mem::MemTries;
@@ -893,8 +893,7 @@ mod tests {
     impl TestTries {
         fn new(check_deleted_keys: bool) -> Self {
             let mem = MemTries::new(100 * 1024 * 1024, ShardUId::single_shard());
-            let store = create_test_store();
-            let disk = ShardTries::test(store, 1);
+            let disk = TestTriesBuilder::new().build();
             Self {
                 mem,
                 disk,